ARIA: Adaptive Region-Based Importance Allocation for Conditional Diffusion Distillation

# ARIA: Adaptive Region-Based Importance Allocation for Conditional Diffusion Distillation
Source: [https://arxiv.org/html/2606.23898](https://arxiv.org/html/2606.23898)
Loay Mualem1,2Vinh Tong1Samir Darouich1,3Mathias Niepert1 1Institute for AI, University of Stuttgart2IMPRS\-IS 3Institute of Theoretical Chemistry, University of Stuttgart\. Please send any questions to:[loaymua@gmail\.com](https://arxiv.org/html/2606.23898v1/mailto:[email protected])

###### Abstract

Distilling conditional diffusion models aims to transfer the behavior of a large teacher to a smaller student while preserving alignment across conditioning inputs\. Unlike recognition tasks, knowledge distillation in conditional diffusion often struggles to transfer knowledge beyond the training distribution, since the predicted noise strongly depends on the conditioning signal\. As a result, effective distillation requires exploring a large conditioning space\. In practical settings, this creates a major bottleneck\. Paired image–condition data may be limited, and generating synthetic images for every available condition is often computationally infeasible, while the pool of conditions, such as text prompts, can be extremely large\. Recent work addresses this issue by switching conditions during training, exposing the student to a broader conditioning space without changing the distillation objective\. Yet this raises a complementary question: once a large conditioning corpus is available, how should the training effort be allocated? In this work, we introduce ARIA, a framework that adaptively allocates training effort across coarse regions of the conditioning space\. By maintaining online estimates of teacher–student discrepancy at the region level, ARIA focuses updates where misalignment persists while preserving the original distillation objective\. Empirically, ARIA improves over RC across most architectures and settings, with the clearest gains observed in unseen and underrepresented regimes\. We also provide a theoretical analysis showing that the proposed tracking mechanism follows the evolving discrepancy during training under bounded variance and drift assumptions\.

## 1Introduction

Diffusion models have achieved strong generative performance across many domains, including images\[[21](https://arxiv.org/html/2606.23898#bib.bib2),[51](https://arxiv.org/html/2606.23898#bib.bib3),[12](https://arxiv.org/html/2606.23898#bib.bib4),[2](https://arxiv.org/html/2606.23898#bib.bib7),[42](https://arxiv.org/html/2606.23898#bib.bib5),[38](https://arxiv.org/html/2606.23898#bib.bib6)\], video\[[20](https://arxiv.org/html/2606.23898#bib.bib8),[15](https://arxiv.org/html/2606.23898#bib.bib9),[3](https://arxiv.org/html/2606.23898#bib.bib10)\], audio\[[27](https://arxiv.org/html/2606.23898#bib.bib11),[34](https://arxiv.org/html/2606.23898#bib.bib12)\], and robotics\[[9](https://arxiv.org/html/2606.23898#bib.bib13),[6](https://arxiv.org/html/2606.23898#bib.bib43),[9](https://arxiv.org/html/2606.23898#bib.bib13)\]\. These models iteratively transform noise into samples from a target distribution\. Large\-scale text\-to\-image systems such as Stable Diffusion\[[41](https://arxiv.org/html/2606.23898#bib.bib34)\]generate images that follow natural\-language prompts, however, their strong performance typically requires large models and many sampling steps, motivating efforts to develop more efficient variants\.

Knowledge distillation \(KD\) is a widely used approach for compressing large models by training a smaller student to mimic the predictions or intermediate representations of a teacher\[[19](https://arxiv.org/html/2606.23898#bib.bib18),[60](https://arxiv.org/html/2606.23898#bib.bib19),[63](https://arxiv.org/html/2606.23898#bib.bib20),[66](https://arxiv.org/html/2606.23898#bib.bib21),[61](https://arxiv.org/html/2606.23898#bib.bib49),[36](https://arxiv.org/html/2606.23898#bib.bib50),[59](https://arxiv.org/html/2606.23898#bib.bib51),[67](https://arxiv.org/html/2606.23898#bib.bib52),[44](https://arxiv.org/html/2606.23898#bib.bib53),[37](https://arxiv.org/html/2606.23898#bib.bib54),[11](https://arxiv.org/html/2606.23898#bib.bib55)\]\. In text\-to\-image diffusion models, this typically involves sampling image–text pairs, generating noisy intermediates along the diffusion trajectory, and training the student to match the teacher’s outputs at each timestep\[[19](https://arxiv.org/html/2606.23898#bib.bib18),[60](https://arxiv.org/html/2606.23898#bib.bib19),[63](https://arxiv.org/html/2606.23898#bib.bib20),[66](https://arxiv.org/html/2606.23898#bib.bib21)\]or feature representations\[[60](https://arxiv.org/html/2606.23898#bib.bib19),[63](https://arxiv.org/html/2606.23898#bib.bib20),[66](https://arxiv.org/html/2606.23898#bib.bib21),[8](https://arxiv.org/html/2606.23898#bib.bib22),[30](https://arxiv.org/html/2606.23898#bib.bib23)\]\. While KD can transfer knowledge beyond the training distribution in recognition models\[[8](https://arxiv.org/html/2606.23898#bib.bib22),[30](https://arxiv.org/html/2606.23898#bib.bib23)\], this effect is weaker in conditional diffusion models\[[26](https://arxiv.org/html/2606.23898#bib.bib17)\], where the predicted noise is strongly condition\-dependent\. As a result, effectively distilling the teacher often requires exploring a large portion of the conditioning space\.

In practice, access to large\-scale paired image–text datasets is often limited due to copyright, privacy, and licensing constraints\. Generating synthetic images with the teacher model is a possible alternative, but large\-scale diffusion sampling is computationally expensive and storage\-intensive\[[67](https://arxiv.org/html/2606.23898#bib.bib52),[44](https://arxiv.org/html/2606.23898#bib.bib53)\]\. In contrast, textual data is abundant: large collections of prompts and captions are widely available at negligible cost\. This creates a fundamental imbalance—while the conditioning space \(text\) can be extremely large, the number of available images for distillation is typically constrained by computational resources\.

To address this mismatch, Kim et al\.\[[26](https://arxiv.org/html/2606.23898#bib.bib17)\]proposed Random Conditioning \(RC\), which pairs noisy images with randomly sampled text conditions, enabling exploration of the conditioning space without generating images for every prompt\. RC answers an important feasibility question: can auxiliary text improve diffusion distillation when paired image–prompt data is limited? Our work asks the complementary allocation question: once a large auxiliary text pool is available, how should the training budget be distributed over it? A static coverage strategy treats the auxiliary conditioning space as a fixed resource and samples it uniformly\. However, the usefulness of a region is student\-dependent and changes during training, as some regions are learned quickly, whereas others continue to exhibit a large teacher–student gap\. To this end we introduce ARIA \(Adaptive Region\-based Importance Allocation\), which uses auxiliary text not only to expand coverage, but to provide adaptive feedback\. ARIA groups auxiliary conditions into coarse regions, tracks their teacher–student discrepancy online, and reallocates sampling toward regions where supervision is currently most needed\.

![Refer to caption](https://arxiv.org/html/2606.23898v1/x1.png)Figure 1:Side\-by\-side comparison of images generated by ARIA, RC, and Teacher models under three settings: training without animals, vehicles, or food; using25%25\\%training checkpoints; and fully trained on the complete dataset\.Concretely, ARIA is a lightweight replacement for the condition\-selection step in RC\. The paired\-image cache, teacher–student objective, and model architectures remain unchanged, while the auxiliary prompt sampler becomes discrepancy\-aware\. We instantiate ARIA for text\-conditioned diffusion model distillation in the regime of limited paired images and abundant auxiliary text\. ARIA leaves the underlying distillation objective unchanged and only modifies how conditions are sampled\. We provide theoretical analysis of the region\-level discrepancy tracking mechanism and demonstrate that ARIA improves over random conditioning across most architectures and data regimes, with the strongest gains in unseen and underrepresented settings, and faster convergence\.

Contributions\.Our main contributions are as follows: \(1\) We introduce ARIA, a region\-based framework for adaptive importance allocation that prioritizes coarse condition regions with higher discrepancy\. Unlike per\-sample importance sampling, ARIA tracks discrepancy at the level of coarse condition regions, enabling scalable adaptive allocation over extremely large input spaces without per\-sample scoring\. \(2\) We instantiate ARIA for text\-to\-image diffusion distillation with limited paired images and abundant auxiliary text, achieving consistent improvements over random sampling, particularly for unseen prompts and imbalanced conditioning distributions\. \(3\) We present extensive empirical evaluation demonstrating stable training behavior and robustness across architectures and region constructions including SD 1\.4, SD 2\.1, and SDXL teachers with block\-pruned, channel\-pruned, and KOALA students\. \(4\) We provide theoretical analysis of ARIA’s EMA\-based region scoring rule, establishing finite\-time guarantees for tracking evolving discrepancy under noise and temporal drift\.

Organization\.Section[2](https://arxiv.org/html/2606.23898#S2)reviews related work, and Section[3](https://arxiv.org/html/2606.23898#S3)formalizes the distillation setting and RC baseline\. Section[4](https://arxiv.org/html/2606.23898#S4)introduces ARIA and presents its algorithmic formulation together with the tracking analysis\. Section[5](https://arxiv.org/html/2606.23898#S5)applies ARIA to text\-to\-image distillation under two experimental setups, demonstrating performance improvements and robustness\. We conclude in Section[6](https://arxiv.org/html/2606.23898#S6)with a discussion and future directions\.

## 2Related Work

Compressed Diffusion Models\.Modern diffusion models are typically large and require substantial computational and memory resources for both training and inference\. Several works aim to reduce the computational and memory cost of diffusion models through model compression\. Existing approaches include quantization\[[48](https://arxiv.org/html/2606.23898#bib.bib33)\], architecture optimization\[[32](https://arxiv.org/html/2606.23898#bib.bib37)\], pruning\[[14](https://arxiv.org/html/2606.23898#bib.bib44),[28](https://arxiv.org/html/2606.23898#bib.bib45),[68](https://arxiv.org/html/2606.23898#bib.bib46),[58](https://arxiv.org/html/2606.23898#bib.bib47),[7](https://arxiv.org/html/2606.23898#bib.bib48)\]and knowledge distillation\[[57](https://arxiv.org/html/2606.23898#bib.bib30),[26](https://arxiv.org/html/2606.23898#bib.bib17),[25](https://arxiv.org/html/2606.23898#bib.bib1)\]\. For example, BK\-SDM\[[25](https://arxiv.org/html/2606.23898#bib.bib1)\]applies knowledge distillation together with block pruning to compress Stable Diffusion\[[41](https://arxiv.org/html/2606.23898#bib.bib34)\], while KOALA\[[29](https://arxiv.org/html/2606.23898#bib.bib35)\]performs layer\-wise compression and knowledge distillation for SDXL\[[39](https://arxiv.org/html/2606.23898#bib.bib36)\]\.

Knowledge Distillation for Model Compression\.Knowledge distillation \(KD\)\[[19](https://arxiv.org/html/2606.23898#bib.bib18)\]is widely used for training compact models by transferring knowledge from a larger teacher through soft predictions or intermediate representations\[[60](https://arxiv.org/html/2606.23898#bib.bib19),[63](https://arxiv.org/html/2606.23898#bib.bib20),[66](https://arxiv.org/html/2606.23898#bib.bib21),[8](https://arxiv.org/html/2606.23898#bib.bib22),[30](https://arxiv.org/html/2606.23898#bib.bib23)\]\. KD has been successfully applied across many domains, including language models, vision transformers, and diffusion models\[[52](https://arxiv.org/html/2606.23898#bib.bib26),[23](https://arxiv.org/html/2606.23898#bib.bib27),[16](https://arxiv.org/html/2606.23898#bib.bib28),[55](https://arxiv.org/html/2606.23898#bib.bib29),[25](https://arxiv.org/html/2606.23898#bib.bib1),[26](https://arxiv.org/html/2606.23898#bib.bib17),[57](https://arxiv.org/html/2606.23898#bib.bib30)\]\. In diffusion models, KD is commonly used to train compressed architectures\. However, Kim et al\.\[[26](https://arxiv.org/html/2606.23898#bib.bib17)\]show that conventional distillation strategies may struggle to transfer knowledge for underrepresented or uncovered concepts, motivating improved data allocation strategies during training\.

Knowledge Distillation for Decreasing Sampling Steps\.Several works focused on accelerating the denoising process in diffusion models without retraining\[[35](https://arxiv.org/html/2606.23898#bib.bib58),[64](https://arxiv.org/html/2606.23898#bib.bib59),[62](https://arxiv.org/html/2606.23898#bib.bib60),[65](https://arxiv.org/html/2606.23898#bib.bib61)\], reducing sampling steps from thousands to as few as 10–25\. However, pushing this reduction further typically leads to substantial performance degradation\. Distillation\-based acceleration methods\[[43](https://arxiv.org/html/2606.23898#bib.bib14),[50](https://arxiv.org/html/2606.23898#bib.bib15),[54](https://arxiv.org/html/2606.23898#bib.bib16),[11](https://arxiv.org/html/2606.23898#bib.bib55),[37](https://arxiv.org/html/2606.23898#bib.bib54),[67](https://arxiv.org/html/2606.23898#bib.bib52),[36](https://arxiv.org/html/2606.23898#bib.bib50),[59](https://arxiv.org/html/2606.23898#bib.bib51),[61](https://arxiv.org/html/2606.23898#bib.bib49)\]address this by training student models to compress multi\-step denoising trajectories into fewer steps, sometimes even a single step\. Importantly, these methods focus on reducing sampling steps rather than compressing model capacity\. In contrast, our work targets the compression of the base diffusion model itself, providing a compact foundation that can naturally complement and strengthen step\-acceleration techniques\.

## 3Problem Setting

We study the distillation of a pretrained text\-to\-image diffusion model \(theteacher\) into a smaller student model as presented in\[[26](https://arxiv.org/html/2606.23898#bib.bib17),[25](https://arxiv.org/html/2606.23898#bib.bib1)\]\. For clarity, we use the noise prediction objective, though the same derivation applies to score, velocity, or data prediction\. Letϵ𝒯​\(xt,t,c\)\\epsilon\_\{\\mathcal\{T\}\}\(x\_\{t\},t,c\)denote the teacher andϵ𝒮​\(xt,t,c\)\\epsilon\_\{\\mathcal\{S\}\}\(x\_\{t\},t,c\)the student, wherextx\_\{t\}is a noisy latent at timestepttandccis a text condition\.

Conditional Distillation Loss\.A straightforward strategy for image\-free distillation is to first synthesize images conditioned on text prompts and construct a paired dataset𝒟=\{\(𝐱n,cn\)\}n=1N\\mathcal\{D\}=\\\{\(\\mathbf\{x\}^\{n\},c^\{n\}\)\\\}\_\{n=1\}^\{N\}, where𝐱n\\mathbf\{x\}^\{n\}denotes the generated image corresponding to the text conditioncnc^\{n\}\. The generated image serves as the original clean sample𝐱0\\mathbf\{x\}\_\{0\}, from which we can produce a noisy input𝐱t\\mathbf\{x\}\_\{t\}for any timestepttunder conditioncnc^\{n\}\. Since diffusion models require substantial computational cost for image generation, these synthetic images are typically generated and stored in advance to form the training dataset\. The teacher model can subsequently be distilled into a student model using the following objective:

ℒout=𝔼\(𝐱t,c\)∈𝒟,t​\[‖ϵ𝒯​\(𝐱t,c,t\)−ϵ𝒮​\(𝐱t,c,t\)‖22\],\\mathcal\{L\}\_\{\\mathrm\{out\}\}=\\mathbb\{E\}\_\{\(\\mathbf\{x\}\_\{t\},c\)\\in\\mathcal\{D\},\\,t\}\\left\[\\left\\\|\\epsilon\_\{\\mathcal\{T\}\}\(\\mathbf\{x\}\_\{t\},c,t\)\-\\epsilon\_\{\\mathcal\{S\}\}\(\\mathbf\{x\}\_\{t\},c,t\)\\right\\\|\_\{2\}^\{2\}\\right\],\(1\)
whereϵ𝒯\\epsilon\_\{\\mathcal\{T\}\}andϵ𝒮\\epsilon\_\{\\mathcal\{S\}\}represent the noise predictions of the teacher and student networks, respectively\. Here,\(𝐱t,c\)\(\\mathbf\{x\}\_\{t\},c\)is sampled from the dataset𝒟\\mathcal\{D\}, with noise injected according to timesteptt, which is drawn uniformly from the interval\[0,T\]\[0,T\]\.

Furthermore, one may augment the training objective with a feature\-level distillation term defined as

ℒfeat=𝔼\(𝐱t,c\)∈𝒟,t​\[∑l‖f𝒯l​\(𝐱t,c,t\)−f𝒮l​\(𝐱t,c,t\)‖22\],\\mathcal\{L\}\_\{\\mathrm\{feat\}\}=\\mathbb\{E\}\_\{\(\\mathbf\{x\}\_\{t\},c\)\\in\\mathcal\{D\},\\,t\}\\left\[\\sum\_\{l\}\\left\\\|f\_\{\\mathcal\{T\}\}^\{l\}\(\\mathbf\{x\}\_\{t\},c,t\)\-f\_\{\\mathcal\{S\}\}^\{l\}\(\\mathbf\{x\}\_\{t\},c,t\)\\right\\\|\_\{2\}^\{2\}\\right\],\(2\)
wheref𝒯lf\_\{\\mathcal\{T\}\}^\{l\}andf𝒮lf\_\{\\mathcal\{S\}\}^\{l\}denote the feature representations extracted at layerllfrom the teacher and student models, respectively\.

Random Conditioning\.Let𝒞aux=\{cj\}j=1M\\mathcal\{C\}\_\{\\text\{aux\}\}=\\\{c\_\{j\}\\\}\_\{j=1\}^\{M\}denote a large corpus of text conditions\. Generating paired images for all conditions requires repeated diffusion sampling and is therefore computationally infeasible\. Instead, training relies on a smaller cached dataset𝒟=\{\(𝐱n,cn\)\}n=1N\\mathcal\{D\}=\\\{\(\\mathbf\{x\}^\{n\},c^\{n\}\)\\\}\_\{n=1\}^\{N\}withN≪MN\\ll M\. Kim et al\.\[[26](https://arxiv.org/html/2606.23898#bib.bib17)\]showed that both sources can be leveraged via*Random Conditioning \(RC\)*, which replaces the condition associated with a cached image during training\. Specifically, RC samples\(𝐱n,cn\)∼𝒟\(\\mathbf\{x\}^\{n\},c^\{n\}\)\\sim\\mathcal\{D\}and a timestept∼Unif​\(\[0,T\]\)t\\sim\\mathrm\{Unif\}\(\[0,T\]\), constructs𝐱t\\mathbf\{x\}\_\{t\}via the forward process, and sets the training condition

c^=\{cn,with probability​1−p​\(t\),c~∼Unif​\(𝒞aux\),with probability​p​\(t\),\\hat\{c\}=\\begin\{cases\}c^\{n\},&\\text\{with probability \}1\-p\(t\),\\\\ \\tilde\{c\}\\sim\\mathrm\{Unif\}\(\\mathcal\{C\}\_\{\\text\{aux\}\}\),&\\text\{with probability \}p\(t\),\\end\{cases\}\(3\)wherep​\(t\)p\(t\)is a timestep\-dependent schedule\.

Uniform sampling enables broad exploration of𝒞aux\\mathcal\{C\}\_\{\\text\{aux\}\}, but implicitly assumes that all regions are equally informative\. As training progresses, discrepancies across conditions become heterogeneous: some regions quickly align while others remain difficult\. Uniform exploration therefore wastes iterations on already\-aligned regions and under\-allocates effort to persistently misaligned ones\.

![Refer to caption](https://arxiv.org/html/2606.23898v1/x2.png)Figure 2:Overview of ARIA\.The conditioning domain𝒵\\mathcal\{Z\}is partitioned into coarse regionsRkR\_\{k\}\. ARIA maintains an EMA discrepancy scoreDkD\_\{k\}per region and converts these scores into a sampling distributionpt=𝒮​\(D1,…,DK\)p\_\{t\}=\\mathcal\{S\}\(D\_\{1\},\\dots,D\_\{K\}\)\. At each step, a regionkt∼ptk\_\{t\}\\sim p\_\{t\}and a samplezt∼πktz\_\{t\}\\sim\\pi\_\{k\_\{t\}\}are drawn, the discrepancyΔ​\(zt\)\\Delta\(z\_\{t\}\)is observed, and the corresponding score is updated via EMA\.
## 4ARIA: Adaptive Region\-based Importance Allocation

In what follows, we introduce ARIA \(Algorithm[1](https://arxiv.org/html/2606.23898#alg1), Figure[2](https://arxiv.org/html/2606.23898#S3.F2)\), our approach for adaptive importance allocation that operates at the level of coarse regions in the input space\. By tracking discrepancy statistics at the region level rather than per sample, ARIA enables adaptive reallocation of training effort toward persistently misaligned areas while remaining fully compatible with the underlying learning objective\. For clarity, we first describe ARIA in a general form and then instantiate it for diffusion distillation in Sec\.[5](https://arxiv.org/html/2606.23898#S5)\.

ARIA assumes that the training domain is partitioned into coarse regions\. These regions may be defined in any user\-specified representation \(e\.g\., clustering in embedding space, binning in latent space, or grouping by metadata\)\. During training, ARIA maintains a scalar discrepancy score for each region using an exponential moving average \(EMA\) and biases sampling toward regions with larger tracked discrepancy\. Importantly, ARIA does not modify the underlying learning objective; it only alters how training samples are selected\.

### 4\.1Regions and discrepancy signal

Let𝒵\\mathcal\{Z\}denote the set of training samples relevant to the learning objective\. In this work,𝒵\\mathcal\{Z\}corresponds to the auxiliary prompt set𝒞aux\\mathcal\{C\}\_\{\\text\{aux\}\}introduced in Sec\.[3](https://arxiv.org/html/2606.23898#S3)\. ARIA assumes a partition of𝒵\\mathcal\{Z\}intoKKcoarse regions of the training domain

ℛ=\{Rk\}k=1K,∪k=1KRk=𝒵,Ri∩Rj=∅​\(i≠j\)\.\\mathcal\{R\}=\\\{R\_\{k\}\\\}\_\{k=1\}^\{K\},\\quad\\cup\_\{k=1\}^\{K\}R\_\{k\}=\\mathcal\{Z\},\\quad R\_\{i\}\\cap R\_\{j\}=\\emptyset\\ \(i\\neq j\)\.This partition is used as a computational device for maintaining region\-level statistics, not as an assumption that the underlying factors of variation are mutually exclusive\. Individual samples may involve multiple factors, but assigning each sample to a single region enables scalable discrepancy tracking without per\-sample scoring\. Regions can be constructed in any representation space\. In our text\-to\-image instantiation, they are obtained by clustering prompt embeddings; details are provided in Sec\.[5](https://arxiv.org/html/2606.23898#S5)\. To quantify the learning signal associated with a sampled training instance, ARIA relies on a scalar*discrepancy*functionΔ​\(⋅\)\\Delta\(\\cdot\)\. At iterationtt,Δt\\Delta\_\{t\}denotes the scalar value used to update the score of the selected region, e\.g\., the minibatch\-averaged loss from samples drawn from that region\.

### 4\.2Region score initialization and tracking

For each regionRkR\_\{k\}, ARIA maintains a scoreDkD\_\{k\}intended to approximate its current mean discrepancy\.

#### Initialization\.

Before adaptive sampling begins, we estimate an initial score for each region by evaluating the discrepancy on a small random subset\. Specifically, for eachkk, we sample a batchℬk⊂Rk\\mathcal\{B\}\_\{k\}\\subset R\_\{k\}and set

Dk\(0\)=1\|ℬk\|​∑z∈ℬkΔ​\(z\)\.D\_\{k\}^\{\(0\)\}\\;=\\;\\frac\{1\}\{\|\\mathcal\{B\}\_\{k\}\|\}\\sum\_\{z\\in\\mathcal\{B\}\_\{k\}\}\\Delta\(z\)\.\(1\)This provides a coarse estimate of which regions are initially under\-aligned\.

#### Online tracking via EMA\.

Whenever regionkkis selected at training stepttand discrepancyΔt\\Delta\_\{t\}is observed, its score is updated using an exponential moving average:

Dk←\(1−α\)​Dk\+α​Δt,α∈\(0,1\)\.D\_\{k\}\\leftarrow\(1\-\\alpha\)D\_\{k\}\+\\alpha\\,\\Delta\_\{t\},\\qquad\\alpha\\in\(0,1\)\.\(2\)All other region scores remain unchanged at that step\. EMA provides a stable tracker in the presence of stochasticity and non\-stationarity\. Our theoretical analysis in subsection[4\.4](https://arxiv.org/html/2606.23898#S4.SS4)quantifies its tracking behavior\.

### 4\.3Score\-based region sampling

At each training steptt, ARIA forms a probability distribution over regions based on the current scores\{Dk\}k=1K\\\{D\_\{k\}\\\}\_\{k=1\}^\{K\}\. Formally, let

pt=𝒮​\(D1,…,DK\),p\_\{t\}=\\mathcal\{S\}\(D\_\{1\},\\dots,D\_\{K\}\),where𝒮:ℝK→ΔK−1\\mathcal\{S\}:\\mathbb\{R\}^\{K\}\\to\\Delta^\{K\-1\}maps region scores to a probability distribution over\{1,…,K\}\\\{1,\\dots,K\\\}\.

The only requirement is that regions with larger scores receive larger sampling probability\. In practice, the mapping should balance prioritization of high\-discrepancy regions with sufficient coverage of the remaining conditioning space\. Beyond this monotonicity condition, the mapping𝒮\\mathcal\{S\}can be chosen according to practical considerations\. In our experiments, we instantiate𝒮\\mathcal\{S\}using a softmax transformation of the region scores\.

Each regionRkR\_\{k\}is associated with a samplerΠk\\Pi\_\{k\}that draws training instances from that region\. In this work, we instantiateΠk\\Pi\_\{k\}as uniform sampling over elements ofRkR\_\{k\}\. More generally,Πk\\Pi\_\{k\}may implement any user\-defined sampling strategy within the region\.

Algorithm 1ARIA: Adaptive Region\-based Importance Allocation1:Regions

ℛ=\{R1,…,RK\}\\mathcal\{R\}=\\\{R\_\{1\},\\dots,R\_\{K\}\\\}with samplers

\{Πk\}\\\{\\Pi\_\{k\}\\\}, discrepancy

Δ​\(⋅\)\\Delta\(\\cdot\), EMA parameter

α∈\(0,1\)\\alpha\\in\(0,1\)\.

2:Sampling mapping

𝒮:ℝK→ΔK−1\\mathcal\{S\}:\\mathbb\{R\}^\{K\}\\to\\Delta^\{K\-1\}\(maps scores to a distribution\)\.

3:Initialization fraction

ρ∈\(0,1\)\\rho\\in\(0,1\)\.

4:\(Initialization\)For each region

kk, sample

ℬk⊂Rk\\mathcal\{B\}\_\{k\}\\subset R\_\{k\}and set

Dk\(0\)←1\|ℬk\|​∑z∈ℬkΔ​\(z\)D\_\{k\}^\{\(0\)\}\\leftarrow\\frac\{1\}\{\|\\mathcal\{B\}\_\{k\}\|\}\\sum\_\{z\\in\\mathcal\{B\}\_\{k\}\}\\Delta\(z\)\.

5:for

t=1,2,…,Tt=1,2,\\dots,Tdo

6:Form region distribution

pt←𝒮​\(D1,…,DK\)p\_\{t\}\\leftarrow\\mathcal\{S\}\(D\_\{1\},\\dots,D\_\{K\}\)\.

7:Sample region

kt∼ptk\_\{t\}\\sim p\_\{t\}\.

8:Sample instance

zt∼Πkz\_\{t\}\\sim\\Pi\_\{k\}\.

9:Compute discrepancy

Δt←Δ​\(zt\)\\Delta\_\{t\}\\leftarrow\\Delta\(z\_\{t\}\)and perform the standard learning update\.

10:Update the selected region score:

Dkt←\(1−α\)​Dkt\+α​Δt\.D\_\{k\_\{t\}\}\\leftarrow\(1\-\\alpha\)D\_\{k\_\{t\}\}\+\\alpha\\,\\Delta\_\{t\}\.
11:endfor

Givenptp\_\{t\}, a region indexkt∼ptk\_\{t\}\\sim p\_\{t\}is sampled and a training instancezt∼Πktz\_\{t\}\\sim\\Pi\_\{k\_\{t\}\}is drawn from the corresponding region\. The discrepancyΔ​\(zt\)\\Delta\(z\_\{t\}\)is then computed and used both for the learning update and for updating the EMA score of regionktk\_\{t\}\. ARIA therefore modifies only the data\-selection policy while leaving the underlying learning objective unchanged\.

### 4\.4Algorithm and tracking analysis

Sampling probabilities depend on region\-level scores, thus these must reflect the evolving discrepancy\. We show that ARIA’s EMA scores track the conditional mean discrepancy with a finite\-time guarantee whose error decomposes into three terms: an initialization term, a drift\-induced bias, and a stochastic term that scales withα\\sqrt\{\\alpha\}\. In addition, a high\-probability bound ensures stability over finite horizons\.

These guarantees support the use of EMA scores as lightweight proxies for region difficulty under bounded variance and bounded drift assumptions\.

For brevity, we present an informal statement below \(Theorem[1](https://arxiv.org/html/2606.23898#Thmtheorem1)\) and refer to Sec\. B in the Supp\. Mat\. for the full theorem and proof\.

###### Theorem 1\(Tracking Guarantee, Informal\)\.

Fix a regionssand letDs\(n\)D\_\{s\}^\{\(n\)\}denote ARIA’s EMA score after itsnn\-th update\. LetJs\(n\):=𝔼​\[Δs\(n\)∣ℱn−1\]J\_\{s\}^\{\(n\)\}:=\\mathbb\{E\}\[\\Delta\_\{s\}^\{\(n\)\}\\mid\\mathcal\{F\}\_\{n\-1\}\]denote the expected discrepancy for regionss, where the expectation is taken over the within\-region sampler \(uniform overRsR\_\{s\}in this work\)\. Define the tracking erroren:=Ds\(n\)−Js\(n\)e\_\{n\}:=D\_\{s\}^\{\(n\)\}\-J\_\{s\}^\{\(n\)\}\. Assume bounded conditional varianceσ2\\sigma^\{2\}and bounded discounted driftVV\. Then for alln≥1n\\geq 1,

𝔼​\[en2\]≤\(1−α\)n​\|e0\|\+\(1−α\)​V\+σ​α\.\\sqrt\{\\mathbb\{E\}\[e\_\{n\}^\{2\}\]\}\\;\\leq\\;\(1\-\\alpha\)^\{n\}\|e\_\{0\}\|\+\(1\-\\alpha\)V\+\\sigma\\sqrt\{\\alpha\}\.
The error decomposes into an exponentially decaying initialization term, a drift\-induced bias term, and a stochastic term scaling withα\\sqrt\{\\alpha\}\. A uniform high\-probability bound further ensures stability over finite horizons\.

## 5Experiments

In this section, we evaluate ARIA across multiple student architectures and data regimes\. Following\[[26](https://arxiv.org/html/2606.23898#bib.bib17)\], we use the exponential schedulep​\(t\)=e−λ​\(1−tT\)p\(t\)=e^\{\-\\lambda\\left\(1\-\\frac\{t\}\{T\}\\right\)\}\. As the discrepancy signalΔ​\(⋅\)\\Delta\(\\cdot\), we use the output\-level distillation lossℒout\\mathcal\{L\}\_\{\\mathrm\{out\}\}\.

For each regionRkR\_\{k\}in the conditioning space, ARIA tracks the expected output\-level loss

Jk:=𝔼c∈Rk,xt​\[ℒout​\(xt,c\)\],J\_\{k\}:=\\mathbb\{E\}\_\{c\\in R\_\{k\},\\,x\_\{t\}\}\\left\[\\mathcal\{L\}\_\{\\mathrm\{out\}\}\(x\_\{t\},c\)\\right\],whereℒout​\(xt,c\)\\mathcal\{L\}\_\{\\mathrm\{out\}\}\(x\_\{t\},c\)denotes the denoising loss evaluated for conditionccand noisy latentxtx\_\{t\}sampled from the forward diffusion process\. In practice, this quantity is estimated online using an exponential moving average over sampled instances fromRkR\_\{k\}\. Specifically, the score of regionRkR\_\{k\}is updated whenever a conditionc∈Rkc\\in R\_\{k\}is selected as a replacement during training, using the observed distillation loss\. We use a softmax mapping for𝒮\\mathcal\{S\}and set the EMA parameter toα=0\.1\\alpha=0\.1\.

To construct the regions, we extract prompt embeddings from the auxiliary corpus using a pretrained CLIP\[[40](https://arxiv.org/html/2606.23898#bib.bib41)\]text encoder and apply k\-means clustering in the embedding space\. The number of clusters is selected using the Silhouette score to balance inter and intra\-cluster cohesion\. Further details on the clustering procedure are provided in Sec\.[D](https://arxiv.org/html/2606.23898#A4)of the Supp\. Mat\., while sensitivity to these design choices is analyzed in Sec\.[E](https://arxiv.org/html/2606.23898#A5)and Sec\.[C\.4](https://arxiv.org/html/2606.23898#A3.SS4)of the Supp\. Mat\.

#### Roadmap\.

Our evaluation is structured around two complementary scenarios\. First, in subsection[5\.1](https://arxiv.org/html/2606.23898#S5.SS1)we describe the experimental setup\. Next, in subsection[5\.2](https://arxiv.org/html/2606.23898#S5.SS2)we consider a balanced setting that reproduces the standard knowledge distillation setup as in\[[25](https://arxiv.org/html/2606.23898#bib.bib1),[26](https://arxiv.org/html/2606.23898#bib.bib17)\]to assess convergence and final performance\. Finally, in subsection[5\.3](https://arxiv.org/html/2606.23898#S5.SS3)we construct an imbalanced concept\-removal setting to test robustness under distributional gaps\. Additional ablation studies of our method are provided in Section[E](https://arxiv.org/html/2606.23898#A5)\.

### 5\.1Experimental Setup

We evaluate ARIA across three teacher models: Stable Diffusion v1\.4\[[42](https://arxiv.org/html/2606.23898#bib.bib5)\], Stable Diffusion v2\.1, and SDXL\[[39](https://arxiv.org/html/2606.23898#bib.bib36)\]\. We distill these teachers into multiple student architectures across several settings\. For Stable Diffusion v1\.4, we consider \(i\) BK\-SDM students\[[25](https://arxiv.org/html/2606.23898#bib.bib1)\], trained for 75K steps, and \(ii\) channel\-pruned U\-Net variants\[[26](https://arxiv.org/html/2606.23898#bib.bib17)\], trained for 300K steps111RC\[[26](https://arxiv.org/html/2606.23898#bib.bib17)\]reports results up to400400K iterations\. In our experiments, channel\-pruned models saturate around300300K, with negligible gains from further training\.\. For Stable Diffusion v2\.1, we distill into two block\-pruned students, BK\-SDM\-v2\-S and BK\-SDM\-v2\-T, trained for 25K steps\. For SDXL, we use KOALA\-700M\[[29](https://arxiv.org/html/2606.23898#bib.bib35)\]as the student and train for 100K steps\. Table[1](https://arxiv.org/html/2606.23898#S5.T1)reports the sizes of the variants in each family\.

For each model, we report the mean and standard deviation computed over three random seeds\.

#### Training Dataset\.

For the distillation of Stable Diffusion v2\.1, Stable Diffusion v1\.4, and SDXL, we follow\[[26](https://arxiv.org/html/2606.23898#bib.bib17)\]and construct the paired training set from LAION\-Aesthetics V2\[[45](https://arxiv.org/html/2606.23898#bib.bib38)\], a subset of the LAION corpus\[[46](https://arxiv.org/html/2606.23898#bib.bib39)\]\. Specifically, we sample212212K image–text pairs from the LAION\-Aesthetics V2 6\.5\+ subset and use them as the training set\. To simulate the image\-free setting, we discard the original images and retain only the associated text prompts; synthetic images are then generated from these prompts using the corresponding teacher model\. Separately, we sample2020M text prompts from the full LAION corpus to construct the auxiliary text pool used by our method\.

For the filtered setting described below, we remove training samples belonging to selected semantic categories, namely animals, vehicles, and food, using BLIP\[[31](https://arxiv.org/html/2606.23898#bib.bib42)\]together with keyword\-based filtering\. After filtering, the paired training set contains approximately168168K image–text pairs\.

#### Evaluation\.

Following\[[25](https://arxiv.org/html/2606.23898#bib.bib1),[26](https://arxiv.org/html/2606.23898#bib.bib17)\], we use3030K image–text pairs sampled from the MS\-COCO\[[33](https://arxiv.org/html/2606.23898#bib.bib40)\]validation set\. Each image is associated with five human\-annotated captions; following\[[25](https://arxiv.org/html/2606.23898#bib.bib1)\], we use a single caption from\[[25](https://arxiv.org/html/2606.23898#bib.bib1)\]for evaluation\. We report the following evaluation metrics: Fréchet Inception Distance \(FID,↓\\downarrow\), Inception Score \(IS,↑\\uparrow\), and CLIP score \(↑\\uparrow\)\. We use the Inception\-v3 model for computing FID and IS, and the Vit\-g/14 model for calculating CLIP score\. We verify that none of the MS\-COCO validation captions used for evaluation appear in the LAION auxiliary prompt pool\.

#### Computational overhead\.

The only additional preprocessing cost in ARIA stems from embedding and clustering the auxiliary text corpus, a process that handles 20M text prompts using only5\.35\.3A100100GPU hours, a small fraction of the400400–600600GPU hours required for student training\. The computational overhead across student models is summarized in Fig\.[3](https://arxiv.org/html/2606.23898#S5.F3), and during training ARIA adds negligible cost, as region scores are updated from distillation losses via lightweight EMA operations\.

![Refer to caption](https://arxiv.org/html/2606.23898v1/figures/efficiency_tiny.png)Figure 3:Training efficiency and convergence comparison\.\(a\) Relative computational overhead of ARIA compared to the base model, showing negligible additional cost\. \(b\) FID across training checkpoints\. \(c\) CLIP score across training checkpoints\. \(d\) IS across training checkpoints\. ARIA converges faster than RC, achieving lower FID and higher CLIP at earlier stages of training\.

### 5\.2Balanced Distillation Setting

Table 1:Balanced distillation setting\.Students are trained using212212K cached LAION pairs and an auxiliary text pool𝒞aux\\mathcal\{C\}\_\{\\text\{aux\}\}of2020M prompts\. ARIA improves over Random Conditioning across most metrics and student architectures\.We start by evaluating ARIA on the full212212K cached LAION paired dataset and the2020M\-prompt𝒞aux\\mathcal\{C\}\_\{\\text\{aux\}\}\. This regime evaluates performance under a broadly representative and semantically diverse prompt distribution\.

In this setting, we study: \(i\) final generative quality, \(ii\) convergence speed, and \(iii\) robustness across student architectures\. Since all training components remain unchanged except prompt selection, any observed improvements can be attributed to effective allocation of training effort across the conditioning space\.

Table 2:Filtered concept\-removal setting on SD 1\.4\. We report performance on seen and unseen concepts\. All students are trained without images from the removed animal, vehicle, and food categories\. ARIA generally improves over RC across multiple metrics, with the largest and most consistent gains naturally appearing in the unseen regime\.Table[1](https://arxiv.org/html/2606.23898#S5.T1)reports results in the balanced regime\. The teacher model \(Stable Diffusion v1\.41\.4\) is fixed and therefore shown only once for reference\. Across student architectures, ARIA improves over RC in most cases, with the clearest gains appearing in FID and CLIP, while maintaining competitive or improved IS\. Notably, the relative gains are more pronounced for smaller student models, indicating that ARIA is particularly effective in capacity\-constrained settings\.

To further evaluate ARIA beyond the SD 1\.4 setting, we distill SDXL\[[39](https://arxiv.org/html/2606.23898#bib.bib36)\]into KOALA\-700M\[[29](https://arxiv.org/html/2606.23898#bib.bib35)\]\. We compare ARIA against RC\[[26](https://arxiv.org/html/2606.23898#bib.bib17)\]and a static region\-sampling policy that uses the same regions as ARIA but samples them uniformly\. This baseline isolates the effect of online adaptive allocation from the effect of region construction alone\. As shown in Table[3](https://arxiv.org/html/2606.23898#S5.T3), ARIA outperforms both RC and static region sampling, indicating that the gains come from discrepancy\-aware allocation rather than clustering alone\.

Table 3:Distilling KOALA\-700M in the unfiltered setting\.RC: Random Conditioning\.Static: uniform sampling over clusters followed by uniform sampling within the selected cluster\.One phenomenon we observe is that ARIA significantly accelerates convergence for models trained from random initialization \(the C\-based architectures\)\. For example, as shown in Fig\.[3](https://arxiv.org/html/2606.23898#S5.F3), the C\-Tiny student trained with ARIA reaches comparable FID and CLIP scores after approximately175175K iterations, which corresponds to60%60\\%of the full training budget, compared to a model trained with RC for300300K iterations\. Similar trends are observed across other C\-based architectures; additional results are provided in Sec\.[C\.4](https://arxiv.org/html/2606.23898#A3.SS4)of the Supp\. Mat\.

Subfigures \(b\)–\(d\) plot FID, CLIP, and IS across training checkpoints for the C\-Tiny architecture\. ARIA achieves lower FID and higher CLIP earlier, demonstrating faster convergence than RC\. Importantly, these improvements are obtained with negligible computational overhead, as illustrated in Fig\.[3](https://arxiv.org/html/2606.23898#S5.F3)\(a\)\.

### 5\.3Filtered Concept\-Removal Setting

Tables[2](https://arxiv.org/html/2606.23898#S5.T2),[4](https://arxiv.org/html/2606.23898#S5.T4)report results in a more challenging regime where samples related to*animals, vehicles, and food*are removed from the cached image–text pairs\. Students are therefore trained using168168K cached pairs while retaining the same auxiliary prompt pool𝒞aux\\mathcal\{C\}\_\{\\text\{aux\}\}of size2020M\. For evaluation, we split the3030K image–text pairs sampled from the MS\-COCO validation set into two subsets: \(i\)20,11720\{,\}117samples that do not contain*animals, vehicles, or food*concepts \(seen\), and \(ii\)9,8839\{,\}883samples that contain only these removed concepts \(unseen\)\. In this setting, the gap between ARIA and RC widens, particularly on the unseen subset corresponding to the removed semantic categories\. Intuitively, when the cached image–text pairs lack entire semantic categories, parts of the conditioning space become poorly covered by training data\. Prompts from such regions often produce larger disagreement between the student and the teacher, since the student has not observed sufficient examples of these concepts\. Since ARIA tracks the student–teacher discrepancy at the region level, these regions naturally accumulate higher discrepancy scores and are therefore sampled more frequently\.

This behavior allows ARIA to focus training effort on parts of the conditioning space that remain under\-distilled\. In contrast, uniform sampling allocates the same budget to both easy and difficult regions, which can lead to wasted updates on already well\-aligned concepts\. As a result, ARIA yields noticeably larger improvements in this filtered setting than in the balanced regime, highlighting the benefit of adaptive allocation when the training distribution is imbalanced or systematically missing semantic categories\.

Table 4:Filtered concept\-removal setting on SD 2\.1 and SDXL; unseen\-concept results only\. ARIA improves over RC across backbones and student architectures\.For the SD 2\.1 and SDXL experiments in Table[4](https://arxiv.org/html/2606.23898#S5.T4), we report only unseen\-concept results due to space constraints\. In the seen regime, ARIA remains comparable to RC and achieves slight gains across most metrics\.

## 6Conclusion and Future Work

Conclusion\.In this work, we introduced*ARIA*, a framework for adaptive importance allocation in large\-scale training regimes\. ARIA operates at the region level, dynamically reallocating training effort toward coarse regions that exhibit higher model discrepancy, without modifying the underlying optimization objective\. We instantiate ARIA for text\-to\-image diffusion distillation, where large conditioning spaces make uniform sampling inefficient\. We provide a theoretical analysis of the proposed scoring mechanism, establishing finite\-time tracking guarantees for region\-level discrepancy estimates\. Empirically, ARIA improves over uniform random sampling in most evaluated settings, with the strongest gains appearing under imbalanced or underrepresented conditioning regimes, and achieves faster convergence through efficient allocation\.

Future Work\.Two directions emerge from this work\. First, adaptive importance allocation is increasingly relevant in large\-scale learning systems, where data pools are massive and unevenly informative\. Extending ARIA to domains such as large language models and multimodal generative systems is a natural direction\. Second, our region\-based formulation opens the door to alternative allocation strategies\. Beyond discrepancy\-driven allocation, future work may explore uncertainty\-aware or dynamically learned region constructions, potentially leading to more principled and robust training mechanisms\.

## 7Limitations

In our instantiation, ARIA constructs regions using fixed clustering over CLIP text embeddings\. Although effective in our experiments, this choice may not capture fine\-grained prompt differences, compositional structure, or task\-specific notions of similarity\. ARIA is not tied to CLIP or to fixed clusters, and future work could explore alternative encoders, task\-specific or model\-derived representations, soft or hierarchical regions, and dynamically updated partitions\. Second, ARIA improves the allocation of a limited training budget, but it does not remove the cost of constructing the cached image set used for distillation\. Following prior work, we assume a fixed cache of generated images and focus on selecting auxiliary text conditions more effectively\. Reducing or eliminating this dependence on cached images remains an important direction\. Finally, ARIA introduces practical design choices, including the embedding model, number of regions, EMA parameter, and score\-to\-probability mapping\. Our ablations suggest that ARIA is stable, but different datasets or teacher–student pairs may benefit from different settings\. More principled or learned allocation rules could ease deployment\.

## Acknowledgments and Disclosure of Funding

This research was funded by the Ministry of Science, Research and the Arts Baden\-Wuerttemberg in the Artificial Intelligence Software Academy \(AISA\)\. We also acknowledge the support of the Stuttgart Center for Simulation Science \(SimTech\) and thank the International Max Planck Research School for Intelligent Systems \(IMPRS\-IS\) for support\. L\. Mualem was supported by a postdoctoral scholarship from the Planning and Budgeting Committee \(PBC\) of the Council for Higher Education in Israel\. The authors gratefully acknowledge the computing time provided on the high\-performance computer HoreKa by the National High\-Performance Computing Center at KIT \(NHR@KIT\)\. This center is jointly supported by the Federal Ministry of Education and Research and the Ministry of Science, Research and the Arts of Baden\-Württemberg, as part of the National High\-Performance Computing \(NHR\) joint funding program \(https://www\.nhr\-verein\.de/en/our\-partners\)\. HoreKa is partly funded by the German Research Foundation \(DFG\)\.

## References

- \[1\]\(2007\)K\-means\+\+ the advantages of careful seeding\.InProceedings of the eighteenth annual ACM\-SIAM symposium on Discrete algorithms,pp\. 1027–1035\.Cited by:[§D\.1](https://arxiv.org/html/2606.23898#A4.SS1.SSS0.Px1.p1.5)\.
- \[2\]S\. Batifol, A\. Blattmann, F\. Boesel, S\. Consul, C\. Diagne, T\. Dockhorn, J\. English, Z\. English, P\. Esser, S\. Kulal,et al\.\(2025\)Flux\. 1 kontext: flow matching for in\-context image generation and editing in latent space\.arXiv e\-prints,pp\. arXiv–2506\.Cited by:[§1](https://arxiv.org/html/2606.23898#S1.p1.1)\.
- \[3\]A\. Blattmann, T\. Dockhorn, S\. Kulal, D\. Mendelevitch, M\. Kilian, D\. Lorenz, Y\. Levi, Z\. English, V\. Voleti, A\. Letts,et al\.\(2023\)Stable video diffusion: scaling latent video diffusion models to large datasets\.arXiv preprint arXiv:2311\.15127\.Cited by:[§1](https://arxiv.org/html/2606.23898#S1.p1.1)\.
- \[4\]D\. Bouneffouf, I\. Rish, and C\. Aggarwal\(2020\)Survey on applications of multi\-armed and contextual bandits\.In2020 IEEE congress on evolutionary computation \(CEC\),pp\. 1–8\.Cited by:[Appendix A](https://arxiv.org/html/2606.23898#A1.SS0.SSS0.Px1.p1.1)\.
- \[5\]V\. Braverman, D\. Feldman, H\. Lang, A\. Statman, and S\. Zhou\(2021\)Efficient coreset constructions via sensitivity sampling\.InAsian Conference on Machine Learning,pp\. 948–963\.Cited by:[§D\.1](https://arxiv.org/html/2606.23898#A4.SS1.SSS0.Px2.p1.2),[§D\.1](https://arxiv.org/html/2606.23898#A4.SS1.SSS0.Px2.p2.6),[§D\.1](https://arxiv.org/html/2606.23898#A4.SS1.SSS0.Px2.p3.1)\.
- \[6\]J\. Carvalho, A\. T\. Le, P\. Kicki, D\. Koert, and J\. Peters\(2025\)Motion planning diffusion: learning and adapting robot motion planning with diffusion models\.IEEE Transactions on Robotics\.Cited by:[§1](https://arxiv.org/html/2606.23898#S1.p1.1)\.
- \[7\]J\. Chen, Y\. Zhao, J\. Yu, R\. Chu, J\. Chen, S\. Yang, X\. Wang, Y\. Pan, D\. Zhou, H\. Ling,et al\.\(2025\)Sana\-video: efficient video generation with block linear diffusion transformer\.arXiv preprint arXiv:2509\.24695\.Cited by:[§2](https://arxiv.org/html/2606.23898#S2.p1.1)\.
- \[8\]P\. Chen, S\. Liu, H\. Zhao, and J\. Jia\(2021\)Distilling knowledge via knowledge review\.InProceedings of the IEEE/CVF conference on computer vision and pattern recognition,pp\. 5008–5017\.Cited by:[§1](https://arxiv.org/html/2606.23898#S1.p2.1),[§2](https://arxiv.org/html/2606.23898#S2.p2.1)\.
- \[9\]C\. Chi, Z\. Xu, S\. Feng, E\. Cousineau, Y\. Du, B\. Burchfiel, R\. Tedrake, and S\. Song\(2025\)Diffusion policy: visuomotor policy learning via action diffusion\.The International Journal of Robotics Research44\(10\-11\),pp\. 1684–1704\.Cited by:[§1](https://arxiv.org/html/2606.23898#S1.p1.1)\.
- \[10\]V\. Cohen\-Addad, K\. Green Larsen, D\. Saulpic, C\. Schwiegelshohn, and O\. A\. Sheikh\-Omar\(2022\)Improved coresets for euclideankk\-means\.Advances in Neural Information Processing Systems35,pp\. 2679–2694\.Cited by:[Appendix A](https://arxiv.org/html/2606.23898#A1.SS0.SSS0.Px2.p1.14)\.
- \[11\]T\. Dao, T\. H\. Nguyen, T\. Le, D\. Vu, K\. Nguyen, C\. Pham, and A\. Tran\(2024\)Swiftbrush v2: make your one\-step diffusion model better than its teacher\.InEuropean Conference on Computer Vision,pp\. 176–192\.Cited by:[§1](https://arxiv.org/html/2606.23898#S1.p2.1),[§2](https://arxiv.org/html/2606.23898#S2.p3.1)\.
- \[12\]P\. Dhariwal and A\. Nichol\(2021\)Diffusion models beat gans on image synthesis\.Advances in neural information processing systems34,pp\. 8780–8794\.Cited by:[§1](https://arxiv.org/html/2606.23898#S1.p1.1)\.
- \[13\]D\. Feldman, A\. Fiat, H\. Kaplan, and K\. Nissim\(2009\)Private coresets\.InProceedings of the forty\-first annual ACM symposium on Theory of computing,pp\. 361–370\.Cited by:[Appendix A](https://arxiv.org/html/2606.23898#A1.SS0.SSS0.Px2.p1.14)\.
- \[14\]B\. Guo, S\. Tang, C\. Zeng, and Z\. Shen\(2025\)Mosaicdiff: training\-free structural pruning for diffusion model acceleration reflecting pretraining dynamics\.InProceedings of the IEEE/CVF International Conference on Computer Vision,pp\. 1655–1664\.Cited by:[§2](https://arxiv.org/html/2606.23898#S2.p1.1)\.
- \[15\]A\. Gupta, L\. Yu, K\. Sohn, X\. Gu, M\. Hahn, F\. Li, I\. Essa, L\. Jiang, and J\. Lezama\(2024\)Photorealistic video generation with diffusion models\.InEuropean Conference on Computer Vision,pp\. 393–411\.Cited by:[§1](https://arxiv.org/html/2606.23898#S1.p1.1)\.
- \[16\]Z\. Hao, J\. Guo, D\. Jia, K\. Han, Y\. Tang, C\. Zhang, H\. Hu, and Y\. Wang\(2022\)Learning efficient vision transformers via fine\-grained manifold distillation\.Advances in Neural Information Processing Systems35,pp\. 9164–9175\.Cited by:[§2](https://arxiv.org/html/2606.23898#S2.p2.1)\.
- \[17\]S\. Har\-Peled and A\. Kushal\(2005\)Smaller coresets for k\-median and k\-means clustering\.InProceedings of the twenty\-first annual symposium on Computational geometry,pp\. 126–134\.Cited by:[Appendix A](https://arxiv.org/html/2606.23898#A1.SS0.SSS0.Px2.p1.14)\.
- \[18\]S\. Har\-Peled and S\. Mazumdar\(2004\)On coresets for k\-means and k\-median clustering\.InProceedings of the thirty\-sixth annual ACM symposium on Theory of computing,pp\. 291–300\.Cited by:[Appendix A](https://arxiv.org/html/2606.23898#A1.SS0.SSS0.Px2.p1.14)\.
- \[19\]G\. Hinton\(2014\)Distilling the knowledge in a neural network\.InDeep Learning and Representation Learning Workshop in Conjunction with NIPS,Cited by:[§1](https://arxiv.org/html/2606.23898#S1.p2.1),[§2](https://arxiv.org/html/2606.23898#S2.p2.1)\.
- \[20\]J\. Ho, W\. Chan, C\. Saharia, J\. Whang, R\. Gao, A\. Gritsenko, D\. P\. Kingma, B\. Poole, M\. Norouzi, D\. J\. Fleet,et al\.\(2022\)Imagen video: high definition video generation with diffusion models\.arXiv preprint arXiv:2210\.02303\.Cited by:[§1](https://arxiv.org/html/2606.23898#S1.p1.1)\.
- \[21\]J\. Ho, A\. Jain, and P\. Abbeel\(2020\)Denoising diffusion probabilistic models\.Advances in neural information processing systems33,pp\. 6840–6851\.Cited by:[§1](https://arxiv.org/html/2606.23898#S1.p1.1)\.
- \[22\]L\. Huang, S\. Jiang, and N\. Vishnoi\(2019\)Coresets for clustering with fairness constraints\.Advances in neural information processing systems32\.Cited by:[Appendix A](https://arxiv.org/html/2606.23898#A1.SS0.SSS0.Px2.p1.14)\.
- \[23\]X\. Jiao, Y\. Yin, L\. Shang, X\. Jiang, X\. Chen, L\. Li, F\. Wang, and Q\. Liu\(2020\)Tinybert: distilling bert for natural language understanding\.InFindings of the association for computational linguistics: EMNLP 2020,pp\. 4163–4174\.Cited by:[§2](https://arxiv.org/html/2606.23898#S2.p2.1)\.
- \[24\]I\. Jubran, A\. Maalouf, and D\. Feldman\(2019\)Introduction to coresets: accurate coresets\.arXiv preprint arXiv:1910\.08707\.Cited by:[Appendix A](https://arxiv.org/html/2606.23898#A1.SS0.SSS0.Px2.p1.14)\.
- \[25\]B\. Kim, H\. Song, T\. Castells, and S\. Choi\(2024\)Bk\-sdm: a lightweight, fast, and cheap version of stable diffusion\.InEuropean Conference on Computer Vision,pp\. 381–399\.Cited by:[§2](https://arxiv.org/html/2606.23898#S2.p1.1),[§2](https://arxiv.org/html/2606.23898#S2.p2.1),[§3](https://arxiv.org/html/2606.23898#S3.p1.5),[§5](https://arxiv.org/html/2606.23898#S5.SS0.SSS0.Px1.p1.1),[§5\.1](https://arxiv.org/html/2606.23898#S5.SS1.SSS0.Px2.p1.4),[§5\.1](https://arxiv.org/html/2606.23898#S5.SS1.p1.1)\.
- \[26\]D\. Kim, S\. Park, G\. Han, S\. W\. Kim, and P\. H\. Seo\(2025\)Random conditioning with distillation for data\-efficient diffusion model compression\.In2025 IEEE/CVF Conference on Computer Vision and Pattern Recognition \(CVPR\),pp\. 18607–18618\.Cited by:[§C\.1](https://arxiv.org/html/2606.23898#A3.SS1.p1.1),[§1](https://arxiv.org/html/2606.23898#S1.p2.1),[§1](https://arxiv.org/html/2606.23898#S1.p4.1),[§2](https://arxiv.org/html/2606.23898#S2.p1.1),[§2](https://arxiv.org/html/2606.23898#S2.p2.1),[§3](https://arxiv.org/html/2606.23898#S3.p1.5),[§3](https://arxiv.org/html/2606.23898#S3.p8.6),[§5](https://arxiv.org/html/2606.23898#S5.SS0.SSS0.Px1.p1.1),[§5\.1](https://arxiv.org/html/2606.23898#S5.SS1.SSS0.Px1.p1.2),[§5\.1](https://arxiv.org/html/2606.23898#S5.SS1.SSS0.Px2.p1.4),[§5\.1](https://arxiv.org/html/2606.23898#S5.SS1.p1.1),[§5\.2](https://arxiv.org/html/2606.23898#S5.SS2.p4.1),[§5](https://arxiv.org/html/2606.23898#S5.p1.3),[footnote 1](https://arxiv.org/html/2606.23898#footnote1)\.
- \[27\]Z\. Kong, W\. Ping, J\. Huang, K\. Zhao, and B\. Catanzaro\(2020\)Diffwave: a versatile diffusion model for audio synthesis\.arXiv preprint arXiv:2009\.09761\.Cited by:[§1](https://arxiv.org/html/2606.23898#S1.p1.1)\.
- \[28\]Y\. D\. Kwon, R\. Li, S\. Li, D\. Li, S\. Bhattacharya, and S\. I\. Venieris\(2025\)HierarchicalPrune: position\-aware compression for large\-scale diffusion models\.arXiv preprint arXiv:2508\.04663\.Cited by:[§2](https://arxiv.org/html/2606.23898#S2.p1.1)\.
- \[29\]Y\. Lee, K\. Park, Y\. Cho, Y\. Lee, and S\. J\. Hwang\(2024\)Koala: empirical lessons toward memory\-efficient and fast diffusion models for text\-to\-image synthesis\.Advances in Neural Information Processing Systems37,pp\. 51597–51633\.Cited by:[§2](https://arxiv.org/html/2606.23898#S2.p1.1),[§5\.1](https://arxiv.org/html/2606.23898#S5.SS1.p1.1),[§5\.2](https://arxiv.org/html/2606.23898#S5.SS2.p4.1)\.
- \[30\]G\. Li, X\. Li, Y\. Wang, S\. Zhang, Y\. Wu, and D\. Liang\(2022\)Knowledge distillation for object detection via rank mimicking and prediction\-guided feature imitation\.InProceedings of the AAAI conference on artificial intelligence,pp\. 1306–1313\.Cited by:[§1](https://arxiv.org/html/2606.23898#S1.p2.1),[§2](https://arxiv.org/html/2606.23898#S2.p2.1)\.
- \[31\]J\. Li, D\. Li, C\. Xiong, and S\. Hoi\(2022\)Blip: bootstrapping language\-image pre\-training for unified vision\-language understanding and generation\.InInternational conference on machine learning,pp\. 12888–12900\.Cited by:[§5\.1](https://arxiv.org/html/2606.23898#S5.SS1.SSS0.Px1.p2.1)\.
- \[32\]Y\. Li, H\. Wang, Q\. Jin, J\. Hu, P\. Chemerys, Y\. Fu, Y\. Wang, S\. Tulyakov, and J\. Ren\(2023\)Snapfusion: text\-to\-image diffusion model on mobile devices within two seconds\.Advances in Neural Information Processing Systems36,pp\. 20662–20678\.Cited by:[§2](https://arxiv.org/html/2606.23898#S2.p1.1)\.
- \[33\]T\. Lin, M\. Maire, S\. Belongie, J\. Hays, P\. Perona, D\. Ramanan, P\. Dollár, and C\. L\. Zitnick\(2014\)Microsoft coco: common objects in context\.InEuropean conference on computer vision,pp\. 740–755\.Cited by:[§5\.1](https://arxiv.org/html/2606.23898#S5.SS1.SSS0.Px2.p1.4)\.
- \[34\]H\. Liu, Z\. Chen, Y\. Yuan, X\. Mei, X\. Liu, D\. Mandic, W\. Wang, and M\. D\. Plumbley\(2023\)Audioldm: text\-to\-audio generation with latent diffusion models\.arXiv preprint arXiv:2301\.12503\.Cited by:[§1](https://arxiv.org/html/2606.23898#S1.p1.1)\.
- \[35\]C\. Lu, Y\. Zhou, F\. Bao, J\. Chen, C\. Li, and J\. Zhu\(2022\)Dpm\-solver: a fast ode solver for diffusion probabilistic model sampling in around 10 steps\.Advances in neural information processing systems35,pp\. 5775–5787\.Cited by:[§2](https://arxiv.org/html/2606.23898#S2.p3.1)\.
- \[36\]W\. Luo, Z\. Huang, Z\. Geng, J\. Z\. Kolter, and G\. Qi\(2024\)One\-step diffusion distillation through score implicit matching\.Advances in Neural Information Processing Systems37,pp\. 115377–115408\.Cited by:[§1](https://arxiv.org/html/2606.23898#S1.p2.1),[§2](https://arxiv.org/html/2606.23898#S2.p3.1)\.
- \[37\]T\. H\. Nguyen and A\. Tran\(2024\)Swiftbrush: one\-step text\-to\-image diffusion model with variational score distillation\.InProceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition,pp\. 7807–7816\.Cited by:[§1](https://arxiv.org/html/2606.23898#S1.p2.1),[§2](https://arxiv.org/html/2606.23898#S2.p3.1)\.
- \[38\]A\. Nichol, P\. Dhariwal, A\. Ramesh, P\. Shyam, P\. Mishkin, B\. McGrew, I\. Sutskever, and M\. Chen\(2021\)Glide: towards photorealistic image generation and editing with text\-guided diffusion models\.arXiv preprint arXiv:2112\.10741\.Cited by:[§1](https://arxiv.org/html/2606.23898#S1.p1.1)\.
- \[39\]D\. Podell, Z\. English, K\. Lacey, A\. Blattmann, T\. Dockhorn, J\. Müller, J\. Penna, and R\. Rombach\(2023\)SDXL: improving latent diffusion models for high\-resolution image synthesis\.InThe Twelfth International Conference on Learning Representations,Cited by:[§2](https://arxiv.org/html/2606.23898#S2.p1.1),[§5\.1](https://arxiv.org/html/2606.23898#S5.SS1.p1.1),[§5\.2](https://arxiv.org/html/2606.23898#S5.SS2.p4.1)\.
- \[40\]A\. Radford, J\. W\. Kim, C\. Hallacy, A\. Ramesh, G\. Goh, S\. Agarwal, G\. Sastry, A\. Askell, P\. Mishkin, J\. Clark,et al\.\(2021\)Learning transferable visual models from natural language supervision\.InInternational conference on machine learning,pp\. 8748–8763\.Cited by:[§5](https://arxiv.org/html/2606.23898#S5.p3.1)\.
- \[41\]R\. Rombach, A\. Blattmann, D\. Lorenz, P\. Esser, and B\. Ommer\(2022\-06\)High\-resolution image synthesis with latent diffusion models\.InProceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition \(CVPR\),pp\. 10684–10695\.Cited by:[§1](https://arxiv.org/html/2606.23898#S1.p1.1),[§2](https://arxiv.org/html/2606.23898#S2.p1.1)\.
- \[42\]R\. Rombach, A\. Blattmann, D\. Lorenz, P\. Esser, and B\. Ommer\(2022\)High\-resolution image synthesis with latent diffusion models\.InProceedings of the IEEE/CVF conference on computer vision and pattern recognition,pp\. 10684–10695\.Cited by:[§1](https://arxiv.org/html/2606.23898#S1.p1.1),[§5\.1](https://arxiv.org/html/2606.23898#S5.SS1.p1.1)\.
- \[43\]T\. Salimans and J\. Ho\(2022\)Progressive distillation for fast sampling of diffusion models\.arXiv preprint arXiv:2202\.00512\.Cited by:[§2](https://arxiv.org/html/2606.23898#S2.p3.1)\.
- \[44\]A\. Sauer, D\. Lorenz, A\. Blattmann, and R\. Rombach\(2024\)Adversarial diffusion distillation\.InEuropean Conference on Computer Vision,pp\. 87–103\.Cited by:[§1](https://arxiv.org/html/2606.23898#S1.p2.1),[§1](https://arxiv.org/html/2606.23898#S1.p3.1)\.
- \[45\]C\. Schuhmann and R\. Beaumont\(2022\)LAION\-aesthetics\.Cited by:[§5\.1](https://arxiv.org/html/2606.23898#S5.SS1.SSS0.Px1.p1.2)\.
- \[46\]C\. Schuhmann, R\. Kaczmarczyk, A\. Komatsuzaki, A\. Katta, R\. Vencu, R\. Beaumont, J\. Jitsev, T\. Coombes, and C\. Mullis\(2021\)LAION\-400m: open dataset of clip\-filtered 400 million image\-text pairs\.InNeurIPS Workshop Datacentric AI,Cited by:[§5\.1](https://arxiv.org/html/2606.23898#S5.SS1.SSS0.Px1.p1.2)\.
- \[47\]K\. R\. Shahapure and C\. Nicholas\(2020\)Cluster quality analysis using silhouette score\.In2020 IEEE 7th international conference on data science and advanced analytics \(DSAA\),pp\. 747–748\.Cited by:[§D\.1](https://arxiv.org/html/2606.23898#A4.SS1.p2.9)\.
- \[48\]Y\. Shang, Z\. Yuan, B\. Xie, B\. Wu, and Y\. Yan\(2023\)Post\-training quantization on diffusion models\.InProceedings of the IEEE/CVF conference on computer vision and pattern recognition,pp\. 1972–1981\.Cited by:[§2](https://arxiv.org/html/2606.23898#S2.p1.1)\.
- \[49\]A\. Slivkins\(2019\)Introduction to multi\-armed bandits\.Foundations and Trends® in Machine Learning12\(1\-2\),pp\. 1–286\.Cited by:[Appendix A](https://arxiv.org/html/2606.23898#A1.SS0.SSS0.Px1.p1.1)\.
- \[50\]Y\. Song, P\. Dhariwal, M\. Chen, and I\. Sutskever\(2023\)Consistency models\.InProceedings of the 40th International Conference on Machine Learning,pp\. 32211–32252\.Cited by:[§2](https://arxiv.org/html/2606.23898#S2.p3.1)\.
- \[51\]Y\. Song, J\. Sohl\-Dickstein, D\. P\. Kingma, A\. Kumar, S\. Ermon, and B\. Poole\(2020\)Score\-based generative modeling through stochastic differential equations\.arXiv preprint arXiv:2011\.13456\.Cited by:[§1](https://arxiv.org/html/2606.23898#S1.p1.1)\.
- \[52\]S\. Sun, Y\. Cheng, Z\. Gan, and J\. Liu\(2019\)Patient knowledge distillation for bert model compression\.InProceedings of the 2019 conference on empirical methods in natural language processing and the 9th international joint conference on natural language processing \(EMNLP\-IJCNLP\),pp\. 4323–4332\.Cited by:[§2](https://arxiv.org/html/2606.23898#S2.p2.1)\.
- \[53\]X\. Tang, J\. Li, N\. Du, and S\. Xie\(2025\)Adapting to non\-stationary environments: multi\-armed bandit enhanced retrieval\-augmented generation on knowledge graphs\.InProceedings of the AAAI Conference on Artificial Intelligence,Vol\.39,pp\. 12658–12666\.Cited by:[Appendix A](https://arxiv.org/html/2606.23898#A1.SS0.SSS0.Px1.p1.1)\.
- \[54\]V\. Tong, T\. Hoang, A\. Liu, G\. Van den Broeck, and M\. Niepert\(2025\)Learning to discretize denoising diffusion odes\.InProceedings of the 13th International Conference on Learning Representations,Cited by:[§2](https://arxiv.org/html/2606.23898#S2.p3.1)\.
- \[55\]H\. Touvron, M\. Cord, M\. Douze, F\. Massa, A\. Sablayrolles, and H\. Jégou\(2021\)Training data\-efficient image transformers & distillation through attention\.InInternational conference on machine learning,pp\. 10347–10357\.Cited by:[§2](https://arxiv.org/html/2606.23898#S2.p2.1)\.
- \[56\]M\. Tukan, L\. Mualem, E\. Netzer, and L\. Sigalat\(2025\)Improving model classification by optimizing the training dataset\.arXiv preprint arXiv:2507\.16729\.Cited by:[Appendix A](https://arxiv.org/html/2606.23898#A1.SS0.SSS0.Px2.p1.14)\.
- \[57\]Q\. Xiang, M\. Zhang, Y\. Shang, J\. Wu, Y\. Yan, and L\. Nie\(2025\)Dkdm: data\-free knowledge distillation for diffusion models with any architecture\.InProceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition,pp\. 2955–2965\.Cited by:[§2](https://arxiv.org/html/2606.23898#S2.p1.1),[§2](https://arxiv.org/html/2606.23898#S2.p2.1)\.
- \[58\]E\. Xie, J\. Chen, Y\. Zhao, J\. Yu, L\. Zhu, Y\. Lin, Z\. Zhang, M\. Li, J\. Chen, H\. Cai,et al\.\(2025\)SANA 1\.5: efficient scaling of training\-time and inference\-time compute in linear diffusion transformer\.InInternational Conference on Machine Learning,pp\. 68578–68598\.Cited by:[§2](https://arxiv.org/html/2606.23898#S2.p1.1)\.
- \[59\]S\. Xie, Z\. Xiao, D\. Kingma, T\. Hou, Y\. N\. Wu, K\. P\. Murphy, T\. Salimans, B\. Poole, and R\. Gao\(2024\)Em distillation for one\-step diffusion models\.Advances in Neural Information Processing Systems37,pp\. 45073–45104\.Cited by:[§1](https://arxiv.org/html/2606.23898#S1.p2.1),[§2](https://arxiv.org/html/2606.23898#S2.p3.1)\.
- \[60\]J\. Yang, B\. Martinez, A\. Bulat, and G\. Tzimiropoulos\(2021\)Knowledge distillation via softmax regression representation learning\.InInternational conference on learning representations,Cited by:[§1](https://arxiv.org/html/2606.23898#S1.p2.1),[§2](https://arxiv.org/html/2606.23898#S2.p2.1)\.
- \[61\]T\. Yin, M\. Gharbi, R\. Zhang, E\. Shechtman, F\. Durand, W\. T\. Freeman, and T\. Park\(2024\)One\-step diffusion with distribution matching distillation\.InProceedings of the IEEE/CVF conference on computer vision and pattern recognition,pp\. 6613–6623\.Cited by:[§1](https://arxiv.org/html/2606.23898#S1.p2.1),[§2](https://arxiv.org/html/2606.23898#S2.p3.1)\.
- \[62\]Q\. Zhang and Y\. Chen\(2022\)Fast sampling of diffusion models with exponential integrator\.arXiv preprint arXiv:2204\.13902\.Cited by:[§2](https://arxiv.org/html/2606.23898#S2.p3.1)\.
- \[63\]B\. Zhao, Q\. Cui, R\. Song, Y\. Qiu, and J\. Liang\(2022\)Decoupled knowledge distillation\.InProceedings of the IEEE/CVF Conference on computer vision and pattern recognition,pp\. 11953–11962\.Cited by:[§1](https://arxiv.org/html/2606.23898#S1.p2.1),[§2](https://arxiv.org/html/2606.23898#S2.p2.1)\.
- \[64\]W\. Zhao, L\. Bai, Y\. Rao, J\. Zhou, and J\. Lu\(2023\)Unipc: a unified predictor\-corrector framework for fast sampling of diffusion models\.Advances in Neural Information Processing Systems36,pp\. 49842–49869\.Cited by:[§2](https://arxiv.org/html/2606.23898#S2.p3.1)\.
- \[65\]K\. Zheng, C\. Lu, J\. Chen, and J\. Zhu\(2023\)Dpm\-solver\-v3: improved diffusion ode solver with empirical model statistics\.Advances in Neural Information Processing Systems36,pp\. 55502–55542\.Cited by:[§2](https://arxiv.org/html/2606.23898#S2.p3.1)\.
- \[66\]H\. Zhou and L\. Song\(2021\)Rethinking soft labels for knowledge distillation: a bias–variance tradeoff perspective\.InProceedings of International Conference on Learning Representations \(ICLR\),Cited by:[§1](https://arxiv.org/html/2606.23898#S1.p2.1),[§2](https://arxiv.org/html/2606.23898#S2.p2.1)\.
- \[67\]M\. Zhou, H\. Zheng, Z\. Wang, M\. Yin, and H\. Huang\(2024\)Score identity distillation: exponentially fast distillation of pretrained diffusion models for one\-step generation\.InForty\-first International Conference on Machine Learning,Cited by:[§1](https://arxiv.org/html/2606.23898#S1.p2.1),[§1](https://arxiv.org/html/2606.23898#S1.p3.1),[§2](https://arxiv.org/html/2606.23898#S2.p3.1)\.
- \[68\]H\. Zhu, D\. Tang, J\. Liu, M\. Lu, J\. Zheng, J\. Peng, D\. Li, Y\. Wang, F\. Jiang, L\. Tian,et al\.\(2024\)Dip\-go: a diffusion pruner via few\-step gradient optimization\.Advances in Neural Information Processing Systems37,pp\. 92581–92604\.Cited by:[§2](https://arxiv.org/html/2606.23898#S2.p1.1)\.

## Appendix AAdditional Related Work

#### Relation to bandit\-style allocation\.

ARIA can be interpreted through a multi\-armed bandit lens\[[49](https://arxiv.org/html/2606.23898#bib.bib64),[53](https://arxiv.org/html/2606.23898#bib.bib72),[4](https://arxiv.org/html/2606.23898#bib.bib71)\], where each region corresponds to an arm and the observed discrepancy acts as a reward signal used to guide future sampling\. However, unlike classical bandit settings, the reward distribution is non\-stationary and depends on the evolving model parameters\. Our analysis therefore focuses on the stability and tracking properties of the region\-level estimator rather than cumulative regret\. This perspective connects ARIA to adaptive data selection methods while highlighting the distinct challenges of training\-time non\-stationarity\.

#### Coresets\.

Coresets\[[24](https://arxiv.org/html/2606.23898#bib.bib74),[13](https://arxiv.org/html/2606.23898#bib.bib68),[22](https://arxiv.org/html/2606.23898#bib.bib69),[56](https://arxiv.org/html/2606.23898#bib.bib66)\]have been widely studied for large\-scale machine learning and clustering problems, in particular for objectives such askk\-means andkk\-median\[[18](https://arxiv.org/html/2606.23898#bib.bib65),[10](https://arxiv.org/html/2606.23898#bib.bib73),[17](https://arxiv.org/html/2606.23898#bib.bib67)\]\. In machine learning, we are \(usually\) given an input setP⊆ℝdP\\subseteq\\mathbb\{R\}^\{d\}ofnnpoints, its corresponding weight functionw:P→ℝw:P\\rightarrow\\mathbb\{R\}, a feasible set of queriesXX, and a loss functionϕ:P×X→\[0,∞\)\\phi:P\\times X\\rightarrow\[0,\\infty\)\. The tuple\(P,w,X,ϕ\)\(P,w,X,\\phi\)is called a*query space*, and it defines the optimization problem at hand\. For a given problem that is defined by its query space\(P,w,X,ϕ\)\(P,w,X,\\phi\), and an error parameterε∈\(0,1\)\\varepsilon\\in\(0,1\), anε\\varepsilon\-coreset is a small weighted subset of the input points that approximates the loss of the input setPPfor every feasible queryxx, up to a provable bound of1\+ε1\+\\varepsilon\.

## Appendix BTheorem Proof

In this section, we present the full tracking guarantee theorem and provide its complete proof\.

#### Why a theoretical guarantee is needed\.

ARIA maintains an exponential moving average \(EMA\) score for each semantic region to estimate its current training difficulty\. These scores determine the sampling distribution used during training, and therefore inaccurate estimates could cause the algorithm to focus on the wrong regions\. However, the discrepancy signal observed during training is stochastic, and the true region difficulty evolves as the student model improves\. It is therefore important to understand whether the maintained EMA scores reliably track the underlying region difficulty\.

#### Interpretation of the guarantee\.

Our analysis shows that the EMA score provably tracks the conditional mean discrepancy of each region\. The resulting bound decomposes the tracking error into three components: an initialization term that decays geometrically, a drift term capturing how the true region difficulty evolves during training, and a stochastic term reflecting sampling noise\.

In the bound,σ\\sigmameasures the variance of the stochastic discrepancy observations arising from minibatch sampling and diffusion noise\. The quantityVVmeasures the cumulative discounted drift of the true region difficulty, defined through the changes inJs\(n\)J\_\{s\}^\{\(n\)\}over time\. In realistic training dynamics, the discrepancy of a region evolves as the student model updates its parameters and progressively improves\. However, stochastic optimization typically produces gradual parameter updates, and therefore the difficulty of a region tends to evolve smoothly across training iterations\. The discounted drift assumption therefore allows the analysis to capture this non\-stationary behavior while remaining mild in practice\.

The theorem provides two complementary guarantees\. The RMS bound controls the expected tracking error, while the high\-probability bound shows that the tracking error remains uniformly bounded over an entire training horizon\. Together, these results justify the use of EMA scores in ARIA as reliable estimates of region difficulty for guiding the sampling distribution during training\.

### B\.1Preliminaries for the analysis

We analyze the behavior of the exponential moving average \(EMA\) score for a fixed region \(cluster\)ss\.

Since the score ofssis updated only whenssis selected, we consider the subsequence of training steps in whichssis chosen\. Letn=1,2,…n=1,2,\\dotsindex thenn\-th update of regionss\(not global training iterations\)\.

LetΔs\(n\)\\Delta\_\{s\}^\{\(n\)\}denote the scalar discrepancy observed at thenn\-th update ofss\. Let\{ℱn−1\}n≥1\\\{\\mathcal\{F\}\_\{n\-1\}\\\}\_\{n\\geq 1\}be the filtration representing all information available just before that update, including previously sampled data, past discrepancies, and the current model parameters\.

Define the conditional mean discrepancy

Js\(n\):=𝔼​\[Δs\(n\)∣ℱn−1\],J\_\{s\}^\{\(n\)\}:=\\mathbb\{E\}\\\!\\left\[\\Delta\_\{s\}^\{\(n\)\}\\mid\\mathcal\{F\}\_\{n\-1\}\\right\],and the noise term

ξn:=Δs\(n\)−Js\(n\)\.\\xi\_\{n\}:=\\Delta\_\{s\}^\{\(n\)\}\-J\_\{s\}^\{\(n\)\}\.
and we define the tracking error

en:=Ds\(n\)−Js\(n\)\.e\_\{n\}:=D\_\{s\}^\{\(n\)\}\-J\_\{s\}^\{\(n\)\}\.
###### Theorem 2\.

Fix a clusterssand consider the subsequence of training iterations in whichssis selected, indexed byn=1,2,…n=1,2,\\dots\. Let\{ℱn−1\}n≥1\\\{\\mathcal\{F\}\_\{n\-1\}\\\}\_\{n\\geq 1\},Δs\(n\)\\Delta\_\{s\}^\{\(n\)\},Js\(n\)J\_\{s\}^\{\(n\)\},ξn\\xi\_\{n\},Ds\(n\)D\_\{s\}^\{\(n\)\}, anden:=Ds\(n\)−Js\(n\)e\_\{n\}:=D\_\{s\}^\{\(n\)\}\-J\_\{s\}^\{\(n\)\}be as defined in the preliminaries, and letα∈\(0,1\)\\alpha\\in\(0,1\)\.

Assume that for alln≥1n\\geq 1:

1. \(A1\)Unbiased noise:𝔼​\[ξn∣ℱn−1\]=0\\mathbb\{E\}\[\\xi\_\{n\}\\mid\\mathcal\{F\}\_\{n\-1\}\]=0\.
2. \(A2\)Bounded conditional variance:𝔼​\[ξn2∣ℱn−1\]≤σ2\\mathbb\{E\}\[\\xi\_\{n\}^\{2\}\\mid\\mathcal\{F\}\_\{n\-1\}\]\\leq\\sigma^\{2\}for someσ\>0\\sigma\>0\.
3. \(A3\)Bounded discounted drift:for someV≥0V\\geq 0and everyn≥1n\\geq 1, Vs\(α\)​\(n\):=∑k=1n\(1−α\)n−k​\|Js\(k\)−Js\(k−1\)\|≤V,V\_\{s\}^\{\(\\alpha\)\}\(n\):=\\sum\_\{k=1\}^\{n\}\(1\-\\alpha\)^\{\\,n\-k\}\\,\\big\|J\_\{s\}^\{\(k\)\}\-J\_\{s\}^\{\(k\-1\)\}\\big\|\\leq V,\(3\)whereJs\(0\)J\_\{s\}^\{\(0\)\}is any fixed reference value\.

Then the following statements hold\.

1. \(i\)\(RMS tracking bound\)\.For everyn≥1n\\geq 1, 𝔼​\[en2\]≤\(1−α\)n​\|e0\|\+\(1−α\)​V\+σ​α\.\\sqrt\{\\mathbb\{E\}\[e\_\{n\}^\{2\}\]\}\\;\\leq\\;\(1\-\\alpha\)^\{n\}\|e\_\{0\}\|\\;\+\\;\(1\-\\alpha\)\\,V\\;\+\\;\\sigma\\sqrt\{\\alpha\}\.\(4\)
2. \(ii\)\(High\-probability tracking bound, uniform up to a horizon\)\.Assume in addition that 1. \(A1\)Bounded noise:\|ξn\|≤b\|\\xi\_\{n\}\|\\leq balmost surely for someb\>0b\>0\. Fix any horizonN≥1N\\geq 1and anyδ∈\(0,1\)\\delta\\in\(0,1\)\. Then with probability at least1−δ1\-\\delta, simultaneously for alln∈\{1,…,N\}n\\in\\\{1,\\dots,N\\\}, \|en\|≤\(1−α\)n​\|e0\|\+\(1−α\)​V\+2​α​σ2​log⁡2​Nδ\+23​α​b​log⁡2​Nδ\.\|e\_\{n\}\|\\;\\leq\\;\(1\-\\alpha\)^\{n\}\|e\_\{0\}\|\\;\+\\;\(1\-\\alpha\)V\\;\+\\;\\sqrt\{2\\,\\alpha\\,\\sigma^\{2\}\\log\\\!\\frac\{2N\}\{\\delta\}\}\\;\+\\;\\frac\{2\}\{3\}\\,\\alpha\\,b\\,\\log\\\!\\frac\{2N\}\{\\delta\}\.\(5\)

###### Proof\.

\(i\) RMS bound\.Leten:=Ds\(n\)−Js\(n\)e\_\{n\}:=D\_\{s\}^\{\(n\)\}\-J\_\{s\}^\{\(n\)\}denote the tracking error\. Starting from the EMA update for clusterssand using the decompositionΔs\(n\)=Js\(n\)\+ξn\\Delta\_\{s\}^\{\(n\)\}=J\_\{s\}^\{\(n\)\}\+\\xi\_\{n\}, we can write

Ds\(n\)=\(1−α\)​Ds\(n−1\)\+α​\(Js\(n\)\+ξn\)\.D\_\{s\}^\{\(n\)\}=\(1\-\\alpha\)D\_\{s\}^\{\(n\-1\)\}\+\\alpha\\bigl\(J\_\{s\}^\{\(n\)\}\+\\xi\_\{n\}\\bigr\)\.SubtractingJs\(n\)J\_\{s\}^\{\(n\)\}and adding/subtracting\(1−α\)​Js\(n−1\)\(1\-\\alpha\)J\_\{s\}^\{\(n\-1\)\}yields

en=\(1−α\)​en−1\+\(1−α\)​\(Js\(n−1\)−Js\(n\)\)\+α​ξn\.e\_\{n\}=\(1\-\\alpha\)e\_\{n\-1\}\+\(1\-\\alpha\)\\bigl\(J\_\{s\}^\{\(n\-1\)\}\-J\_\{s\}^\{\(n\)\}\\bigr\)\+\\alpha\\xi\_\{n\}\.\(6\)
Define the drift incrementδn:=Js\(n−1\)−Js\(n\)\\delta\_\{n\}:=J\_\{s\}^\{\(n\-1\)\}\-J\_\{s\}^\{\(n\)\}\. Repeated substitution of \([6](https://arxiv.org/html/2606.23898#A2.E6)\) yields

en=\(1−α\)n​e0\+\(1−α\)​∑k=1n\(1−α\)n−k​δk\+α​∑k=1n\(1−α\)n−k​ξk\.e\_\{n\}=\(1\-\\alpha\)^\{n\}e\_\{0\}\+\(1\-\\alpha\)\\sum\_\{k=1\}^\{n\}\(1\-\\alpha\)^\{\\,n\-k\}\\,\\delta\_\{k\}\+\\alpha\\sum\_\{k=1\}^\{n\}\(1\-\\alpha\)^\{\\,n\-k\}\\,\\xi\_\{k\}\.\(7\)For convenience, denote the three terms in \([7](https://arxiv.org/html/2606.23898#A2.E7)\) by

An:=\(1−α\)n​e0,Bn:=\(1−α\)​∑k=1n\(1−α\)n−k​δk,Cn:=α​∑k=1n\(1−α\)n−k​ξk,A\_\{n\}:=\(1\-\\alpha\)^\{n\}e\_\{0\},\\qquad B\_\{n\}:=\(1\-\\alpha\)\\sum\_\{k=1\}^\{n\}\(1\-\\alpha\)^\{\\,n\-k\}\\,\\delta\_\{k\},\\qquad C\_\{n\}:=\\alpha\\sum\_\{k=1\}^\{n\}\(1\-\\alpha\)^\{\\,n\-k\}\\,\\xi\_\{k\},so thaten=An\+Bn\+Cne\_\{n\}=A\_\{n\}\+B\_\{n\}\+C\_\{n\}\.

The initialization term satisfies

\|An\|=\(1−α\)n​\|e0\|\.\|A\_\{n\}\|=\(1\-\\alpha\)^\{n\}\|e\_\{0\}\|\.For the drift term, apply the triangle inequality andδk=Js\(k−1\)−Js\(k\)\\delta\_\{k\}=J\_\{s\}^\{\(k\-1\)\}\-J\_\{s\}^\{\(k\)\}:

\|Bn\|≤\(1−α\)​∑k=1n\(1−α\)n−k​\|δk\|=\(1−α\)​∑k=1n\(1−α\)n−k​\|Js\(k\)−Js\(k−1\)\|\.\|B\_\{n\}\|\\leq\(1\-\\alpha\)\\sum\_\{k=1\}^\{n\}\(1\-\\alpha\)^\{\\,n\-k\}\|\\delta\_\{k\}\|=\(1\-\\alpha\)\\sum\_\{k=1\}^\{n\}\(1\-\\alpha\)^\{\\,n\-k\}\\bigl\|J\_\{s\}^\{\(k\)\}\-J\_\{s\}^\{\(k\-1\)\}\\bigr\|\.By assumption \(A3\) this discounted drift sum is at mostVV, hence

\|Bn\|≤\(1−α\)​V\.\|B\_\{n\}\|\\leq\(1\-\\alpha\)V\.\(8\)
We next control the stochastic termCnC\_\{n\}in mean square\. From the definition ofCnC\_\{n\},

Cn=α​∑k=1nwk,n​ξk,wk,n:=\(1−α\)n−k\.C\_\{n\}=\\alpha\\sum\_\{k=1\}^\{n\}w\_\{k,n\}\\,\\xi\_\{k\},\\qquad w\_\{k,n\}:=\(1\-\\alpha\)^\{n\-k\}\.ExpandingCn2C\_\{n\}^\{2\}and taking expectations gives

𝔼​\[Cn2\]=α2​𝔼​\[∑k=1n∑j=1nwk,n​wj,n​ξk​ξj\]\.\\mathbb\{E\}\[C\_\{n\}^\{2\}\]=\\alpha^\{2\}\\,\\mathbb\{E\}\\\!\\left\[\\sum\_\{k=1\}^\{n\}\\sum\_\{j=1\}^\{n\}w\_\{k,n\}w\_\{j,n\}\\,\\xi\_\{k\}\\xi\_\{j\}\\right\]\.We show that cross terms vanish: forj<kj<k,ξj\\xi\_\{j\}isℱk−1\\mathcal\{F\}\_\{k\-1\}\-measurable, hence

𝔼​\[ξk​ξj\]=𝔼​\[𝔼​\[ξk​ξj∣ℱk−1\]\]=𝔼​\[ξj​𝔼​\[ξk∣ℱk−1\]\]=0,\\mathbb\{E\}\[\\xi\_\{k\}\\xi\_\{j\}\]=\\mathbb\{E\}\\\!\\left\[\\mathbb\{E\}\[\\xi\_\{k\}\\xi\_\{j\}\\mid\\mathcal\{F\}\_\{k\-1\}\]\\right\]=\\mathbb\{E\}\\\!\\left\[\\xi\_\{j\}\\,\\mathbb\{E\}\[\\xi\_\{k\}\\mid\\mathcal\{F\}\_\{k\-1\}\]\\right\]=0,where the last equality uses \(A1\)\. Therefore only diagonal terms remain:

𝔼​\[Cn2\]=α2​∑k=1nwk,n2​𝔼​\[ξk2\]\.\\mathbb\{E\}\[C\_\{n\}^\{2\}\]=\\alpha^\{2\}\\sum\_\{k=1\}^\{n\}w\_\{k,n\}^\{2\}\\,\\mathbb\{E\}\[\\xi\_\{k\}^\{2\}\]\.Using \(A2\) and iterated expectation,

𝔼​\[ξk2\]=𝔼​\[𝔼​\[ξk2∣ℱk−1\]\]≤σ2,\\mathbb\{E\}\[\\xi\_\{k\}^\{2\}\]=\\mathbb\{E\}\\\!\\left\[\\mathbb\{E\}\[\\xi\_\{k\}^\{2\}\\mid\\mathcal\{F\}\_\{k\-1\}\]\\right\]\\leq\\sigma^\{2\},so

𝔼​\[Cn2\]≤α2​σ2​∑k=1n\(1−α\)2​\(n−k\)\.\\mathbb\{E\}\[C\_\{n\}^\{2\}\]\\leq\\alpha^\{2\}\\sigma^\{2\}\\sum\_\{k=1\}^\{n\}\(1\-\\alpha\)^\{2\(n\-k\)\}\.The remaining sum is geometric:

∑k=1n\(1−α\)2​\(n−k\)=∑j=0n−1\(1−α\)2​j≤11−\(1−α\)2=12​α−α2≤1α\.\\sum\_\{k=1\}^\{n\}\(1\-\\alpha\)^\{2\(n\-k\)\}=\\sum\_\{j=0\}^\{n\-1\}\(1\-\\alpha\)^\{2j\}\\leq\\frac\{1\}\{1\-\(1\-\\alpha\)^\{2\}\}=\\frac\{1\}\{2\\alpha\-\\alpha^\{2\}\}\\leq\\frac\{1\}\{\\alpha\}\.Hence

𝔼​\[Cn2\]≤α​σ2⇒𝔼​\[Cn2\]≤σ​α\.\\mathbb\{E\}\[C\_\{n\}^\{2\}\]\\leq\\alpha\\sigma^\{2\}\\qquad\\Rightarrow\\qquad\\sqrt\{\\mathbb\{E\}\[C\_\{n\}^\{2\}\]\}\\leq\\sigma\\sqrt\{\\alpha\}\.\(9\)
Combining the above bounds, fromen=An\+Bn\+Cne\_\{n\}=A\_\{n\}\+B\_\{n\}\+C\_\{n\}and the triangle inequality,

\|en\|≤\|An\|\+\|Bn\|\+\|Cn\|\.\|e\_\{n\}\|\\leq\|A\_\{n\}\|\+\|B\_\{n\}\|\+\|C\_\{n\}\|\.TakingL2L\_\{2\}\-norms \(i\.e\.,‖X‖2:=𝔼​\[X2\]\\\|X\\\|\_\{2\}:=\\sqrt\{\\mathbb\{E\}\[X^\{2\}\]\}\) and using Minkowski’s inequality,

𝔼​\[en2\]≤𝔼​\[An2\]\+𝔼​\[Bn2\]\+𝔼​\[Cn2\]\.\\sqrt\{\\mathbb\{E\}\[e\_\{n\}^\{2\}\]\}\\leq\\sqrt\{\\mathbb\{E\}\[A\_\{n\}^\{2\}\]\}\+\\sqrt\{\\mathbb\{E\}\[B\_\{n\}^\{2\}\]\}\+\\sqrt\{\\mathbb\{E\}\[C\_\{n\}^\{2\}\]\}\.SinceAnA\_\{n\}andBnB\_\{n\}are deterministic givene0e\_\{0\}and\{Js\(k\)\}\\\{J\_\{s\}^\{\(k\)\}\\\},𝔼​\[An2\]=\|An\|\\sqrt\{\\mathbb\{E\}\[A\_\{n\}^\{2\}\]\}=\|A\_\{n\}\|and𝔼​\[Bn2\]=\|Bn\|\\sqrt\{\\mathbb\{E\}\[B\_\{n\}^\{2\}\]\}=\|B\_\{n\}\|\. Applying the deterministic bounds above and \([9](https://arxiv.org/html/2606.23898#A2.E9)\) gives

𝔼​\[en2\]≤\(1−α\)n​\|e0\|\+\(1−α\)​V\+σ​α,\\sqrt\{\\mathbb\{E\}\[e\_\{n\}^\{2\}\]\}\\leq\(1\-\\alpha\)^\{n\}\|e\_\{0\}\|\+\(1\-\\alpha\)V\+\\sigma\\sqrt\{\\alpha\},which is exactly \([4](https://arxiv.org/html/2606.23898#A2.E4)\)\.

#### \(ii\) Uniform high\-probability bound\.

We now prove the uniform high\-probability bound\. Fix ann∈\{1,…,N\}n\\in\\\{1,\\dots,N\\\}\. By the decompositionen=An\+Bn\+Cne\_\{n\}=A\_\{n\}\+B\_\{n\}\+C\_\{n\}and the deterministic bounds above,

\|en\|≤\|An\|\+\|Bn\|\+\|Cn\|≤\(1−α\)n​\|e0\|\+\(1−α\)​V\+\|Cn\|\.\|e\_\{n\}\|\\leq\|A\_\{n\}\|\+\|B\_\{n\}\|\+\|C\_\{n\}\|\\leq\(1\-\\alpha\)^\{n\}\|e\_\{0\}\|\+\(1\-\\alpha\)V\+\|C\_\{n\}\|\.Thus it suffices to upper bound\|Cn\|\|C\_\{n\}\|with high probability\.

RecallCn=∑k=1nXkC\_\{n\}=\\sum\_\{k=1\}^\{n\}X\_\{k\}where we define

Xk:=α​wk,n​ξk,wk,n:=\(1−α\)n−k\.X\_\{k\}:=\\alpha\\,w\_\{k,n\}\\,\\xi\_\{k\},\\qquad w\_\{k,n\}:=\(1\-\\alpha\)^\{n\-k\}\.We claim that\(Xk\)k=1n\(X\_\{k\}\)\_\{k=1\}^\{n\}is a martingale difference sequence with respect to\(ℱk\)k=0n\(\\mathcal\{F\}\_\{k\}\)\_\{k=0\}^\{n\}\. Indeed, sincewk,nw\_\{k,n\}is deterministic andξk\\xi\_\{k\}isℱk\\mathcal\{F\}\_\{k\}\-measurable,XkX\_\{k\}isℱk\\mathcal\{F\}\_\{k\}\-measurable\. Moreover, by \(A1\),

𝔼​\[Xk∣ℱk−1\]=α​wk,n​𝔼​\[ξk∣ℱk−1\]=0\.\\mathbb\{E\}\[X\_\{k\}\\mid\\mathcal\{F\}\_\{k\-1\}\]=\\alpha w\_\{k,n\}\\,\\mathbb\{E\}\[\\xi\_\{k\}\\mid\\mathcal\{F\}\_\{k\-1\}\]=0\.By \(A4\) andwk,n≤1w\_\{k,n\}\\leq 1,

\|Xk\|=α​wk,n​\|ξk\|≤α​b\.\|X\_\{k\}\|=\\alpha w\_\{k,n\}\|\\xi\_\{k\}\|\\leq\\alpha b\.Finally, by \(A2\),

𝔼​\[Xk2∣ℱk−1\]=α2​wk,n2​𝔼​\[ξk2∣ℱk−1\]≤α2​wk,n2​σ2\.\\mathbb\{E\}\[X\_\{k\}^\{2\}\\mid\\mathcal\{F\}\_\{k\-1\}\]=\\alpha^\{2\}w\_\{k,n\}^\{2\}\\,\\mathbb\{E\}\[\\xi\_\{k\}^\{2\}\\mid\\mathcal\{F\}\_\{k\-1\}\]\\leq\\alpha^\{2\}w\_\{k,n\}^\{2\}\\,\\sigma^\{2\}\.Summing overk=1,…,nk=1,\\dots,n,

∑k=1n𝔼​\[Xk2∣ℱk−1\]≤α2​σ2​∑k=1nwk,n2=α2​σ2​∑k=1n\(1−α\)2​\(n−k\)≤α2​σ2⋅1α=α​σ2,\\sum\_\{k=1\}^\{n\}\\mathbb\{E\}\[X\_\{k\}^\{2\}\\mid\\mathcal\{F\}\_\{k\-1\}\]\\leq\\alpha^\{2\}\\sigma^\{2\}\\sum\_\{k=1\}^\{n\}w\_\{k,n\}^\{2\}=\\alpha^\{2\}\\sigma^\{2\}\\sum\_\{k=1\}^\{n\}\(1\-\\alpha\)^\{2\(n\-k\)\}\\leq\\alpha^\{2\}\\sigma^\{2\}\\cdot\\frac\{1\}\{\\alpha\}=\\alpha\\sigma^\{2\},where we used the same geometric\-sum bound as above\. Thus the predictable quadratic variation is at mostv:=α​σ2v:=\\alpha\\sigma^\{2\}\.

Freedman’s inequality \(for martingale differences with bounded increments\) states that if∑k=1n𝔼​\[Xk2∣ℱk−1\]≤v\\sum\_\{k=1\}^\{n\}\\mathbb\{E\}\[X\_\{k\}^\{2\}\\mid\\mathcal\{F\}\_\{k\-1\}\]\\leq vwith high probability and\|Xk\|≤c\|X\_\{k\}\|\\leq cwith high probability, then for anyη∈\(0,1\)\\eta\\in\(0,1\),

Pr⁡\(∑k=1nXk≥2​v​log⁡\(1/η\)\+23​c​log⁡\(1/η\)\)≤η\.\\Pr\\\!\\left\(\\sum\_\{k=1\}^\{n\}X\_\{k\}\\geq\\sqrt\{2v\\log\(1/\\eta\)\}\+\\frac\{2\}\{3\}c\\log\(1/\\eta\)\\right\)\\leq\\eta\.We apply this withc=α​bc=\\alpha b,v=α​σ2v=\\alpha\\sigma^\{2\}, and∑k=1nXk=Cn\\sum\_\{k=1\}^\{n\}X\_\{k\}=C\_\{n\}\. This gives

Pr⁡\(Cn≥2​α​σ2​log⁡\(1/η\)\+23​α​b​log⁡\(1/η\)\)≤η\.\\Pr\\\!\\left\(C\_\{n\}\\geq\\sqrt\{2\\alpha\\sigma^\{2\}\\log\(1/\\eta\)\}\+\\frac\{2\}\{3\}\\alpha b\\log\(1/\\eta\)\\right\)\\leq\\eta\.Applying the same bound to−Cn\-C\_\{n\}yields

Pr⁡\(\|Cn\|≥2​α​σ2​log⁡\(2/η\)\+23​α​b​log⁡\(2/η\)\)≤η,\\Pr\\\!\\left\(\|C\_\{n\}\|\\geq\\sqrt\{2\\alpha\\sigma^\{2\}\\log\(2/\\eta\)\}\+\\frac\{2\}\{3\}\\alpha b\\log\(2/\\eta\)\\right\)\\leq\\eta,where we used a union bound for the two tails\.

Setη:=δ/N\\eta:=\\delta/Nand apply the above inequality to eachn∈\{1,…,N\}n\\in\\\{1,\\dots,N\\\}\. Then for each fixednn,

Pr⁡\(\|Cn\|≥2​α​σ2​log⁡2​Nδ\+23​α​b​log⁡2​Nδ\)≤δN\.\\Pr\\\!\\left\(\|C\_\{n\}\|\\geq\\sqrt\{2\\alpha\\sigma^\{2\}\\log\\\!\\frac\{2N\}\{\\delta\}\}\+\\frac\{2\}\{3\}\\alpha b\\log\\\!\\frac\{2N\}\{\\delta\}\\right\)\\leq\\frac\{\\delta\}\{N\}\.Taking a union bound overn=1,…,Nn=1,\\dots,Nshows that with probability at least1−δ1\-\\delta, simultaneously for alln∈\{1,…,N\}n\\in\\\{1,\\dots,N\\\},

\|Cn\|≤2​α​σ2​log⁡2​Nδ\+23​α​b​log⁡2​Nδ\.\|C\_\{n\}\|\\leq\\sqrt\{2\\alpha\\sigma^\{2\}\\log\\\!\\frac\{2N\}\{\\delta\}\}\+\\frac\{2\}\{3\}\\alpha b\\log\\\!\\frac\{2N\}\{\\delta\}\.\(10\)
On the same event, for everyn≤Nn\\leq N,

\|en\|≤\(1−α\)n​\|e0\|\+\(1−α\)​V\+\|Cn\|\.\|e\_\{n\}\|\\leq\(1\-\\alpha\)^\{n\}\|e\_\{0\}\|\+\(1\-\\alpha\)V\+\|C\_\{n\}\|\.Substituting \([10](https://arxiv.org/html/2606.23898#A2.E10)\) gives exactly \([5](https://arxiv.org/html/2606.23898#A2.E5)\)\. This completes the proof\. ∎

## Appendix CAdditional Experiments

### C\.1Implementation Details

We adopt the same experimental setup used in\[[26](https://arxiv.org/html/2606.23898#bib.bib17)\]\. All models are trained using the AdamW optimizer with a learning rate of5×10−55\\times 10^\{\-5\}\. Training is performed on four NVIDIA A100 GPUs with 80GB memory using a batch size of 256\.

We employ the same loss functions described in Eq\. \(1\) and Eq\. \(2\) in the main paper, with equal weights of 1\. The feature\-level distillation loss is applied after each block of the U\-Net architecture\.

### C\.2Embedding Similarity and Training Signal Consistency

ARIA clusters the conditioning space using text embeddings\. To validate that this geometry is meaningful for optimization, we study whether prompts that are close in embedding space also induce similar distillation signals\.

For a noisy latentxtx\_\{t\}and conditioningcc, the distillation residual is

r​\(xt,c\)=ϵ^S​\(xt,c\)−ϵ^T​\(xt,c\),r\(x\_\{t\},c\)=\\hat\{\\epsilon\}\_\{S\}\(x\_\{t\},c\)\-\\hat\{\\epsilon\}\_\{T\}\(x\_\{t\},c\),whereϵ^S\\hat\{\\epsilon\}\_\{S\}andϵ^T\\hat\{\\epsilon\}\_\{T\}are the student and teacher denoising predictions\. Since this residual determines the update applied to the student, we analyze both the residual similarity itself and the similarity between the corresponding gradient directions\.

Figure[4](https://arxiv.org/html/2606.23898#A3.F4)shows the cosine distance between training signals as a function of cosine similarity in text embedding space\. The left subfigure reports denoising residual distance and the right subfigure reports gradient direction distance, each at two noise levels\. In all cases, higher embedding similarity corresponds to smaller training\-signal distance\. This indicates that nearby prompts in the conditioning space tend to produce similar optimization signals, supporting the use of embedding\-based clustering for defining regions in ARIA\.

![Refer to caption](https://arxiv.org/html/2606.23898v1/figures/paper_ready_noise.png)\(\(a\)\)Denoising residual distance versus embedding similarity\.
![Refer to caption](https://arxiv.org/html/2606.23898v1/figures/paper_ready_gradients.png)\(\(b\)\)Gradient direction distance versus embedding similarity\.

Figure 4:Relationship between text embedding similarity and the similarity of the distillation training signal\. Each subfigure already contains two panels corresponding toσ=0\.9\\sigma=0\.9andσ=1\.0\\sigma=1\.0\. The left subfigure shows denoising residual distance, while the right subfigure shows gradient direction distance\. Across both noise levels and both signal types, prompts that are closer in embedding space consistently induce more similar training signals\. The green curves show the average over selected clusters, and the shaded regions indicate the interquartile range of pairwise distances\.
### C\.3Sampling Distribution Analysis

After clustering the text corpus into semantic regions, ARIA maintains a difficulty scoreDsD\_\{s\}for each cluster using the EMA statistics described in Sec\. 2\. The remaining design choice is how these difficulty scores are converted into a sampling distribution over clusters\.

To control the sharpness of the distribution, we first apply a power transformation to the difficulty scores

ws=Dsβ,w\_\{s\}=D\_\{s\}^\{\\beta\},\(11\)withβ=0\.6\\beta=0\.6in our experiments\. This transformation amplifies relative differences between clusters and produces a smoother long\-tailed distribution over difficulty levels\.

We then convert the transformed scores into sampling probabilities using one of two normalization schemes\. The first is a power\-law allocation

ps=ws∑j=1Swj,p\_\{s\}=\\frac\{w\_\{s\}\}\{\\sum\_\{j=1\}^\{S\}w\_\{j\}\},\(12\)which directly allocates probability proportionally to the transformed difficulty\. The second uses a temperature\-scaled softmax

ps=exp⁡\(ws/T\)∑j=1Sexp⁡\(wj/T\),p\_\{s\}=\\frac\{\\exp\(w\_\{s\}/T\)\}\{\\sum\_\{j=1\}^\{S\}\\exp\(w\_\{j\}/T\)\},\(13\)where the temperatureTTcontrols the concentration of probability mass\. Large values ofTTproduce distributions close to uniform sampling, while smaller values increasingly concentrate probability on the hardest clusters\.

Figure[5](https://arxiv.org/html/2606.23898#A3.F5)illustrates the resulting sampling behavior\. Figure[5](https://arxiv.org/html/2606.23898#A3.F5)\(a\) shows the probability distributions over clusters ranked by difficulty\. Uniform sampling allocates equal probability to all clusters, whereas decreasing the softmax temperature or using power\-law weighting increasingly concentrates probability on the hardest clusters\.

However, overly aggressive focusing can reduce coverage of easier regions\. Fig\.[5](https://arxiv.org/html/2606.23898#A3.F5)\(b\) shows that sharper distributions \(e\.g\.,T=0\.05T\{=\}0\.05or power\-law weighting\) allocate a larger fraction of samples to the top 10% hardest clusters while reducing the share assigned to the easiest clusters\. In our experiments, power\-law weighting improved performance for teacher\-initialized models, but only matched or sometimes degraded performance for randomly initialized students\.

Finally, Fig\.[5](https://arxiv.org/html/2606.23898#A3.F5)\(c\) reports the Spearman correlation between cluster difficulty and sampling frequency over training\. Lower temperatures and power\-law sampling produce stronger alignment between the sampling probability and the estimated cluster difficulty\.

Based on this analysis, we adoptT=0\.1T\{=\}0\.1, which provides a moderate focusing effect while preserving sufficient sampling coverage across all clusters\.

![Refer to caption](https://arxiv.org/html/2606.23898v1/x3.png)Figure 5:Analysis of sampling distributions over semantic clusters\. \(a\) Probability distribution over clusters ranked by difficulty\. Lower temperatures concentrate probability mass on harder clusters, while power\-law weighting produces the sharpest allocation\. \(b\) Fraction of samples allocated to the hardest and easiest 10% of clusters during training\. Aggressive focusing increases coverage of difficult regions but reduces sampling of easy clusters\. \(c\) Spearman correlation between cluster difficultyDsD\_\{s\}and sampling frequency\. Lower temperatures and power\-law sampling produce stronger alignment with cluster difficulty\.
### C\.4Additional Empirical Evaluation

In what follows, we present additional empirical evaluations that further illustrate the training dynamics of ARIA compared to Random Conditioning \(RC\)\. Figure[6](https://arxiv.org/html/2606.23898#A3.F6)reports the evolution of standard generation metrics across training checkpoints for the C\-Micro student architecture\.

Specifically, we track the Fréchet Inception Distance \(FID\), CLIP score, and Inception Score \(IS\) as training progresses from2525K to300300K iterations\. These curves provide a more detailed view of the convergence behavior discussed in the main paper\.

As shown in Fig\.[6](https://arxiv.org/html/2606.23898#A3.F6), ARIA consistently achieves better performance earlier during training\. In particular, ARIA reaches lower FID and higher CLIP scores at substantially earlier checkpoints, indicating faster alignment between the student and teacher models\. A similar trend is observed for the Inception Score, where ARIA maintains higher values throughout most of the training trajectory\.

These results further support the observation that ARIA improves the efficiency of the distillation process by allocating training updates to more informative regions of the conditioning space\. As a result, the student model converges faster while incurring negligible computational overhead compared to RC\.

![Refer to caption](https://arxiv.org/html/2606.23898v1/x4.png)Figure 6:Training dynamics for the C\-Micro student architecture\. We report \(a\) FID, \(b\) CLIP score, and \(c\) Inception Score \(IS\) across training checkpoints from2525K to300300K iterations\. ARIA consistently achieves lower FID and higher CLIP scores earlier in training compared to Random Conditioning \(RC\), demonstrating faster convergence of the distilled student model\.Key TakeawayARIA significantly improves training efficiency by accelerating convergence of randomly initialized student models\.Across training checkpoints, ARIA consistently achieves lower FID and higher CLIP scores earlier than RC, indicating that the student reaches strong performance with substantially fewer optimization steps\.This property is particularly important in practical distillation settings where the student architecture differs from the teacher, for example when training compressed or lightweight models that cannot reuse the teacher’s weights\. In such scenarios the student must be trained from random initialization, making convergence speed a critical factor\.By directing updates toward harder regions of the conditioning space, ARIA allocates the training budget more effectively, allowing high\-quality students to be obtained with reduced compute\. This is especially valuable in resource\-constrained environments where training large diffusion models is expensive and iterative experimentation \(e\.g\., architecture search or hyperparameter tuning\) is required\.

![Refer to caption](https://arxiv.org/html/2606.23898v1/x5.png)Figure 7:Qualitative comparison on animal\-related prompts on models that were trained on animal\-free images\. From left to right, we show outputs from the channel\-pruned student trained with ARIA, the same channel\-pruned student trained with RC, the block\-pruned student trained with ARIA, the same block\-pruned student trained with RC, and the teacher\.![Refer to caption](https://arxiv.org/html/2606.23898v1/x6.png)Figure 8:Qualitative comparison on vehicles\-related prompts on models that were trained on images without vehicles\.![Refer to caption](https://arxiv.org/html/2606.23898v1/x7.png)Figure 9:Qualitative comparison on food\-related prompts on models that were trained on images without food\.![Refer to caption](https://arxiv.org/html/2606.23898v1/x8.png)Figure 10:In our experiments, we observed that training with ARIA significantly improves convergence speed\. This figure provides a qualitative comparison across training checkpoints, illustrating that ARIA converges faster than RC\. We compare generations at 25%, 60%, and 100% of training for two students trained from random initialization: C\-Micro and C\-Tiny\.

## Appendix DClustering the data

### D\.1Constructing Semantic Regions via CLIP Clustering

Recall that our framework is agnostic to the specific procedure used to divide the training samples space into regions\. In principle, any method that partitions the training samples space into meaningful groups can be used\. In our implementation, we instantiate regions by clustering text prompts in the embedding space of a pretrained CLIP text encoder\. The main idea is that prompts that are close in CLIP space tend to share semantic content, and therefore clustering this space provides a practical way to define regions over which our cluster\-level discrepancy statistics can be maintained\.

More concretely, let𝒞=\{c1,…,cn\}\\mathcal\{C\}=\\\{c\_\{1\},\\dots,c\_\{n\}\\\}denote the set of text conditions used to define the conditioning pool, and letei∈ℝde\_\{i\}\\in\\mathbb\{R\}^\{d\}be the CLIP embedding of conditioncic\_\{i\}\. We partition\{ei\}i=1n\\\{e\_\{i\}\\\}\_\{i=1\}^\{n\}intoKKclusters usingkk\-means\. Since the quality of the partition depends on the choice ofKK, we do not fix it a priori\. Instead, we selectKKbased on the silhouette score\[[47](https://arxiv.org/html/2606.23898#bib.bib70)\], which is a widely used criterion for evaluating the quality of a clustering by balancing within\-cluster cohesion against between\-cluster separation\. Intuitively, a good value ofKKshould produce clusters whose members are close to one another while remaining well separated from points assigned to other clusters\.

A direct sweep over many candidate values ofKKon the full embedding set is very expensive\. Runningkk\-means repeatedly on a very large prompt collection or training samples is already costly, and computing clustering\-quality diagnostics on top of that adds further overhead\. To make this step practical, we use akk\-means coreset\. A coreset is a weighted subset of points that approximately preserves thekk\-means objective of the original dataset, and therefore can be used as a much smaller proxy for downstream clustering computations\. In our setting, this allows us to evaluate many candidate values ofKKusing only a compact weighted representative subset, instead of repeatedly clustering the full prompt pool\.

We want to use the sensitivity sampling framework to compute a coreset for a set of pointsPPinℛd\\mathcal\{R\}^\{d\}\. Sensitivity\-based coreset constructions typically rely on an initial approximate clustering solution in order to bound the importance \(sensitivity\) of each point with respect to thekk\-means objective\. This is commonly obtained through a bicriteria solution, which provides a constant\-factor approximation to the optimalkk\-means cost while possibly using more thankkcenters\. Such a solution allows one to derive upper bounds on point sensitivities that can be used for importance sampling\.

#### Bicriteria initialization viakk\-means\+\+\.

To build the coreset, we first obtain a bicriteria solution usingkk\-means\+\+ seeding\.kk\-means\+\+ chooses the first center uniformly at random and then samples each next center with probability proportional to the squared distance to the nearest previously chosen center\. Arthur and Vassilvitskii\[[1](https://arxiv.org/html/2606.23898#bib.bib63)\]showed that thisD2D^\{2\}\-sampling procedure yields anO​\(log⁡K\)O\(\\log K\)approximation in expectation to the optimalkk\-means objective; specifically, their Theorem 3\.1 proves

𝔼​\[ϕ\]≤8​\(ln⁡K\+2\)​ϕOPT\.\\mathbb\{E\}\[\\phi\]\\leq 8\(\\ln K\+2\)\\phi\_\{\\mathrm\{OPT\}\}\.They also note that the running time of the seeding procedure isO​\(n​K​d\)O\(nKd\)fornnpoints inℝd\\mathbb\{R\}^\{d\}\. In the language of coreset construction, this provides an\(α,β\)\(\\alpha,\\beta\)bicriteria solution withα=O​\(log⁡K\)\\alpha=O\(\\log K\)andβ=1\\beta=1\.

#### Sensitivity\-based coreset construction\.

Given this bicriteria solution, we follow the sensitivity\-sampling framework of Braverman et al\.\[[5](https://arxiv.org/html/2606.23898#bib.bib62)\]\. In particular, plugging the\(α,β\)\(\\alpha,\\beta\)assignment induced bykk\-means\+\+ into their coreset construction yields an efficiently computable upper bound on the sensitivity of each point\. Their framework samples points proportionally to such upper bounds, and Algorithm 1 constructs a weighted coreset whose size depends on the total sensitivity rather than directly on the original dataset size\.

In our implementation, after obtaining thekk\-means\+\+ centers, each pointppis assigned to its closest center, inducing a cluster𝒞​\(p\)\\mathcal\{C\}\(p\)with centerμ​\(𝒞​\(p\)\)\\mu\(\\mathcal\{C\}\(p\)\)\. We then use the following sensitivity upper bound:

s​\(p\)=4​\(8​\(log⁡K\+2\)\+1\)\|𝒞​\(p\)\|\+16​\(log⁡K\+2\)​‖p−μ​\(𝒞​\(p\)\)‖22∑q∈𝒞​\(p\)‖q−μ​\(𝒞​\(p\)\)‖22\.s\(p\)=\\frac\{4\\bigl\(8\(\\log K\+2\)\+1\\bigr\)\}\{\|\\mathcal\{C\}\(p\)\|\}\+\\frac\{16\(\\log K\+2\)\\,\\\|p\-\\mu\(\\mathcal\{C\}\(p\)\)\\\|\_\{2\}^\{2\}\}\{\\sum\_\{q\\in\\mathcal\{C\}\(p\)\}\\\|q\-\\mu\(\\mathcal\{C\}\(p\)\)\\\|\_\{2\}^\{2\}\}\.This is exactly the form obtained by combining thekk\-means\+\+ bicriteria guarantee with the sensitivity framework of Braverman et al\.\[[5](https://arxiv.org/html/2606.23898#bib.bib62)\]\. The first term increases the importance of small clusters, preventing them from being ignored, while the second term increases the importance of points that lie far from their assigned cluster center, which are precisely the points that contribute more to thekk\-means cost\. Sampling according to these sensitivities therefore yields a weighted subset that better preserves the geometry relevant to clustering than uniform subsampling\.

Braverman et al\.\[[5](https://arxiv.org/html/2606.23898#bib.bib62)\]further show that, once such sensitivity upper bounds are available, one can construct anε\\varepsilon\-coreset of size

O​\(tε2​\(d​log⁡t\+log⁡\(1/δ\)\)\),O\\\!\\left\(\\frac\{t\}\{\\varepsilon^\{2\}\}\\bigl\(d\\log t\+\\log\(1/\\delta\)\\bigr\)\\right\),wherettis the total sensitivity upper bound\. For thekk\-means setting induced by the bicriteria above, this gives a coreset size scaling on the order of

O​\(K​d​log⁡Kε2\)O\\\!\\left\(\\frac\{Kd\\log K\}\{\\varepsilon^\{2\}\}\\right\)up to logarithmic factors, and Theorem 34 in their supplement guarantees that the resulting weighted sample is anε\\varepsilon\-coreset with high probability\. Importantly for our use case, that theorem applies simultaneously to everyk′<Kk^\{\\prime\}<Konce the sensitivities are computed using a fixed upper boundKK\.

#### Selecting the number of clusters\.

After constructing the coreset, we sweep over candidate values ofKKand run weightedkk\-means on the coreset for each candidate\. For each resulting partition, we compute the silhouette score and select the value with the highest score\. In our experiments, we compute sensitivities once usingKm​a​x=1000K\_\{max\}=1000, construct the coreset with respect to this value, and then sweep over candidate numbers of clusters in the rangeK′∈\[200,1000\]K^\{\\prime\}\\in\[200,1000\]\. The final value ofKKis the one that maximizes the silhouette score on the coreset\. Once this value is chosen, we run the finalkk\-means clustering with thatKKto define the semantic regions used by our method\. Since the coreset is only used for model selection of the clustering granularity, and not as a replacement for the EMA mechanism itself, this preprocessing stage remains fully modular and can be replaced by other region\-construction procedures in future work\.

Algorithm 2Region construction via coreset\-based clustering1. 1\.Embed all prompts using the CLIP text encoder, obtainingX=\{xi\}i=1N⊂ℝdX=\\\{x\_\{i\}\\\}\_\{i=1\}^\{N\}\\subset\\mathbb\{R\}^\{d\}\.
2. 2\.Runkk\-means\+\+ withKmaxK\_\{\\max\}centers onXXto obtain a bicriteria solution\.
3. 3\.Compute sensitivity upper bounds using the bicriteria solution, and sample a weighted coreset\(S,w\)\(S,w\)\.
4. 4\.For each candidate number of clustersk∈𝒦k\\in\\mathcal\{K\}: 1. \(a\)run weightedkk\-means on\(S,w\)\(S,w\)withkkclusters; 2. \(b\)compute the corresponding silhouette scoreSil\(S,w\)​\(k\)\\mathrm\{Sil\}\_\{\(S,w\)\}\(k\)\.
5. 5\.Select k^=arg⁡maxk∈𝒦⁡Sil\(S,w\)​\(k\)\.\\hat\{k\}=\\arg\\max\_\{k\\in\\mathcal\{K\}\}\\mathrm\{Sil\}\_\{\(S,w\)\}\(k\)\.
6. 6\.Runkk\-means on the full embedding set withk=k^k=\\hat\{k\}and use the resulting partition to define the semantic regions\.

## Appendix EAblation Studies

We analyze several design choices in ARIA: \(i\) the mapping used to convert region discrepancy scores into sampling probabilities, \(ii\) the number of regions used to partition the conditioning space, \(iii\) the discrepancy signal used for adaptive allocation, and \(iv\) the behavior of the EMA\-based discrepancy tracker\. All ablations follow the same setup as Sec\.[5](https://arxiv.org/html/2606.23898#S5)and use the BK\-Base student model\.

#### Sampling mapping\.

ARIA converts region discrepancy scores into sampling probabilities using a mapping𝒮\\mathcal\{S\}\. We compare two choices: softmax \(pk∝exp⁡\(Dk\)p\_\{k\}\\propto\\exp\(D\_\{k\}\)\) and power\-law \(pk∝Dkp\_\{k\}\\propto D\_\{k\}\)\. Table[5](https://arxiv.org/html/2606.23898#A5.T5)shows that both mappings outperform uniform sampling\. Softmax yields the best FID and CLIP scores, suggesting moderate prioritization of difficult regions is more effective than overly concentrated allocation\.

#### Number of regions\.

We study the impact of the number of regionsKKused to partition the conditioning space\. The valueK=405K=405is selected by the clustering procedure described in Sec\. D in the Supp\. Matt\. As shown in Table[5](https://arxiv.org/html/2606.23898#A5.T5), ARIA improves over uniform sampling across all testedKK\. Moderate granularity \(K=405K=405\) yields the best FID and IS, while larger partitions \(K=2000K=2000\) slightly improve CLIP\. Performance remains stable across region counts, indicating that ARIA does not require precise tuning ofKK\.

Table 5:Ablation study of ARIA on the BK\-Base student model\. We analyze the effect of the sampling mapping, the number of regions used to partition the conditioning space, and the discrepancy signal used for adaptive allocation\.
#### Tracked discrepancy signal\.

ARIA reallocates training effort according to the magnitude of the teacher–student discrepancy\. We compare three allocation signals: the output\-level lossℒout\\mathcal\{L\}\_\{\\mathrm\{out\}\}, the feature\-level lossℒfeat\\mathcal\{L\}\_\{\\mathrm\{feat\}\}, and their combination\. As shown in Table[5](https://arxiv.org/html/2606.23898#A5.T5), ARIA improves over uniform sampling for all objectives\. However, the largest gains occur when allocation is driven byℒout\\mathcal\{L\}\_\{\\mathrm\{out\}\}, which directly reflects mismatch in the predicted score function\. Feature\-level discrepancies provide a weaker signal for adaptive allocation, as intermediate representations may remain aligned even when output\-level mismatch persists\.

#### EMA calibration of cluster difficulty\.

ARIA maintains cluster difficulty scoresDsD\_\{s\}using an exponential moving average \(EMA\), which determines the sampling distribution over regions\. To verify that these online estimates reflect the true student–teacher discrepancy, we compare them against Monte Carlo estimates of the per\-cluster loss at multiple training checkpoints on the C\-Tiny architecture\. For each of theS=405S\{=\}405clusters, we sample10001000prompts, pair them with random latents, and add noise at timestepttdrawn from the same replacement distributionp​\(t\)p\(t\)used during training\. We then evaluate the KD loss between the student and teacher and average the results to obtain the ground\-truth score\.

As shown in Fig\.[11](https://arxiv.org/html/2606.23898#A5.F11), the EMA estimates closely track ground truth throughout training, with relative error below4%4\\%across checkpoints\. Both quantities decay smoothly from∼2\.4×10−3\{\\sim\}2\.4\{\\times\}10^\{\-3\}to∼1\.0×10−3\{\\sim\}1\.0\{\\times\}10^\{\-3\}, confirming that the lightweight tracker provides a reliable signal for focal sampling\.

![Refer to caption](https://arxiv.org/html/2606.23898v1/x9.png)Figure 11:EMA calibration study\.\(a\) Comparison of the EMA\-tracked mean cluster scoreDsD\_\{s\}and the Monte Carlo ground\-truth lossJsJ\_\{s\}at four checkpoints \(25K, 100K, 225K, 300K\)\. Percentages denote relative error\. \(b\) Evolution of both quantities; the shaded band shows the interquartile range \(IQR\) of true per\-cluster losses\. The EMA tracker closely follows the KD loss throughout training, with relative error below4%4\\%\.