RAPID$^3$: Tri-Level Reinforced Acceleration Policies for Diffusion Transformer
Authors: Wangbo Zhao
Yizeng Han
Zhiwei Tang
Jiasheng Tang
Pengfei Zhou
Kai Wang
Bohan Zhuang
Zhangyang Wang
Fan Wang
Yang You
Venue: ICLR 2026 Poster (2026)
Keywords: Diffusion Transformer
Acceleration
Primary area: generative models

Abstract: Diffusion Transformers (DiTs) excel at visual generation yet remain hampered by slow sampling.  Existing training-free accelerators—step reduction, feature caching, and sparse attention—enhance inference speed but typically rely on a uniform heuristic or manually designed adaptive strategy for all images, leaving quality on the table. Alternatively, dynamic neural networks offer per-image adaptive acceleration, but their high fine-tuning costs limit broader applicability. To address these limitations, we introduce RAPID^3: Tri-Level Reinforced Acceleration Policies for Diffusion Transformer framework that delivers image-wise acceleration with zero updates to the base generator.  Specifically, three lightweight policy heads—Step-Skip, Cache-Reuse, and Sparse-Attention—observe the current denoising state and independently decide their corresponding speed-up at each timestep. All policy parameters are trained online via Group Relative Policy Optimization (GRPO) while the generator remains frozen.  Meanwhile, an adversarially learned discriminator augments the reward signal, discouraging reward hacking by boosting returns only when generated samples stay close to the original model’s distribution. Across state-of-the-art DiT backbones including Stable Diffusion 3 and FLUX, RAPID^3 achieves nearly 3$\times$ faster sampling with competitive generation quality.

Published as a conference paper at ICLR 2026
RAPID3: TRI-LEVEL REINFORCED ACCELERATION POLICIES FOR DIFFUSION TRANSFORMER
Wangbo Zhao 1+
Yizeng Han
Zhiwei Tang 2,3,4
Jiasheng Tang 2,31
Pengfei Zhou' 1
Kai Wang'
Bohan Zhuang2 2.4
Zhangyang Wang" Fan Wang?
Yang You' f
National University of Singapore
DDAMO Academy, Alibaba Group The University of Texas at Austin
*Hupan Lab
47IP Lab, Zhejiang University
ABSTRACT
Diffusion Transformers (DiTs) excel at visual generation yet remain hampered by slow sampling. Existing training-free accelerators-ster reduction, feature caching, and sparse attention-enhance inference speed but typically rely on a uniform heuristic manually designed adaptive strategy for all images, leaving quality on the table. Alternatively, dynamic neural networks offer per-image adap- tive acceleration, but their high fine-tuning costs limit broader applicability. To ad- dress these limitations, we introduce RAPID:: Tri-Level Reinforced Acceleration Pollcies for Diffussio Transformer framework that delivers image-wise accelera- tion with zero updates to the base generator. Specifically, three lightweight pol- icy heads-Stery Steuds--ep--kip.. Cache-Reuss Reuse, and Sparse-Attention -ttentioonnobbeeee the current denoising state and independently decide their corresponding speed-up at each timestep. All policy parameters are trained online via Group Relative Policy Optimization (GRPO) while the generator remains frozen. Meanwhile, an ad- versarially learned discriminator augments the reward signal, discouraging re- ward hacking by boosting returns only when generated samples stay close the original model's distribution. Across state-of-the-at DiT backbones includ- ing Stable Diffusion and FLUX, RAPID* achieves nearly 3x faster sampling with competitive generation quality. The code is publicly available at https: //gi thub. com/ NUS HPC-ATTCCTT ATw.ATTTTAAPTT..AA.
INTRODUCTION
Diffusion Transformers (DiTs) have emerged the dominant backbone for high-fidelity visual gen- eration thanks to their scalability and strong generalization (Peebles & Xie, 2023; Bao al., 2023). They now underpin tate-offthheaut systems in diverse downstream tasks-image synthesis (Chen et al., 2023; Esser et al., 2024; Labs, 2024), video generation (OpenAI, 2024; team, 2025), and con- trollable editing (Xiao et al., 2024; Feng et al., 2025). Despite this progress, DiTs remain inefficient during inference, requiring multiple denoising steps with computationally intensive blocks on large latent maps, making real-world deployment challenging.
A first family of training-free free accelerators (Figure 1 (a)) tackles this cost by hard-coded heuristics: reducing the step count (Song al., 2020; Lu et al., 2022a;b; Park al., 2024), reusing intermediate features (Ma et al., 2024b; Chen et al., 2024; Selvaraju et al., 2024; Liu et al., 2024), or sparsifying at- tention (Zhang et al., 2025a; Yuan et al., 2024b). These techniques preserve the frozen generator but apply a uniform or manually designed adaptive policy to every image and timestep; quality therefore fluctuates, and conservative settings are adopted to avoid artifacts, sacrificing potential speed-ups. A second family-dynam dynamic neural networks trained with fine-tuning (Figure 1 (b)) learns routers that adapt width, depth, or spatial resolution on the fly (Han al., 2021; Ganjdanesh et al., 2024; Zhao et al., 2024; You et al., 2024; Anagnostidis et al., 2025; Zhao et al., 2026). While these methods of- fer adaptability and high performance, they require extensive optimization on large-scale image-text datasets, making them impractical for many large generators proprictary models.
Work done during Wangho's 's internship at DAMO Academy, Alihaba Group. wangbo.zhaa96@@gmii.. Cooresponding author.

---

Published as a conference paper at ICLR 2026
(a) training-free methods cache reuse
(b) dynamic neural networks
(c) RAPID* (ours)
step skip (..0)
sparse attention
(6.C2..
(6.Ch2,N
(6:6]+
(6,C))
(e. c)
(c.c),
policias
policies
%6
generator
frozen generator
frozen generator
pre-image adaptivity
pre-image adaptivity
pre-image adaptivity
frozen
trainable
original model
cache
reuse
w/ sperse
attention
dynamic
model
inistal noise
condition
Figure 1: Accelerating Diffusion Transformers: (a) Training-free methods primarily use uniform or manually designed rules (e.g., step skip, cache reuse, or sparse attention) for all images and timesteps, offering speed but little adaptivity. (b) Dynamic fine-tuned models learn routers that tailor acceleration to each input but require costly fine-tuning the generator's parameters. (c) Pro- posed RAPID* keeps the generator frozen and equips it with three lightweight policy heads Sredd-Sepp- Skip, Cache- and Sparse-Artention -rrined via reinforcement learning to make action deci- sions based the latent representations, timestep, and text prompt information.
Sampling in a diffusion transformer, however, naturally a Markov decision process: its forward pass unfolds as a sequence of denoising steps. we consider the latent generated each step step-along with its timestep index and prompt embedding-an the state, then the selection of the next timestep continuous action, while decisions on whether to reuse cached features or apply
can be treated sparse attention are regarded as discrete actions. Since the generator's 's weights remain fixed, the outcome of any action deterministic and fully predictable by the model itself, eliminating the need to learn environment dynamics. Episodes are inherently brief, typically spanning multiple sampling steps, and conclude with complete image whose quality can be quantitatively assessed post hoc, enabling straightforward reward signal that balances image fidelity against computational cost.
Our key insight is that the DiT inference setting, when viewed through the above lens, sits squarely in the comfort zone of modern policy-optimization methods: concise but expressive action space; a stationary and deterministic environment; an inexpensive simulator (the frozen DiT) that delivers abundant rollouts; and scalar reward that directly reflects the user's objective. By leveraging this structure, we can train lightweight policies to dynamically determine which acceleration strategy to invoke invoke-enabling adaptivity without modifying the generator's parameters. The recent TPDM (Ye et al., 2024) relates to our idea: it achieves data-dependent step skipping via reinforcement learning (RL), while negleting the redundant computation from the intra-timeste timestep perspective.
Motivated by this observation, we present RAPID (Figure 1 (c))--Ti-LLeel Reinforced Acceleration Pollcies for Diffusion Transformer. attaches three policy htads-Step-Skip. Cache- Reuse, and Sparse Sparse-Attention Atteerioontttittottoooo to frozen DiT. Each head observes summaries of the latent, timestep, or prompt, and independently decides corresponding acceleration strategies. All policy parameters are trained online with Group Relative Policy Optimisation (GRPO) (Shao et al., 2024). Moreover, instead of using existing evaluation metrics to build the reward function, we train discriminator adversarially to enhance the reward for samples that remain close to the original generator's 's distri- bution, which prevents the reward hacking problem (Skalse et al., 2022).
Experimental results show that RAPID achieves nearly 3x acceleration for Stable Diffusion 3 (Esser et al., 2024) and FLUX (Labs, 2024), while maintaining competitive visual quality. By updating only the policy head, which constitutes merely 0.025% of the generator's parameters, its training process relies solely on readily available text prompts and and consumes only 1% of the GPU hours required for training dynamic neural networks (Zhao al., 2026).
2

---

Published as a conference paper at ICLR 2026
2 RELATED WORK
Diffusion transformers. Diffusion models (Sohl-Dickstein et al., 2015; Ho et al., 2020; Albergo et al., 2023) have achieved remarkable success visual generation tasks (Rombach et al., 2022: OpenAI, 2024). Early diffusion models (Ho et al., 2020; Dhariwal & Nichol, 2021; Rombach al., 2022) are primarily built on U-Net (Ronneberger et al., 2015). However, Transformer based archi- tectures (Vaswani et al., 2017) have gradually replaced U-Net the foundation for these models. DiT (Peebles & Xie, 2023) represents one of the first attempts to integrate Transformers into the diffusion process. Concurrently, U-VIT (Bao al., 2023) combines the advantages of U-Net's skip connections with the capabilities of Transformer architectures. Building on DīT, PixArt-a expands its capabilities to text-to-image generation, while SD3 (Esser et al., 2024) and FLUX (Labs, 2024) further demonstrate its scalability. Beyond image generation, Transformer- els, such as Sora (OpenAI, 2024) and WanX (WanTeamm et al., 2025), have also shown significant promise in video generation. These progresses in foundation models have also facilitated the de- velopment of various applications, such as image and video editing (Ma et al., 2024c;d; 2025e;d;a). Despite these achievements, diffusion models continue suffer from high computational demands and tinecconsuiing generation. The proposed framework addresses this limitation by adaptively adjusting computation for each image generation process, enabling more efficient visual generation.
diffusion mod-
Diffusion transformer acceleration. To enhance inference efficiency of DiTs, researchers have ex- plored training-frres ffree techniques such reducing steps (Song et al., 2020; Lu et al., 2022a; Park et al., 2024; Lu et al., 2022b), feature caching (Ma et al., 2024b; Chen et al., 2024; Liu et al., 2024; Huang et al., 2024), and sparse attention mechanisms (Zhang et al., 2025a;b; Xi et al., 2025). However, these approaches typically use uniform strategies for all images and the entire generation process or depend on manually designed heuristics, limiting their generalization. Inspired by the efficiency gains of dynamic neural networks (Han et al., 2021; Verma et al., 2024; Zhao et al., 2024; 2025a; Wang et al., 2025), some studies (Ganjdanesh et al., 2024; Zhao et al., 2025b; 2026) introduce dy- namic architectures to improve efficiency, while others (You et al., 2024; Anagnostidis et al., 2025) dynamically reduce computation along the spatial dimension. However, these methods require fine- tuning diffusion models, which imposes a significant training burden and becomes impractical when training data is limited. Few-step distillation techniques (Luo al., 2023; Yan et al., 2024; Lin al., 2024; Yin et al., 2024; Shao et al., 2025) accelerate inference but often require extensive parameter tuning and training of around 100M additional parameters, even with methods like LoRA (Hu et al., 2022). In contrast, our approach trains lightweight policy heads with only 3M additional parameters, adaptively selecting acceleration strategies while keeping DiT parameters frozen. This ensures both parameter and data efficiency, offering more practical and elegant solution.
Reinforcement learning in diffusion models. The success of reinforcement learning (RL. in large language models (LLMs) (Ouyang et al., 2022: Shao et al., 2024; Guo et al., 2025; Team, 2025) has encouraged researchers to explore it in diffusion models. DDPO (Black et al., 2023) views denoising as multi-step decision-making task, aligning generated images with downstream objectives. Simi- larly, DPOK (Fan et al., 2023) demonstrates the effectiveness of RL in fine-tuning diffusion models, achieving superior text-image alignment and improved image fidelity. Some approaches (Wallace et al., 2024; Yuan et al., 2024a) have also explored the alignment of generation quality with human preferences in an offline RL framework. However, these methods primarily focus on enhancing the generation quality rather than accelerating the generation speed. TPDM (Ye et al., 2024) allevi- ates this problem by learning an efficient noise scheduler, without exploring its incorporation with other acceleration techniques. Our approach leverages RI. to learn policy heads that dynamically determine step-skipping, cache-reuse, and sparse-attention strategies at each timestep, adaptively improving the generation efficiency.
METHODOLOGY
We begin with foundational concepts of DiT and RL in Section 3. Then, we present the proposed tri-level reinforced acceleration policies and their corresponding policy heads in Section 3.2 and Section 3.3, respectively. Section 3.4 details the adversarial reinforcement learning within the the proposed framework. A comprehensive overview of the methodology is provided in Figure 2.
3

---

Published as a conference paper at ICLR 2026
3.1 PRELIMINARY
Architecture of diffusion transtormers. Diffusion Transformers (DiTs) (Peebles & Xie, 2023; Bao et al., 2023; Chen et al., 2023; Esser et al., 2024; Labs, 2024) has demonstrated remarkable advancements in visual generation thanks to its scalable architecture. It generally consists of stack of layers, each of which integrates a self-attention (SA) block and a multi-layer perceptron (MLP) block. The operation of DiT block can be roughly expressed as:
x cttt F'AVTRR (x/, c)) FM (x) c))).
where x E TRN denotes the input to the 1-th layer at timestep I. Here, N represents the number of tokens and the channel dimension. The condition embedding cf combines timestep information with text-based guidance, which are critical for generation.
Reinforcement learning with Group Relative Policy Optimization. To improve the perfor- mance and alignment of LLMS with human preferences, RI. techniques such as Proximal Policy Optimization (PPO) (Schulman et al., 2017) are widely applied during fine-tuning. To reduce the training cost, Group Relative Policy Optimization (GRPO) (Shao et al., 2024) is proposed, which leverages a group-based strategy to estimate advantages. Specifically, for a given question q from the training set, GRPO samples a group of output answers {o.),- from LLM model Tea with fixed parameters. The target policy Te is then updated by maximizing the following objective:J<Roo(0) = t Ei= (min =u(o=lo) Te (o.19) A;, clip Tu(oolo)
To 93 1 E, 1 +8) A.) -9DDI.).
where and 3 are hyperparameters. The term DKL represents KL-divergence penalty between the current model Ty and the pre-trained reference model Tye-- The advantage A, is computed from a group of rewards {ri] as A; T - meen((rr.....
std([r1:f2."" (aa) Due to its simplicity and effectiveness, integrate GRPO into the generation process of diffusion transformers to train policy heads, as discussed later in Section 3.4.
3.2 TRI-LEVEL REINFORCED ACCELERATION POLICIES
As previously outlined, our goal find policies that dynamically selects acceleration strategies for visual generation. To achieve this, we define three levels of candidate strategies to accelerate from model-external to model-internal perspectives: Step-Skip. Cache-Reuse, and Sparse- Atfention. These strategies are effective and training-free acceleration techniques, but existing approaches often apply them with uniform or manually designed adaptive policies.
Level-1: Step-Skip. The diffusion Transformer's generation involves multi-timester, schedule, where the timestep decreases from tr to to- This process' efficiency is critically dependent on the required of timestep number. A natural idea is to reduce the needed timesteps. However, prior meth- ods (Song al., 2020; Lu et al., 2022a; Park et al., 2024; Lu et al., 2022b) employ fixed schedules, ignoring that different images may require varying timesteps. For instance, highly detailed images often demand more timesteps to achieve high-quality outputs, whereas simpler images can be gen- erated effectively with fewer steps. Applying uniform schedule across all cases risks degrading both generation quality and computational efficiency.
To achieve inference-time dynamic timestep jumping, define policy putepp that selects the current timestep f. based on the state of the previous timestep fprev- Specifically, let G(-) denote the forward computation of DIT, sampler(-) represent the diffusion sampling function, and X X indicate the latent feature at the previous timestep. The feature at the current timestep. X,, can be formulated as
X Sampler(X.,... 9(X (prey) t).
(2)
This approach allows the generation process to adjust the number of required timesteps per input, enabling varying computational resource allocation depending on the the target complexity. As result, the diffusion process achieves a balance between generation quality and efficiency.
Level-2: Cache-Reuse. Feature caching mechanisms (Ma et al., 2024b; Chen et al., 2024; Selvaraju et al., 2024; Liu et al., 2024) leverage temporal coherence in diffusion processes by reusing com- puted feature maps across consecutive timesteps. A prevalent approach in diffusion transformers in-
4

---

Published as a conference paper at ICLR 2026
Diffusion Transformer
O
Diffusion Transformer
sampler
sampler
0 Fyevv
Syurr G Fpeer
stop
Cache
sparse
stop
feature flow
policy heads
pollcy
action flow
X
Cachi
heads
sampling
(a) accelerate with RAPID' during inference
0,
pstep
Buta(a,
peacho
Cat. (pp
X
pepans
Cal. (p, (parse)
pparse
(b) details policy heads
Figure 2:
Overview of RAPID* Accelerate generation with RAPID* during inference. (b) The details of policy heads {psiep, piuche prpuuee. ).
volves caching the residual components of feature maps (Chen et al., 2024; Liu et al., 2024) specif- ically, the difference between model's input and output. For timestep t, this enables computational simplification through reuse of the residual cached at timestep feache expressed as O, R= X, Leate where O, denotes the output of the diffusion transformer and A.... = 10XXttttte leache) resents the cached residual.
Xtuaee rep-
Two fundamental challenges emerge in practical implementations: a) The trade-off between accel- eration and generation quality exhibits strong dependence the temporal interval between cache updates. Current approaches relying on fixed or heuristic update intervals demonstrate limited gen- cralizability, particularly when combined with other acceleration techniques. b) Generating different images may require distinct caching schedules, posing challenges to designing and optimization.
Hence, the 2nd level of our objective centers on developing an adaptive policy peuche that dynami- cally selects optimal caching strategies per timestep t. Formally, the policy determines the compu- tation path through the following conditional operation:
O2 = ((X,, t) if update cache,
X, + A reuse cache.
If the policy model determines to update the cache, we will conduct the computation G(X,, () as the original diffusion transformer and update the cache to A.-
Level-3: Sparse-Attention. The quadratic computational complexity of self-attention mechanisms in transformers (O()2) for n tokens) poses significant latency challenges for high-resolution image generation. To mitigate this bottleneck, recent work (Zhang et al., 2025a; Yuan et al., 2024b) has integrated sparse attention mechanisms into diffusion transformers through conditional computation of attention weights. These approaches prune the computation in attention maps by employing a hyperparameter 9 E_.x to control sparsity for each layer and attention head, where and H represent the number of layers and attention heads, respectively. Existing implementations, such as SpargeAttn (Zhang et al., 2025a), search for over the entire diffusion process of DiT and keep fixed across all timesteps during generation.
However, this static configuration ignores the evolution of attention patterns during the diffusion process. During the diffusion process, different timesteps exhibit distinct properties and may vary in their sensitivity to sparsity. Fixing the 0 could sacrifice potential speed-up or lead to quality degradation. More importantly, the 0 optimized for the original DiT may become unsuitable when DiT is combined with other acceleration methods.
To address this, we propose a dynamic A policy paparse to adaptively identify a optimal hyper- paremter for the current timestep. 0, e from series of pre-defined candidate hyperparemters (01, N..00 The original self-attention block can then be replaced with sparse attention using the identified hyperparameter, which is formulated as:
Fsa (x cf)
SApans (x c/i 0').
(4)
This approach enables the adaptation of computation sparsity on a per-timestep basis.
5

---

Published as a conference paper at ICLR 2026
3.3 DESIGN OF POLICY HEADS
The learning of our tri-level acceleration framework involves joint optimization of three policy heads pakop, ppache PPspuse They determine the acceleration strategy based on the state the current timestep 1. Each policy head begins with a convolution layer for feature projection, followed by AdaLN (Perez et al., 2018) integrating the condition embedding et. The output finally pooled and fed to the corresponding linear head for action prediction. Each policy head is detailed as follows:
a) Step-Skip policy pslep, It predicts jump steps via parametric distribution learning:
(5) where O, denotes the output of the diffusion transformer at the current timestep Inspired by (Ye et al.. 2024), Pseo uses a linear layer to regress two values, and B, which parameterize a Beta distribution. A value from this distribution is subsequently used compute the next timestep fnext-
paup O2 [0, step Beta(a, B), Inxx = -L.. step
b) Cache-Reuse policy pouche This discrete decision network determines whether to perform computation or reuse the cached result based on the discrepancy between the inputs at the current timestep X, and the last cached timestep Xtovr. This process can be expressed as:
peache X
Xtemm + p
cache R2, df eache ~ Categorical (paace))
(6)
with action semantics: afache = 0 indicates that computation should be performed at the current timestep to update the cache, while at aaache correspond reusing the cache.
c) Sparse-Attenntion policy paparse This policy adapts computation through discrete choices:
P Pypprre: X, 1+ P.
sparse E N 04 sparse Categorical (p.ure at
where sparse O indicates full attention where no sparsity is applied, while sparse E [1, N. sparsel de-
notes applying sparse attention with different / values that progressively increase the sparsity ratios. Here, N spurse is set to 3 by default. Although we experiment with additional candidate hyperparam- eters, we observe that an overly aggressive sparsity strategy significantly degrades the generation quality. Notably, Piipame remains inactive when peache decides to reuse the cache.
3.4 ADVERSARIAL REINFORCEMEN LEARNING
To train the policy heads, we propose an adversarial reinforcement learning framework, inspired by adversarial training (Goodfellow et al., 2020; Lin et al., 2024; Sauer et al., 2024; Lin et al., 2025). This approach iteratively trains the policy heads alongside a discriminator model, effectively addressing the challenging issue of reward hacking. The training pipeline is outlined in Algorithm 1.
Training policy heads with reinforcement learning. To train the policy heads, we adopt the diffu- sion transformer G equipped with policy heads (pstep, ,peahee ppspare) } to sample group of images (1:) based on the same text prompt. Subsequently, the pre-trained image reward model Q is employed to evaluate the quality and alignment of each image with the given text condition c, pro- ducing scores {q Simultancously, the discriminator model D (details provided later), is used to estimate the likelihood that generated image originates from the diffusion transformer without acceleration, producing scores {d; G 1=1- This framework encourages the policy heads to identify acceleration strategies that ensure the accelerated model performs comparably to the original.
To estimate the generation cost, we define the concept of equivalent steps K, expressed as K -E Cgachog x (1 Capans)) where Keeep is the total number of steps during gener- ation, determined by peep_ Here, Csache and Cspprse represent the cost reductions (normalized into [0, 1I) achieved by reusing cache and applying sparse attention at the k-th step, respectively. In practice, we round K to its nearest integer. The final reward for an image is then expressed as ><= + wd,), where (0, 1) a decay factor that penalizes higher generation
r; = 1. , costs and w denotes the weight of the discriminator in reward.
k
We can adopt the equation Section 3.1 to obtain the each sample's advantage A,. Based on the formulation of GRPO, our training process can be formulated as:
JORPO-RAPID = G (min (e.A, clip (0r, 1,,,1-1,1 e,1+e)
(8)
6

---

Published as a conference paper at ICLR 2026
where Φ; no (T(()) Tuu (J: Here, Te (I: c) represents the likelihood of generating the image I given the text condition c, under the current parameters 8 of the three policy heads, while Te.. uses the old parameters. The KL divergence term is omitted because the policy heads are randomly initialized, and therefore there is no reference model.
Training the discriminator model. Relying solely the image reward model Q may lead policy heads to exploit or "hack" the reward model, rather than genuinely improving acceleration strategies. To address this, introduce the discriminator model 'D, binary classification model designed to distinguish between images generated by the diffusion transformer G with and without acceleration strategies applied. It enhances the reward of accelerated samples from & that remain close those without acceleration. This makes the discriminator complementary to the reward model, ensuring a more robust training process.
To achieve this, we first allow the generation model G, without acceleration, to sample images and construct the positive dataset Zorigin Next, we initialize the policy heads to accelerate g, sample images again, and construct the negative dataset Iacolle Both datasets are then employed to train the discriminator. During the reinforcement learning process, the images sampled using GRPO can be used to update the negative dataset Iacelle Cross entropy loss is employed during training. As mentioned in the previous paragraph, we employ the discriminator model as an additional reward model to provide supplementary reward signal, which is combined with the reward from the image reward model Q to produce the final reward signal.
This adversarial training process enables joint improvement of the policy models and discriminator. The policy heads are trained to assist G in efficiently producing high-quality images that can fool the discriminator, while the discriminator is optimized to better distinguish between data generated by g with and without acceleration.
EXPERIMENT
Model configurations. We conduct experiments using two diffusion transformers: Stable Diffusion 3 (SD3) (Esser et al., 2024) and its larger counterpart, FLUX (Labs, 2024). For the reward model, we adopt ImageReward (Xu et al., 2023). For the discriminator, we use pre-trained CLIP (Radford et al., 2021) and inset adapters (Chen et al., 2022) for parameter-efficient training. The hyperparameters for the baseline approaches and our method are detailed in Section G and Section H, respectively.
Datasets and evaluation. We use 20K prompts randomly sampled from Byeon et al. (2022) and COCO 2017 (Lin et al., 2014) training set to train our policy heads. For evaluation, we use 5,000 images with prompts from the COCO 2017 validation set, along with metrics such as CLIP-Score (Zhengwentai, 2023), and Aesthetic v2 (Schuhmann, 2022) for a comprehensive as- sessment. Additionally, two comprehensive benchmarks, HPS (Wu et al., 2023) with 1,600 prompts and GenEval Ghosh et al. (2023) with 553 prompts, are introduced to demonstrate the generalization capability of our method. We report the latency measured on an NVIDIA H20 GPU.
4.1 RESULTS ON SD3
We compare RAPID* with other acceleration techniques, including common step-reduction, two cache reuse methods-A-Aa (Chen et al., 2024) and TeaCache (Liu et al., 2024) and sparse attention mechanism, SpargeAtttn (Zhang et al., 2025a). A-DiT employs a uniform interval for all image generations, while TeaCache introduces a manually set threshold to control accumulation errors, enabling a certain degree of adaptivity. For our method, we set the decay factor in our method is set to 0.97, resulting in around 2.92 acceleration. From Table 1, we observe that RAPID achieves the best balance across all metrics while delivering the highest acceleration ratio compared to its counterparts. This demonstrates that the learned adaptive policies outperform both uniform policies and manually designed adaptive policies, validating the effectiveness of our approach.
Additionally, we compare our method with TPDM (Ye et al., 2024), which only adjusts the genera- tion steps per image. The results show that our method surpasses TPDM across all metrics, validat- ing the superiority of our tri-level acceleration design compared single-strategy approach.
7

---

Published as a conference paper at ICLR 2026
Table 1: Results on SD3 (Esser al., 2024). Bold highlights the best results across various ac-
celeration methods. We report latency of the diffusion transformer, excluding the text encoder and VAE decoder. Additional comparison results are provided in Table 8.
COCO
HPS Score
GenEval
Method
Latency (s) + Speed f
CLIP 1
Aesthetie ?
Correct
Overall f
SD3 28-steps
5.77
1.00 x
32.05
5.31
28.83
67.81
69.01
static manually designed adaptive acceleration methods
SD3 9-steps
1.98
2.91 x
31,88
5.21
27.67
60.58
61.67
w/ TeaCache 6=0.15
2.20
2.62 x
32.02
5.25
27.87
61.48
62.81
w/ A-DiT N
3.76
1.53 x
31.91
5.12
27.67
57.69
58.67
w/ SpargeAttn
5.08
1.13 x
31.39
5.02
27.16
43.94
45.01
dynamic acceleration methods
TPDM w/ RAPID' (Ours)
2.32
2.48 x
31.98
5.25
27.75
59.67
60.70
1.97
2.92
32.09
5.26
28.07
62.57
63.48
bunch of cumpkins that carved different faces them.
soood +0.01 Acothetic
X2. S1 spordd
3.01 Aestseti
COCD-Aesthelic
slope.
9-2hpps
5.53
Wi SpargeAthh
wi Twetlache
wi RAPID' (ours)
miksing
Untanaali TeaCache x2.61
5.43
10
13
16
19
22
FLUX 28-steps x1.0
FLUX 9-steps x2.86
SpargeAttn x1.15
RAPID* (Ours) x2.91
latency (9)
Figure 3: Latency vs. quality trade-off.
Figure 4: Visual comparison. The acceleration ratio rel- ative to the original FL.UX 28-steps is reported.
4.2 SCALE UP TO FLUX
Comparison with other acceleration strategies. In Figure 3. we compare the trade-off between generation latency and quality across various acceleration methods on COCO. We set in RAPID to 0.97 and 0.90, respectively, to achieve varying acceleration ratios. Both our method and Tea- Cache (Liu al., 2024), as well as directly reducing steps, significantly accelerate generation speed. However, our method demonstrates greater robustness in maintaining performance, verifying its effectiveness in larger diffusion transformers.
In Figure 4, we compare our approach with other acceleration techniques in terms of visual quality. For this comparison, SpargeAttn, TeaCache, and our method use the fastest point on curve from Figure 3. The results demonstrate that method better preserves visual quality, highlighting the importance of dynamically selecting acceleration strategies.
Comparison with dynamic model. DyFLUX (Zhao et al., 2026) extends DyDiT (Zhao et al., 2024) FLUX (Labs, 2024). It introduces more sophisticated training process to fine-tune the generator, enabling it to learn an acceleration strategy with per-image adaptivity. As shown in Table 2, our RAPID achieves superior generation quality while also delivering faster inference. Notably, in terms of training cost-bot data and computational efficiency method outperforms DyFLUX significantly, highlighting its clear superiority.
Visualization. In Figure 5, illustrate the images generated by the original FLUX and our RAPID*. To facilitate a clear comparison with the original 28-step generation process, we use the equivalent step as metric, defined in Section 3.4, to represent the generation time cost. For images with a single object and a simple scenario (e.g. the image in the first column), our method requires fewer equivalent steps. Conversely, for images with multiple objects and complex scenarios (e.g.
8

---

Published as a conference paper at ICLR 2026
Comparison be- tween RAPID' and DyFLUX. Compared to DyFLUX (Zhao et al., 2026), our method, sig- nificantly improves inference speed, maintains generation quality competitive with the original FLUX, and requires substantially less training cost.
Table 2:
training
inference lateney (s) & Aesthetic
method
GPU hours +
data 4
FLUX
22.15
5.64
DyFLUX
38,000
3M image-text
13.93
5.29
RAPID' (Ours)
400
20K text only
5.63
R/ %
0.7%
-40%
+0.34
tmestep (b) FLUX W/ method (a) FL)X-X288s
1.0
0.8
0.6
0.4
0.2
Liggis slepems usng cande
Luing cacra Stepardd
10
15
5
10 15
10
15
5
10
15
0
5 10
15
20
step index
step index
step Index
step Index
step index
Figure 5: Visual comparison based on FLUX. (a) Images generated by the origin FLUX with the default 28 steps. (b) Images generated by FLUX accelerated with RAPID* (c) The generation processes for (b), where equivalent steps represent the generation time cost for each image.
the image in the last column), it requires more equivalent steps. This demonstrates that method has learned to adaptively adjust its acceleration strategy based on the complexity of each image. Additional visualizations provided in Section J.
4.3 ANALYSIS
For fairness, we adjust the hyperparameters to ensure similar equivalent steps across all experiments. The color denotes the default setting of RAPID*. All evaluations conducted COCO.
Effectiveness of different candidate acceleration strategies. Table 3, we present the results of method using different candidate acceleration strategies, including step-skip, cache-reuse, and sparse-attention, referred to as "step", "cache", and "sparse", respectively. The results show that progressively incorporating these three acceleration strategies leads to increasingly improved generation performance (e.g. CLIP score). Notably, RAPID* with our default setting, which uses all three strategies, achieves the best average performance. This improvement be attributed to the expanded solution space provided by more acceleration strategies, enabling RAPID3 to identify better acceleration approach for each image generation process.
Effectiveness the reward model and discriminator model. We individually remove the reward model Q and the discriminator D from the training of our method to evaluate their effectiveness.

---

Published as a conference paper at ICLR 2026
Table 3: Effectiveness of dif- ferent candidate acceleration strategies. The symbol / indi- cates that the corresponding strat- egy is used in our method.
strategy cache
COCO
HPS Score t
step
sparse
CLIP
Aesthetic
31.98
5.25
27.75
32.04
5.25
27.86
/
31.96
5.13
27.80
32.09
5.26
28.07
Table Effectiveness of reward model Q and discriminator model D. In + w controlling the weight of the reward from D is set to 1.0. IR de- notes the score from ImageReward (Xu al., 2023).
COCO
HPS
Method
IR
CLIP
Aesthetic f
Score
only
0.9605
32.04
5.18
27.72
only Q D
0.9538
31.91
5.24
27.96
0.9574
32.09
5.26
28.07
Results are presented in Table 4. We additionally include the ImageReward score (IR), used during training, additional evaluation metric. When only the reward model Q is employed, training process degrades to standard reinforcement learning with GRPO, achieving the highest ImageRe- ward score. However, its performance on other metrics drops significantly, relies on the leageReward score, making susceptible to reward hacking. As a result, the policy model focuses solely optimizing the reward from Q, neglecting the actual generation quality.
the reward siginal
In contrast, our method incorporates the discriminator D into the training process, effectively al- leviating the reward hacking problem and maintaining strong performance across various metries, highlighting the importance of the proposed adversarial reinforcement learning. For completeness, also conduct an experiment using only the discriminator D during training and find that it does not outperform our method, further highlighting the importance of integrating both Q and D.
CONCLUSION
In this study, we address the challenge of accelerating diffusion transformers in per-image adaptive manner without modifyyng their parameters. To achieve this, we introduce RAPID": Tri-Level Reinforced Accelcration Policies for Diffusion Transformer. RAPID" employs three lightweight policy heads, optimized via Group Relative Policy Optimization, to select Step-Skip. Cache-Reuse, and Sparse-Attention strategies at each timestep, significantly improving the generation speed. To mitigate reward hacking problem, incorporate adversarially learned discriminator to ensure robust policy learning. Extensive experiments demonstrate the effectiveness of RAPID*, and we anticipate our method will inspire further advancements in accelerating diffusion transformers.
Limitations and future work. Our method still relies on training to learn acceleration policies. Incorporating prior knowledge adaptively select acceleration strategies could further ease the training burden. Additionally, extending the proposed method to video generation models (Zheng al., 2024; Yang al., 2024; WanTeam al., 2025) and editing models (Ma al., 2025a;b;c; 2023; 2022) warrants further exploration in the future.
Acknowledgment. This work supported by Damo Academy through Damo Academy Research Intern Program. The computational work involved in this research work is also supported by NUS IT's Research Computing group using grant numbers CFP02-CF-004. Yang You's research group is being sponsored by NUS startup grant (Presidential Young Professorship), Singapore MOE Tier- grant, ByteDance grant, ARCTIC grant, SMI grant (WBS number: A8001 10400000-,0-)0), Alibaba grant, and Google grant for TPU usage.
REFERENCES
Arash Ahmadian, Chris Cremer, Matthias Gallé, Marzich Fadace, Julia Kreutzer, Olivier Pietquin, Ahmet Üstün, and Sara Hooker. Back basics: Revisiting reinforce style optimization for learn-
10

---

Published as a conference paper at ICLR 2026
ing from human feedback in llms. arXiv preprint arXhv:2402.14740, 2024.
Michael S Albergo, Nicholas M Boffi and Eric Vanden-Fijnden Stochastic interpolants: A unifying framework for flows and diffusions. arXiv preprint arXiv:2333300879, 2023.
Sotiris Anagnostidis, Gregor Bachmann, Yeongmin Kim, Jonas Kohler, Markos Georgopoulos. Art- siom Sanakoyeu, Yuming Du, Albert Pumarola, Ali Thabet, and Edgar Schönfeld. Flexidit: Your diffusion transformer can easily generate high-quality samples with less compute. arXiv preprint arXiv:2502.200 20/26, 2025.
Fan Bao, Shen Nie, Kaiwen Xue, Yue Cao, Chongxuan Li, Hang Su, and Jun Zhu. All are worth words: A vit backbone for diffusiin models. In CVPR, pp. 22669-226777 2023.
Kevin Black, Michael Janner, Yilun Du, llya Kostrikov, and Sergey Levine. Training diffusion models with reinforcement learning. arXiv preprint arXiv:2305. 12255.33300 13301, 2023.
Daniel Bolya and Judy Hoffman. Token merging for fast stable diffusion. In Proceedings of the IEEE/CVF conference on computer vision and patlern recognition, pp. 4599-4603, 2023.
Minwoo Byeon, Beomhee Park, Haecheon Kim, Sungjun Lee, Woonhyuk Baek, and Saehoon Kim. Coyo-700m: Image-text pair dataset. https: http://ttttt...oh...ee wtt.////aaaaaaa/////aaoo..
coyo-dataset, 2022.
Junsong Chen, Jincheng Yu, Chongjian Ge, Lewei Yao, Enze Xie, Yue Wu, Zhongdao Wang, James Kwok, Ping Luo. Huchuan Lu, et al. Pixart-alpha: Fast training of diffusion transformer for photorealistic text-to-image synthesis. arXiv preprint arXiv:2310.00426, 2023.
Pengtao Chen, Mingzhu Shen, Peng Ye, Jianjian Cao, Chongjun Tu, Christos-Savvar Bouganis, Yiren Zhao, and Tao Chen. Delta-dit: A training-free acceleration method tailored for diffusion transformers. arXiv preprint arXiv:2406.01125, 2024.
Shoufa Chen, Chongjian Ge, Zhan Tong, Jiangliu Wang, Yibing Song, Jue Wang, and Ping Luo. Adaptformer: Adapting vision transformers for scalable visual recognition. NeurIPS, 35:16664- 16678, 2022.
Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. NeurIPS, 34:8780-8788 2021.
Patrick Esser, Sumith Kulal, Andreas Blattmann, Rahim Entezari, Jonas Müller, Harry Saini, Yam Levi, Dominik Lorenz, Axel Sauer, Frederic Boesel, et al. Scaling rectified flow transformers for high-resolution image synthesis. In ICML, 2024.
Ying Fan, Olivia Watkins, Yuqing Du, Hao Liu, Moonkyung Ryu, Craig Boutilier, Pieter Abbeel, Mohammad Ghavamzadeh, Kangwook Lee, and Kimin Lee. Dpok: Reinforcement learning for fine-tuning text-to-image diffusion models. NeurlPS, (6.7988--79883 2023.
Haipeng Fang, Sheng Tang, Juan Cao, Enshuo Zhang, Fan Tang, and Tong- Lee. Attend to not attended: Structure-then-detail token merging for post-training dit acceleration. In CVPR, pp. 18083-18092, 2025.
Kunyu Feng, Yue Ma, Bingyuan Wang, Chenyang Qi, Haozhe Chen, Qifeng Chen, and Zeyu Wang. Dit4edit: Diffusion transformer for image editing. AAAI, volume 39, pp. 2969-2977, 2025.
Alireza Ganjdanesh, Yan Kang, Yuchen Liu, Richard Zhang, Zhe Lin, and Heng Huang. Mixture of efficient diffusion experts through automatic interval and sub-network selection. In ECCV, pp. 54-71. Springer, 2024.
Dhruba Ghosh, Hannaneh Hajishirzi, and Ludwig Schmidt. Geneval: An object-focused framework for evaluating text-to-image alignment. Advances Neural Information Processing Systems, 36: 52132-52152, 2023.
Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial networks. Communications of the ACM, 63(11 (::111)1339111 1444 2020.
11

---

Published as a conference paper at ICLR 2026
Daya Guo, Dejian Yang, Haowei Zhang, Junxiao Song, Ruoyu Zhang, Runxin Xu, Qihao Zhu, Shirong Ma, Peiyi Wang, Xiao Bi, et al. Deepseek-rl: Incentivizing reasoning capability in llms via reinforcement learning. arXiv preprint arXiv:2501. 2025.
Yizeng Han, Gao Huang, Shiji Song, Yang, Honghui Wang, and Yulin Wang. Dynamic neural networks: A survey. TPAMI, 44(11):7436-744 2021.
Jonathan Ho and Tim Salimans.
Classifier-free diffusion guidance.
arXiv preprint
0rXiv:2207.12598, 2022.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. NeurIPS, 33: 6840-6851, 2020.
Edward J Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, Weizhu Chen, et al. Lora: Low-rank adaptation of large language models. ICLR, 1(2):, 2022.
Xiwei Hu, Rui Wang, Yixiao Fang, Bin Fu, Pei Cheng, and Gang Yu. Ella: Equip diffusion models with Ilm for enhanced semantic alignment. arXiv preprint arXiv:2403.05135, 2024.
Yushi Huang, Zining Wang, Ruihao Gong, Jing Liu, Xinjie Zhang, Jinyang Guo, Xianglong Liu, and Jun Zhang. Harmonica: Harmonizing training and inference for better feature caching in diffusion transformer acceleration. In ICML, 2024.
Black Forest Labs. Flux, https:/ /ttp::///hthhht com/ com/blact bllckk forest labs/ 2024.
Shanchuan Lin, Anran Wang, and Xiao Yang. Sixl-lightning: Progressive adversarial diffusion distillation. arXiv preprint arXiv:2402. 19002.13929 13929, 2024.
Shanchuan Lin, Xin Xia, Yuxi Ren, Ceyuan Yang, Xuefeng Xiao, and Lu Jiang. Diffusion adversar- ial post-training for one-step video generation. arXiv preprint arXiv:2501.08316. 2025.
Tsung- Yi Lin, Michael Maire, Serge Belongie, James Hays, Pietro Perona, Deva Ramanan, Piotr Dollár, and C Lawrence Zitnick. Microsoft coco: Common objects in context. In ECCV, pp- 740-755. Springer, 2014.
Feng Liu, Shiwei Zhang, Xiaofeng Wang, Yujie Wei. Haonan Qiu, Yuzhong Zhao, Yingya Zhang, Qixiang Ye, and Fang Wan. Timestep embedding tells: It's time to cache for video diffusion model. arXiv preprint arXiv:2411.19008, 2024.
Jinming Lou, Wenyang Luo, Yufan Liu, Bing Li, Xinmiao Ding, Weiming Hu, Jiajiong Cao, Yuming Li. and Chenguang Ma. Token caching for diffusion transformer acceleration. arXiv preprint MXX1:2240911823.3 2024.
Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. Dpm-solver: fast ode solver for diffusion probabilistic model sampling in around 10 steps. NeurlPS, 35:17755775 2022a.
Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. Dpm-solvert+: Fast solver for guided sampling of diffusion probabilistic models. arXiv preprint arX1v:2211.01095, 2022b.
Simian Luo, Yiqin Tan, Longbo Huang, Jian Li, and Hang Zhao. Latent consistency models: Synthe- sizing high-resolution images with few-step inference. arXiv preprint arXiv:2310.04378. 2023.
Xinyin Ma, Gongfan Fang, Michael Bi Mi, and Xinchao Wang. Searning-to-cachee Accelerating
diffusion transformer via layer caching. NeurlPS, 37:133282-133304 2024a.
Xinyin Ma, Gongfan Fang, and Xinchao Wang. Deepcache: Accelerating diffusion models for free. In CVPR, pp. 15762-15772, 2024b.
Yue Ma, Yali Wang, Yue Wu, Ziyu Lyu, Siran Chen, Xiu Li, and Yu Qiao. Visual knowledge graph for human action reasoning in videos. In Proceedings of the 30th ACM International Conference Mullimedia, pp. 4132-4141, 2022.
12

---

Published as a conference paper at ICLR 2026
Yue Ma, Xiaodong Cun, Yingqing He, Chenyang Qi, Xintao Wang, Ying Shan, Xiu Li, and Qifeng Chen. Magicstick: Controllable video editing via control handle transformations. arXiv preprint arXiv:2312.03044, 2023.
Yue Ma, Yingqing He, Xiaodong Cun, Xintao Wang, Siran Chen, Xiu Li, and Qifeng Chen. Follow your pose: Pose-guided text-to-video generation using pose-free videos. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 38, pp. 4117-4125, 2024c.
Yue Ma, Hongyu Liu, Hongfa Wang. Heng Pan, Yingqing He, Junkun Yuan, Ailing Zeng, Chengfei Cai, Hcung- Yeung Shum, Wei Liu, et al. yourow-ouu--emojt cont-oontrol and expressive freestyle portrait animation. In SIGGRAPH Asia 2024 Conference Papers, pp. 1-12, 2024d.
Yue Ma, Kunyu Feng, Zhongyuan Hu, Xinyu Wang. Yucheng Wang, Mingzhe Zheng, Xuanhua He, Chenyang Zhu, Hongyu Liu, Yingqing He, et al. Controllable video generation: A survey, arXiv preprint arXiv:2507. 16869, 2025a.
Yue Ma, Kunyu Feng, Xinhua Zhang, Hongyu Liu, David Junhao Zhang, Jinbo Xing, Yinhan Zhang, Ayden Yang. Zeyu Wang, and Qifeng Chen. Follow-yourr-reattoon Empowering 4d creation through video inpainting. arXiv preprint arXiv:2506.0453 04590, 2025b.
Yue Ma. Yingqing He, Hongfa Wang. Andong Wang. Leqi Shen, Chenyang Qi, Jixuan Ying, Chengfei Cai, Zhifeng Li, Heung-Yeung Shum, al. image animation via motion prompts. In Proceedings of the AAAI Conference on Artificial Intel- ligence, volume 39, pp. 6018-6026, 2025c.
Open-domain regional
Yue Ma, Yulong Liu, Qiyuan Zhu, Ayden Yang. Kunyu Feng, Xinhua Zhang, Zhifeng Li, Sirui Han, Chenyang Qi, and Qifeng Chen. sillow-your-motionn Video motion transfer via efficient spatial-temporal decoupled finetuning. arXiv preprint arXiv:2506.05207, 2025d.
Yue Ma, Zexuan Yan, Hongyu Liu, Hongfa Wang, Heng Pan, Yingqing He, Junkun Yuan, Ailing Zeng, Chengfei Cai, Heung-Yer Yeung Shum, et al. follow-your-emoji-ee fine-controllable, and expressive freestyle portrait animation. arXiv preprint arXiv:2509. 2025e.
Towards efficient,
OpenAL. Video generation models as world simulators. pttps://openai.com www./////////////ccccocccc/ccece vi prideo-generation-noo s-ass as-wasooe woss-wordd--- wwwlddsimulatooeeooooooeoooooooo simulators/, 2024. Accessed: 2025-04- 10.
Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, al. Training language models to follow instructions with human feedback. NewrlPS, 35:27700-77744, 2022.
Yong-Hyun Park. Chieh-Hsin Lai, Satoshi Hayakawa, Yuhta Takida, and Yuki Mitsufuji. Jump your steps: Optimizing sampling schedule of discrete diffusion models. In ICLR, 2024.
William Peebles and Saining Xie. Scalable diffusion models with transformers. In ICCV, pp. 4195- 4205, 2023.
Ethan Perez, Florian Strub, Harm De Vries, Vincent Dumoulin, and Aaron Courville. Film: Visual reasoning with a general conditioning layer. In AAAI, volume 32, 2018.
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. In ICML, pp. 7748-8763. PmLR, 2021.
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm Ommer. High- resolution image synthesis with latent diffusion models. In CVPR, pp. 10684-10695, 2022.
Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomed- ical image segmentation. In Medical image computing and cosputer-assisted intervention- MICCAI 2015: 18th international conference, Munich, Germany, October 5-9, 2015, proceed- ings, part III 18, pp. 241. Springer, 2015.
13

---

Published as a conference paper at ICLR 2026
Axel Sauer, Dominik Lorenz, Andreas Blattmann, and Robin Rombach. Adversarial diffusion dis- tillation. In ECCV, pp. 87-103. Springer, 2024.
Christoph Schuhmann. Laion-aesthetics/ 2022. Accessed: 2025-04-19.
Laion-aesthetics.
hetps://laioon.al/blog.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv: 1707006377 2017.
Pratheba Selvaraju, Tianyu Ding, Tianyi Chen, llya Zharkov, and Luming Liang. Fora: Fast-forward caching in diffusion transformer acceleration. arXiv preprint vrXiv:2407.01425. 2024.
Huiyang Shao, Xin Xia, Yuhong Yang, Yuxi Ren, Xing Wang, and Xuefeng Xiao. Rayflow: Instance-aware diffusion acceleration via adaptive flow trajectories. arXiv preprint SxXiv:2003.0769, 2025.
Zhihong Shao, Peiyi Wang, Qihao Zhu, Runxin Xu, Junxiao Song. Xiao Bi, Haowei Zhang, Mingchuan Zhang, YK Li, Y Wu, et al. Deepseekmath: Pushing the limits of mathematical reasoning in open language models. arXiv preprint arXiv:2402.03300 2024.
Joar Skalse, Nikolaus Howe, Dmitrii Krasheninnikov, and David Krueger. Defining and character- izing reward gaming. NeurIPS, 35:9460-9471, 2022.
Jascha Sohl-Dickss Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In ICMI., pp. 2256-2265, pmir, 2015.
Jiaming Song. Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. arXiv
preprint arXiv:2010.02502, 2020.
Qwen Team. Qwq-32b:
Embracing the power of reinforcement learning, March 2025, URL btttps://wwenin hrtphnnlm..ihhub..... github.uh..... / blog/ /loog/qq---ooo 32b/.
The Movie Gen team. Movie gen: A cast of media foundation models, 2025. URL https: //aaa/ii.rrg/as//////.e 11/224111112000 13720.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. NeurlPS, 30, 2017.
Preeti Raj Verma, Navneet Pratap Singh, Deepika Pantola, and Xiaochun Cheng. Neural network developments: detailed survey from static to dynamic models. Computers and Electrical Engi- neering, 120: 109100 2024.
Bram Wallace, Meihua Dang. Rafael Rafailov, Linqi Zhou, Aaron Lou, Senthil Purushwalkam, Stefano Ermon, Caiming Xiong, Shafiq Joty, and Nikhil Naik. Diffusion model alignment using direct preference optimization. In CVPR, pp. 8228 8238, 2024.
Hongjie Wang, Difan Liu, Yan Kang, Yijun Li, Zhe Lin, Niraj K Jha, and Yuchen Liu. Attention- driven training-free efficiency enhancement of diffusion models. In CVPR, pp. 16080-16089, 2024.
Yulin Wang, Yang Yue, Yang Yue, Huanqian Wang, Haojun Jiang, Yizeng Han, Zanlin Ni, Yifan Pu, Minglei Shi, Rui Lu, et al. Emulating human-like adaptive vision for efficient and flexible machine visual perception. Nature Machine Intelligence, pp. 1-19, 2025.
WanTeam, Ang Wang, Baole Ai, Bin Wen, Chaojie Mao, Chen-Wei Xie, Di Chen, Feiwu Yu, Haiming Zhao, Jianxiao Yang. Jianyuan Zeng, Jiayu Wang, Jingfeng Zhang. Jingren Zhou, Jinkai Wang, Jixuan Chen, Kai Zhu, Kang Zhao, Keyu Yan, Lianghua Huang, Mengyang Feng, Ningyi Zhang, Pandeng Li, Pingyu Wu, Ruihang Chu, Ruili Feng, Shiwei Zhang, Siyang Sun, Tao Fang, Tianxing Wang. Tianyi Gui, Tingyu Weng, Tong Shen, Wei Lin, Wei Wang, Wei Wang, Wenmeng Zhou, Wente Wang. Wenting Shen, Wenyuan Yu, Xianzhong Shi, Xiaoming Huang, Xin Xu, Yan Kou, Yangyu Lv, Yifei Li, Yijing Liu, Yiming Wang. Yingya Zhang, Yitong Huang, Yong Li, You Wu, Yu Liu, Yulin Pan, Yun Zheng, Yuntao Hong, Yupeng Shi, Yutong Feng, Zeyinzi Jiang, Zhen Han, Zhi-Fan Wu, and Ziyu Liu. Wan: Open and advanced large-scale video generative models. arXiv preprint 2503 2503.2300033344 14, 2025.
14

---

Published as a conference paper at ICLR 2026
Xiaoshi Wu, Yiming Hao, Keqiang Sun, Yixiong Chen, Feng Zhu, Rui Zhao, and Hongsheng Li. Human preference score v2: A solid benchmark for evaluating human preferences of text-to- image synthesis. arXiv preprint arXiv:2306. 09341, 2023.
Haocheng Xi, Shuo Yang. Yilong Zhao, Chenfeng Xu, Muyang Li, Xiuyu Li, Yujun Lin, Han Cai, Jintao Zhang, Dacheng Li, al. Sparse videogen: Accelerating video diffusion transformers with 2025.
spatial-temporal sparsity. arXiv preprint arX.::2502.0011
Shitao Xiao, Yueze Wang, Junjie Zhou, Huaying Yuan, Xingrun Xing, Ruiran Yan, Chaofan Li, Shuting Wang, Tiejun Huang, and Zheng Liu. Omnigen: Unified image generation. arXiv preprint arXiv:2409. 1140, 2024.
Jiazheng Xu, Xiao Liu, Yuchen Wu, Yuxuan Tong, Qinkai Li, Ming Ding, Jie Tang, and Yuxiao Dong. Imagereward: Learning and evaluating human preferences for text-to-image generation. NeurlPS, 36:15903-1590 15935, 2023.
Hanshu Yan, Xingchao Liu, Jiachun Pan, Jun Hao Liew, Qiang Liu, and Jiashi Feng. Perflow: Piecewise rectified flow as universal plug-and-play accelerator. arXiv preprint arXiv:2405.07510, 2024.
Zhuoyi Yang, Jiayan Teng, Wendi Zheng, Ming Ding, Shiyu Huang, Jiazheng Xu, Yuanming Yang, Wenyi Hong, Xiaohan Zhang, Guanyu Feng. al. Cogvideox: Text-to-video diffusion models with an expert transformer. arXiv preprint arXiv:2408.0607,, 2024.
Zilyu Ye, Zhiyang Chen, Tiancheng Li, Zemin Huang. Weijian Luo, and Guo-Jun Qi. Schedule the fly: Diffusion time prediction for faster and better image generation. arXiv preprint arXiv:2412.01243. 2024.
Tianwei Yin, Michaël Gharbi, Taesung Park, Richard Zhang, Eli Shechtman, Fredo Durand, and Bill Freeman. Improved distribution matching distillation for fast image synthesis. NeurIPS, 37: 47455 47487, 2024.
Haoran You, Connelly Barnes, Yuqian Zhou, Yan Kang, Zhenbang Du, Wei Zhou, Lingzhi Zhang, Yotam Nitzan, Xiaoyang Liu, Zhe Lin, et al. Layer-and timestep-adaptic differentiable token compression ratios for efficient diffusion transformers. arXiv preprint arXiv:2412.16822 2024.
Huizhuo Yuan, Zixiang Chen, Kaixuan Ji, and Quanquan Gu. Self-play fine-tuning of diffusion models for text-to-image generation. arXiv preprini arXiv:2402. 2024a.
Zhihang Yuan, Hanling Zhang, Lu Pu, Xuefei Ning, Linfeng Zhang, Tianchen Zhao, Shengen Yan, Guohao Dai, and Yu Wang. Ditfastattn: Attention compression for diffusion transformer models. NeurlPS, 37:1196-1219, 2024b.
Jintao Zhang, Chendong Xiang, Haofeng Huang. Jia Wei, Haocheng Xi, Jun Zhu, and Jianfei Chen. Spargeattn: Accurate sparse attention accelerating any model inference. arXiv preprint arXiv:2502.18137, 2025a.
Peiyuan Zhang, Yongqi Chen, Runlong Su, Hangliang Ding, Ion Stoica, Zhenghong Liu, and Hao Zhang. Fast video generation with sliding tile attention. arXiv preprint 2502 2025b.
Wangbo Zhao, Jiasheng Tang, Yizeng Han, Yibing Song, Kai Wang, Gao Huang, Fan Wang, and Yang You. Dynamic tuning towards parameter and inference efficiency for vit adaptation. Ad- vances in Neural Information Processing Systems, 37:1 17::14765-111 14796, 2024.
Wangbo Zhao, Yizeng Han, Jiasheng Tang, Zhikai Li, Yibing Song, Kai Wang, Zhangyang Wang, and Yang You. A stitch in time saves nine: Small vlm is precise guidance for accelerating large vlms. In Proceedings of the Computer Vision and Pattern Recognition Conference, pp. 19814 19824, 2025a.
Wangbo Zhao, Yizeng Han, Jiasheng Tang, Kai Wang, Yibing Song, Gao Huang, Fan Wang, and Yang You. Dynamic diffusion transformer. In ICLR, 2025b.
15

---

Published as a conference paper at ICLR 2026
Wangbo Zhao, Yizeng Han, Jiasheng Tang, Kai Wang, Hao Luo, Yibing Song, Gao Huang, Fan Wang, and Yang You. Dydit++: Diffusion transformers with timestep and spatial dynamics for efficient visual generation. TPAMI, 2026.
Zangwei Zheng, Xiangyu Peng. Tianji Yang, Chenhui Shen, Shenggui Li, Hongxin Liu, Yukun Zhou, Tianyi Li, and Yang You. Open-sora: Democratizing efficient video production for all. arXiv preprint arXiv:2412.20400 2024.
SUN Zhengwentai. clip-score: CLIP Score for PyTorch. httpss//ggtthub.commtttttt c March 2023, Version 0.2.1.
Chang Zou, Xuyang Liu, Ting Liu, Siteng Huang, and Linfeng Zhang. Accelerating diffusion transformers with sosen-wise feature caching. arXiv preprint 2xX1:2241000311,22224
16

---

Published as a conference paper at ICLR 2026
ACKNOWLEDGMENT OF LLM USAGE
All writing, visualizations, and experiments are completed by the authors. LLMs (e.g., GPT-4o) used solely to refine the writing.
We organize our appendix as follows:
Additional details of methods and experiments:
Section A: Pipeline of adversarial reinforcement learning in our method.
Section B: Implementation details for compatibility with classifier-free guidance (CFG).
Section C: Details of manually combined strategics.
Section D: Sensitivity RL method.
Section E: Comparison with additional state-of-the-art techniques.
Section F: Impact of training data scale.
Section L: Experiments on DPG-Beng Bench.
Section M: Comparison with learning- based cache.
Experimental settings:
Section G: Details of training- free methods in comparison.
Section H: More implementation details of our method.
Visualizations
Section Additional visual comparison with other acceleration techniques.
Section Additional visualization results.
Section K: Visualization of the distribution patterns in dynamic acceleration strategies.
17

---

Published as a conference paper at ICLR 2026
A PIPELINE OF ADVERSARIAL REINFORCEMENT LEARNING
In Algorithm 1, we present the pipeline of the proposed adversarial reinforcement learning in RAPID*
Algorithm 1: The pipeline of adversarial reinforcement learning in RAPID*
Pre-trained diffusion transformer & and image reward model Q Randomly initialize the discriminator 'D and policy heads p = {Pscp, peache Papprse] while training do
Input:
// Training the discriminator model
If Iniggin is None then: T + 9 samplers w/o acceleration strategies from P end
If Taceele is None then: Taccele f- G samples w/ acceleration strategies from P end
for E Tunigin U T do
update D with cross entropy loss; end
// Training policies with reinforcement learning foreach training iteration do
[T,} t 9 conducts sampling with acceleration strategies from P;
(r.) +- obtain rewards of { based on Q and D;
12
13
update policy models P with JGRPO-RAPID update negative dataset Taccelo with {T,}²11
14
15
end
16
end Output: Policy heads P that can select the acceleration strategy for each image generation
B IMPLEMENTATION DETAILS FOR COMPATIBILITY WITH CLASSIFIER-FRES GUIDANCE
Our method compatible with Classifier-Free Guidance (CFG) (Ho Salimans, 2022) and actually all experiments conducted using CFG. Below, we outline the implementation details:
SD3 (Esser et al., 2024): In the default settings SD3, the CFG scale is set to 7.0. During generation with CFG, the batch size is configured to two, consisting of one sample with textual condition and another with a null condition. The acceleration is dependent on the status of the conditioned sample.
FLUX (Labs, 2024): As the CFG scale has already been distilled into the FLUX, generation can be conducted directly using our method with batch size of 1.
C COMPARISON WITH MANUALLY COMBINED STRATEGIES
To demonstrate the superiority of the proposed learned policy in our method, we compare it against manually combining different acceleration strategies, as shown in Table 6. The details of manually combined strategies are presented in Table 5. In these methods, we manually integrate reduced sampling steps, feature caching, and sparse attention Specifically, for reducing sampling steps, we adjust the sampling schedule directly. For feature caching, employ the hand-crafte adaptive method TeaCache (Liu et al., 2024), while using SpargeAttn (Zhang et al., 2025a) as the sparse mechanism.
We observe that the proposed RAPID3 significantly outperforms all manual strategies, demonstrat- ing that simply combining acceleration strategies does not lead to better performance. In fact, the joint introduction of step skipping, cache reuse, and sparse attention greatly expands the search space, making it difficult manually identify optimal strategies. This often destabilizes the gen-
18

---

Published as a conference paper at ICLR 2026
eration process and leads to unsatisfactory generation quality. This highlights the importance of learnalle policies in our method.
Table Comparison with manual acceleration strategies.
Method
Latency (s)
Step
Cache
Sparse
manual-1
3.27
21
8 0.15
(1 0.05, C2 0.06
manual-2
2.04
28
0.20
Çi 0.05, C 0.06
manual-3
2.38
26
= 0.12
(1 0.20, (2 0.21
COCO
HPS
Method
Latency (8) +
CLIP
Aesthetic f
Score
Comparison with manual acceleration strate- gies. Our method significantly outperforms them.
Table 6:
manual-1
3.27
31.43
4.92
27.16
manual-2
2.04
31.34
4.95
27.48
manual-3
2.38
29.82
4.85
26.94
RAPID
1.97
32.09
5.26
28.07
D SENSITIVITY TO RL METHOD
In addition to GRPO, we also evaluate RLOO (Ahmadian et al., 2024), a RL approach that has proven effective in LLM training. The primary difference between GRPO and RLOO lies in how the advantage is obtained. As shown in Table 7, replacing GRPO with RLOO in our method also achieves competitive performance, demonstrating the robustness of our approach across different reinforcement learning approaches.
COCO
HPS
Table 7: Replacing the GRPO with RLOO (Ahmadian et al., 2024).
RL method
CLIP
Aesthetic
Score t
RAPID demonstrates robustness across two RL approaches.
GRPO RLOO
32.09
5.26
28.07
32.10
5.27
28.06
E ADDITIONAL COMPARISON WITH SOTA METHODS
To supplement Table we incorporate additional methods into our comparison on COCO dataset (Lin et al., 2014), including ToMeSD (Bolya & Hoffman, 2023), AT-EDM (Wang et al., 2024), SDTM (Fang et al., 2025), TokenCache (Lou et al., 2024), DyDiT (Zhao et al., 2024), and ToCa (Zou al., 2024), shown in Table 8. The consistent superiority of method over these state-of-thhr of-therr- acceleration techniques further underscores its effectiveness and significance.
F IMPACT OF TRAINING DATA SCALE
In Table 9, we further investigate the impact of training data scale on the performance of our method. The results show that even with just 5K text-only training samples, our method achieves competitive performance, highlighting its data efficiency. Since approach keeps the original generator frozen to reduce training costs, the relatively modest improvement observed when increasing the training data to 40K expected. To balance performance and training efficiency, adopt 20K samples as the default setting.
G DETAILS OF TRAINING-FREE METHODS
Details of TeaCache. TeaCache (Liu et al., 2024) is representative acceleration method leverag- ing cache reuse. It manually designs strategy to accelerate inference with per-image adaptivity.
19

---

Published as a conference paper at ICLR 2026
Method
Accelcration Category
Speed
CLIP
Aesthetic 1
SD3
1.00
32.05
5.31
ToMeSD
Token Merging
1.50
30.39
5.03
AT-EDM
Token Pruning
1.54
30.27
5.02
SDTM
Token Merging
1.56 x
31.59
5,23
TokenCache
Cache-Reuse
1.49 x
31.43
5,21
DyDiT
Dynamic Neural Network
1.57 x
31.48
5.22
ToCa
Cache-Reuse
2.67
32.05
5.24
RAPID" (ours)
Dynamic Acceleration
2.92
32.09
5.26
Table 8: Comparison of methods across different acceleration categories and evaluation metries.
COCO
HPS
Method
CLIP
Aesthetic t
Score
Table 9: Impact of training data scale. The training data consists of text only. We use 20K samples as the default setting balance perfor- mance and training efficiency.
SD3
32.05
5.31
28.83
5K
32.04
5.23
27.91
20K
32.09
5.26
28.07
40K
32.09
5.27
28.26
The method employs a threshold, ò, decide whether to use the cache perform computation. Specifically, the accumulated difference between the latent maps of two consecutive timesteps exceeds the threshold 5, the model performs full computation and updates the cache. Otherwise, it directly uses the cached residual to skip the model's 's computation. Larger å brings more significant acceleration while also hurts the performance.
For experiments with SD3, we set 5 = 0.15. For FLUX, we use two settings, 5 = 0.15 and 5 = 0.25, to balance the trade-off between latency and generation quality.
Details of A-DIT. DDD (Chen et al., 2024) is representative acceleration method that lever- ages cache reuse while employing uniform strategy accelerate different image generation. The method divides the generation process into two distinct stages. In the first stage, the latter half the network layers versely, in the second stage, the earlier half of the layers can use cached residuals, and the latter layers always conduct full computations.
cached residuals, while the earlier layers perform full computations. Con-
In the experiment on SD3, for the layers that using caching, the interval for performing computations update the cache, N, is set to However, it is challenging to directly apply this approach to FLUX, which consists of two different types of layers and was not explored in the original paper. Therefore, we conduct our experiments using SD3.
Details of SpargeAttn. SpargeAttt (Zhang et al., 2025a) is a representative method for sparsifying the computation of attention. It employs thresholds Çi and Ç2 control the difference between the results of attention with and without sparsification. Larger values of G and & offer better acceleration but inevitably introduce performance degradation.
For experiments with SD3, we set (1 = 0.20 and &e 0.21. For experiments with FLUX, we use three settings: [[1 0.07, Ç2 0.08], Iç1 0.20, Ç2 0.21], and [ςi 0.30, <2 0.31].
H MORE IMPLEMENTATION DETAILS OF RAPID*
In Table 10, present the default implementation details of approach.
20

---

Published as a conference paper at ICLR 2026
Table 10: Default Implementation Details RAPID*.
details of training
device
8x NVIDIA H20 GPU
total batch size
128
learning rate
le-5
weight decay
0.1
optimizer
AdamW W
samples per group in GRPO
details of reward
decay factor À weight the discriminator reward Cj
0.97
1.0
0.95 (measured)
[41 0.07, <2 = 0.08).
candidate sparse attention
[[1 0.10, <2 = 0.11].
[[ 0.20, Ç2 0.21]
Sparse
0.05, 0.07, 0.10 (measured)
I ADDITIONAL VISUAL COMPARISON WITH OTHER ACCELERATION TECHNIQUES
In Figure 6, we provide additional visual comparisons with other acceleration techniques on FLUX (Labs, 2024).
croes country edge
woman hoiding racquet tannis
mis detare
uryealisiic
stuffed animel hokding hockey slick.
FLUX 29 stops
FLUX 9 9-steps
SpargeAlIn x1.15
TeaCache x2.61
RAPID (Ours) x2.91
FLUX 28-steps x1.0
FLUX 9-steps
RAPIDI (Ours)
x1.0
X2.08 86
x2.88
x1.15
A2.61
x2.91
Figure 6: Additional visual comparison with other acceleration techniques.
J ADDITIONAL VISUALIZATION RESULTS
In Figures/ 7 and 8, we present additional visualizations of images generated by the original FLUX (Labs, 2024) and its accelerated counterpart using our method. The results demonstrate that our method preserves visual quality more effectively while achieving the best acceleration ratio, verifying the importance of dynamically selecting acceleration strategies.
Moreover, observe several notable trends emerging from dynamic acceleration strategies:
Step-Skip Stride: The stride of step-skip is relatively small when approaches 1.0 (near the noise distribution) but becomes larger as approaches 0.0 (near the image distribution). This behavior can be attributed to the model's requirements during different stages of gen- eration. At the initial stage, the model focuses on generating the overall shape, structure, and composition of the image, which has significant impact the final image quality. Therefore, our method employs finerggraing steps during this phase to ensure higher image quality.
21

---

Published as a conference paper at ICLR 2026
Cache-Reusee Cache reuse is typically applied at intervals of 0 steps, the computations in some consecutive steps are similar and can be replaced by cached results from previous steps.
Oparse-Attention Sparse attention is applied in an image-dependent manner, meaning the frequency of its usage varies depending on the image. For instance, as shown Figure 7, 5 steps of sparse attention used for the rightmost image, while only 1 step is applied for the middle image.
These results, along with the ablation study presented in Table 3, highlight that the three acceleration policies work synergistically and effectively complement each other.
FLUX-28step
FLUX w/ our method
1.0
0.8
tmmeslap
0.6
0.4
originll stepar13
oiginal alager?
onggholl smaper giNpNn 5150000
L4i1g RO-tts srapend
0.0
step index 10
10
15
5
20
step index
step index
step index
step index
Figure 7: Additional visualizations of images generated by the original FLUX (Labs, 2024) and its accelerated version using method. (1)
K DISTRIBUTION PATTERNS IN DYNAMIC ACCELERATION STRATEGIES
In Figure we present the distribution patterns of dynamic acceleration strategies learned by our policy heads, derived from 5,000 samples of the COCO dataset (Lin et al., 2014). Specifically, Fig- ure 9()) illustrates the distribution of total steps used during generation, while Figures 9(b) and (c) demonstrate the distributions of steps involving the use of cache and sparse attention, respectively. These findings further confirms the ability of our method adaptively select acceleration strategies for each image generation process, leading to diverse range of strategies.
L EXPERIMENTS ON DPG-BENCH
We conduct experiments comprehensive benchmark. DPG-Bench (Hu et al., 2024), to evaluate the proposed method against various baselines. The results, presented in Table 11, demonstrate that the proposed dynamic acceleration strategy remains robust even in such challenging settings.
M COMPARISON WITH LFARNING-BASED CACHE
Recent advancements in feature-caching for accelerating diffusion transformers, such as Harmon- iCa (Huang al., 2024) and Leammin-tt--Cache (Ma et al., 2024a), have started incorporating
22

---

Published as a conference paper at ICLR 2026
FLUX-288h -Bstep
R.IOES
FLUX W our method
100
1.0
0.8
timestep
0.6
using
2lo using 8230m depamo
stopor11
Ling cache 3te30m3 Sqaare aqudraeer
uuing Meggura
uning oport 6Mpoor5
0.0
5 step index
10
10 step index
15
5
10
15
5
10
15
10
15
step index
step index
step index
Figure 8: Additional visualizations of images generated by the original FLUX (Labs, 2024) and its accelerated version using our method. (2)
stops
(b) distribution stops using cache
(c)
staps using sparse
3230
3000
3000
(64.8%)
3000
1781
semppee
2000
(35 (295)
2000
263
(25.
208
580
8%)
1000
(12.00%)
(5.26%)
(4.18%)
86%))
(1.02%)
88%)
(0.
(8.12)
(12.1)
(16,20)
(20,24)
24+
[0.4]
(4.8)
(8.12)
(12,16)
16+
(0.2)
[2.4]
[4,6]
(8.8)
8+
slepe steeval
steps Interval
stops Interval
Figure 9: Visualization of distribution patterns in dynamic acceleration strategies. In (a), we visualize the distribution of total steps for sample generation. In (b) and (c), we demonstrate the distribution of steps using cache and sparse attention. For instance, 3,230 samples are generated using 12-16 steps. 2,694 samples utilize cache reuse with 4-8 steps, and 2,526 samples employ sparse attention with 2-4 steps during the generation process. This experiment is based on our model, which achieves 2.91 speedup over FLUX (Labs, 2024), as illustrated in Figure 3.
Table 11: Experiments on DPG-Bench (Hu et al., 2024).
Method
Global
Entity
Attribute
Relation
Other
Average
SD3 28-steps
88.09
87.81
85.79
86.16
87.00
81.29
SD3 9-steps
85.78
85.38
82.12
83.72
83.70
76.77
A-DiT N=4
68.61
73.69
76.46
79.78
75.69
64.58
w/ TeaCache
84.03
85.56
84.15
88.46
83.20
77.53
w/ TPDM
80.00
78.35
80.26
81.76
79.19
72.46
w/ SpargeAttn
80.39
81.66
81.53
79.48
79.66
73.73
w/ RAPID* (Ours)
88.94
84.91
84,36
86.30
85.77
78.81
learning-basse based processes. These methods typically introduce routers to decide whether to perform computation for block reuse the cache from previous step. To compare our proposed method with such designs, we conduct experiments on HarmoniCa, an improved version of Learning-te to- Cache, with results presented in Table 12. The results demonstrate that method outperforms the
23

---

Published as a conference paper at ICLR 2026
Table 12: Comparison with HarmoniCa (Huang al., 2024).
COCO
HPS Score
GenEval
Method
Speed f
CLIP
Aesthetic
Correct
Overall
SD3 28-steps
1.00 x
32.05
5.31
28.83
67.81
69.01
w/ HarmoniCa
2.01 x
31.75
5.23
28.05
59.67
60.99
w/ RAPID* (Ours)
2.92 x
32.09
5.26
28.07
62.57
63.48
learning- based feature-caching approach, HarmoniCa, underscoring the superiority of our approach. This can be attributed to three key reasons:
Three-level acceleration strategies: Our method incorporates three levels of acceleration strategies, step-skip, cache-reusee and sparse-attention. rather than relying solely on cache reuse. This diverse set of strategies offers greater flexibility and optimization potential compared to caching-based methods like HarmoniCa.
Unconstrained performance upper bound: Traditional learning-based caching paradigms like HarmoniCa are designed to align the performance with the original model, which in- herently limits their upper bound to the performance of the original model. In contrast, our method, trained with reinforcement learning. is not bound by this constraint. Instead, it continuously optimizes the model achieve higher rewards, unlocking the potential for superior performance.
Its router determines whether to use caching solely based on the timestep, independent of each image generation. This results in inferior performance compared our image- adaptive accretion strategy.
We hope our design will inspire future work in feature-caching methods to fully realize their com- bined potential.
24


Reviews:
[{"review_id": "xhEfcNNVnq", "reviewer": "Reviewer_Wtnj", "summary": "This paper proposes RAPID3, a system that attaches three tiny, trainable policy heads to a frozen Diffusion Transformer (DiT) to adaptively accelerate inference per image without fine-tuning the generator. The three policies are:\n\n- Step-Skip — learns per-timestep jumps (a Beta-parameterized jump).\n    \n- Cache-Reuse — decides whether to reuse cached residuals vs recompute.\n    \n- Sparse-Attention — chooses among discrete attention sparsity patterns per timestep.\n    \n\nPolicies are trained with a group based policy-optimization (GRPO) objective and an adversarial discriminator is used alongside an image reward model to reduce reward-hacking. Experiments on Stable Diffusion 3 and FLUX show ≈3× speedups while retaining competitive metrics (CLIP, Aesthetic, GenEval/HPS) and much lower training cost than dynamic fine-tuning approaches.", "questions": "1. **Policy overhead & reproducibility:** Please report the exact inference overhead of running the three policy heads (ms per step) and memory used for caches, and confirm whether the reported latency numbers include policy compute. Also, can you release seeds/hyperparameter configs used for baselines so comparisons are fully reproducible? \n    \n2. **Ablation on discriminator strength & dataset bias:** How sensitive are results to the discriminator model choice (architecture, adapter size) and to the reward weight ω? Table 4 shows some effect, but a sweep and discussion of potential adversarial equilibria (where policies game D instead of Q) would be helpful.\n    \n3. **Generality to other samplers and guidance strengths:** Does RAPID3 interact with different samplers (DPM-solver variants) and classifier-free guidance scales? Appendix claims compatibility, but please show a small table demonstrating sensitivity to sampler/guidance variations.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. **Practical, low-cost adaptivity.**\n    \n    The paper addresses a highly practical problem (fast inference for large DiTs) and presents a solution that keeps the generator frozen, adding only lightweight heads (~0.025% of generator params) — an attractive tradeoff for deployment and proprietary models. The claimed GPU-hour savings vs fully fine-tuned dynamic models is strong evidence of real utility.   \n    \n2. **Tri-level design is well motivated and complementary.**\n    \n    Combining step skipping, cache reuse, and sparse attention covers orthogonal sources of compute (temporal, recomputation, attention cost). Ablations show incremental gains when adding strategies, supporting the architectural choice.   \n    \n3. **Careful training objective to mitigate reward hacking.**\n    \n    The use of an adversarial discriminator alongside an image-reward model is an effective and conceptually clean way to prevent policies from exploiting an automatic metric. The Algorithm and RL setup (GRPO) are reasonable for the deterministic frozen-model environment. \n    \n4. **Empirical evidence on strong backbones.**\n    \n    Results on SD3 and the larger FLUX backbone are provided, including comparisons to training-free baselines (TeaCache, ∆-DiT, SpargeAttn), single-strategy RL (TPDM), and a dynamic fine-tuned baseline (DyFLUX). Quantitative results and visualization suggest RAPID3 attains large speedups with competitive quality.   \n    \n5. **Thorough ablations and implementation detail.**\n    \n    The appendix and tables include implementation defaults, sensitivity to RL method (GRPO vs RLOO), and data-scale experiments, which strengthen reproducibility and trust in robustness.", "weaknesses": "1. **Benchmarks and baselines — potential fairness concerns.**\n    \n    The paper compares to many baselines, but details matter: how were hyperparameters for baselines tuned (equivalent computational budget / equivalent “equivalent steps”)? The paper claims equivalent-step fairness but could be clearer on whether baselines were given the same effort for tuning (especially manual adaptive methods). Without fully transparent tuning procedures and seeds, it’s hard to judge whether gains stem from method novelty or hyperparameter selection. \n    \n2. **Computational cost during deployment & overheads.**\n    \n    Although policy heads are tiny, the per-step decision logic, discriminator CPA during training, and any additional memory for caches add overhead. The latency numbers report the DiT forward pass only (excluding text encoder and VAE decoder); reporting wall-clock end-to-end costs (including policy inference) would give a clearer picture of real speedups. The paper should more explicitly quantify policy inference cost and memory overhead.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "qCJ2cU6QL4", "reviewer": "Reviewer_6idy", "summary": "This paper presents RAPID3, a reinforcement-learning-based acceleration framework for Diffusion Transformers (DiTs). The key idea is to view the diffusion sampling process as a Markov Decision Process (MDP) and attach three lightweight policy heads—Step-Skip, Cache-Reuse, and Sparse-Attention—in a clean and neat design.", "questions": "See Weakness", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The story of this paper is reasonable, and the design is clever. \n\n2. The proposed method does not modify or fine-tune the base generator. Training only the small policy heads with GRPO is resource-efficient, which makes the framework practical even for closed-source large models.", "weaknesses": "1. It is a little bit weird that the main experiments in Table 1 choose 28 steps as the baseline, which is not a commonly used setting. I want to know why.\n\n2. Given that the method is resource-efficient, more validations on video diffusion or image editing can further strengthen the impact of this paper.\n\n3. The ablation studies are solid but remain surface-level. For example, while the tri-level design (Step/Cache/Sparse) shows quantitative gains, there is little insight into *when* and *why* each policy activates.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "luJIGuylpF", "reviewer": "Reviewer_ZdkW", "summary": "This paper introduces RAPID3, a framework that enables per-image adaptive acceleration for large DiT models such as Stable Diffusion 3 and FLUX, without modifying the base generator’s parameters. RAPID3 formulates the sampling process of a DiT as a Markov Decision Process (MDP), in which each denoising step corresponds to a state transition. It attaches three lightweight policy heads, each trained via RL, to dynamically determine Step-Skip Policy, Cache-Reuse Policy, Sparse-Attention Policy. Empirically, RAPID3 achieves nearly 3× faster sampling on SD3 and FLUX with minimal image fidelity degradation.", "questions": "Lack of comparison with learned feature caching approaches\n\nThe experimental comparisons mainly include training-free baselines (e.g., TeaCache, SpargeAttn) and dynamic fine-tuned models (DyFLUX, TPDM), but they omit recent learned feature caching methods such as HarmoniCa [1] and Learning-to-Cache [2], which also learn adaptive caching or reuse strategies through training. These approaches represent a closer baseline to RAPID3, as they similarly attempt to learn data-dependent cache update schedules rather than rely on manual thresholds. Without direct comparison or discussion, it remains unclear how RAPID3’s reinforcement learning–based policy differs from or improves upon these learned caching paradigms in terms of adaptivity, efficiency, and quality preservation. Including such baselines would better position RAPID3 within the broader family of learned, model-aware acceleration methods.\n\n[1] HarmoniCa: HarmonizingTraining and Inference for Better Feature Caching in Diffusion Transformer Acceleration, ICML 2025.\n\n[2] Learning to-cache: Accelerating diffusion transformer via layer caching, NeurIPS 2024.", "limitations": "", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "RAPID3 reframes the inference-time efficiency problem of diffusion transformers as a reinforcement learning (RL) control problem, treating each denoising step as a state-action transition. This perspective is both conceptually elegant and practically impactful. The proposed tri-level policy design—encompassing Step-Skip, Cache-Reuse, and Sparse-Attention strategies—is a novel unification of previously disparate acceleration techniques under a single learnable framework.", "weaknesses": "1. Limited scope of evaluation metrics and qualitative diversity.\n\nThe paper’s evaluation primarily relies on image-level alignment metrics such as CLIP-Score, Aesthetic, and GenEval, which mainly capture prompt consistency or aesthetic preference. However, it overlooks popular benchmarks that directly assess the intrinsic quality and compositional correctness of generated images, such as DPG-Bench or other perception-grounded generation metrics. These benchmarks are designed to measure fidelity, entity consistency, and spatial relations—aspects that are particularly sensitive to dynamic acceleration strategies like cache reuse and sparse attention.\n\n2. Limited exploration of policy coordination and interdependence\n\nWhile RAPID3 introduces three distinct policy heads (Step-Skip, Cache-Reuse, Sparse-Attention), each is trained independently under a shared reward function. The paper does not fully analyze the inter-policy dependencies—for example, how aggressive step-skipping interacts with sparse attention or cache reuse. In practice, these actions may have compounding or conflicting effects on visual fidelity and stability. A more detailed study on joint versus decoupled optimization, or a hierarchical policy controller that adaptively balances the three heads, would strengthen the understanding of policy synergy and improve robustness.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)