REPA Works Until It Doesn’t: Early-Stopped, Holistic Alignment Supercharges Diffusion Training
Authors: Ziqiao Wang
Wangbo Zhao
Yuhao Zhou
Zekai Li
Zhiyuan Liang
Mingjia Shi
Xuanlei Zhao
Pengfei Zhou
Kaipeng Zhang
Zhangyang Wang
Kai Wang
Yang You
Venue: NeurIPS 2025 poster ()
Keywords: Diffusion models
Efficient training
Representation learning
Primary area: deep_learning
TL;DR: We propose HASTE, which combines holistic alignment (feature and attention) with early termination to accelerate diffusion transformer training by 28× while maintaining quality.

Abstract: Diffusion Transformers (DiTs) deliver state-of-the-art image quality, yet their training remains notoriously slow. A recent remedy---representation alignment (REPA) that matches DiT hidden features to those of a non-generative teacher (e.g., DINO)---dramatically accelerates the early epochs but plateaus or even degrades performance later.   We trace this failure to the capacity mismatch: once the generative student begins modeling the joint data distribution, the teacher's lower-dimensional embeddings and attention patterns become a straitjacket rather than a guide. We then introduce HASTE (Holistic Alignment with Stage-wise Termination for Efficient training), a two-phase schedule that keeps the help and drops the hindrance. Phase I applies a holistic alignment loss that simultaneously distills attention maps (relational priors) and feature projections (semantic anchors) from the teacher into mid-level layers of the DiT, yielding rapid convergence. Phase II then performs one-shot termination that deactivates the alignment loss, once a simple trigger such as a fixed iteration is hit, freeing the DiT to focus on denoising and exploit its generative capacity. HASTE speeds up training of diverse DiTs without architecture changes. On ImageNet 256×256, it reaches the vanilla SiT-XL/2 baseline FID in 50 epochs and matches REPA’s best FID in 500 epochs, amounting to a 28× reduction in optimization steps. HASTE also improves text-to-image DiTs on MS-COCO, proving to be a simple yet principled recipe for efficient diffusion training across various tasks.

REPA Works Until It Doesn't: Early-Stopped, Holistic
Alignment Supercharges Diffusion Training
Ziqiao Wang
Wangbo Zhao'
Yuhao Zhou'
Zekai Li' Zhiyuan Liang' Kaipeng Zhang?
Mingjia Shi'
Xuanlei Zhao'
Pengfei Zhou'
Yang You! NUS HPC-AI Lab, Shannhaii AI Laboratory, SUT Austin
Zhangyang Wang
Kai Wang't
Abstract
Diffusion Transformers (DiTs) deliver state-of-the-art image quality, yet their training remains notoriously slow. A recent remeey-representation alignment (REPA) that matches DiT hidden features to those of non-generative teacher (e.g., DINO)-dramatican
accelerates the early epochs but plateaus or even degrades performance later. We trace this failure the capacity mismatch: once the generative student begins modeling the joint data distribution, the teacher's s lower-dimensional embeddings and attention patterns become a straitjacket rather than guide. We then introduce HASTE (Holistic Alignment with Stage-wiss wise Termination for Efficient training), two-phase schedule that keeps the help and drops the hindrance. Phase / applies holistic alignment loss that simultancously distills attention maps (relational priors) and feature projections (semantic anchors) from the teacher into mid-level layers of the DiT, yielding rapid convergence. Phase Il then performs one-shot termination that deactivates the alignment loss, once simple trigger such as fixed iteration is hit, freeing the DiT to focus on denoising and exploit its generative capacity. HASTE speeds up training of diverse DiTs without architecture changes. On ImageNet 256 x 256, it reaches the vanilla SiT-XL/2 baseline FID in 50 epochs and matches REPA's best FID in 500 epochs, amounting to 28 reduction in optimization steps. HASTE also improves text -to-image DiTs on MS- MOCCO proving to be a simple yet principled recipe for efficient diffusion training across various tasks. Our code is available here.
Introduction
Diffusion Transformers (DiTs) stunningly good and stunningly slow. Recent variants such as DiT [37] and SiT [34] achieve ofate-of-the-art visual fidelity across a growing list of generative tasks [8, 29, 2, 9]. Unfortunately, their training incurs vast compute and wall-clock budgets because each update must back-propagate through hundreds of noisy denoising steps. A first wave of accelerators tackles this either by architectural surgery-linearized attention, masking or gating [56, 53, 12, 54, 2 or by training heuristics, e.g. importance re-weighting of timesteps [49]. These interventions help, but often the cost of specialized kernels or fragile aaper-parameter tuning.
Representation alignment: early rocket, late parachute? Recent work has demonstrated the effectiveness of leveraging external representations to accelerate diffusion model training-compleede sidestepping the need for architectural modifications [55, 53, 45, 30). representative method, Representation Alignment (REPA) [55], projects an intermediate DiT feature map onto the embedding
"equal contribution (ziqiaow @u.nus.edu). Ziqiao, Wangbo, Zhangyang, are core contributors. corresponding author.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

space of powerful non-generative vision encoder such as DINOv2 [36], enforcing cosine- similarity loss that bootstraps useful semantics during training. The gain immediate: the student
DiT latches onto global object structure and converges several times faster than vanilla run. Yet REPA's help is not unconditional. Figure removes the alignment loss after either 100K or 400K iterations. Stopping late (400K) improves FID over the always-on baseline; stopping early (100K) hurts-evidence that REPA works until doesn't t. Why?
Our Conjecture: Capacity mismatch incurs the hidden turning point. Diffusion models eventu- ally model the joint data distribution, a harder ob- jective than the marginal/condiitonall targets implicit in a frozen, non gonngeeert encoder. Consequently, once the student has burned in, its own capacity over- takes the teacher's. Our gradient-angle analysis (Sec- tion 2.2) shows alignment and denoising objectives start aligned (acute angles), drift to orthogonality, then turn obtuse siggalling that continued align- ment may become harmful constraint.
|IT+RERA
SIT-RERN Stop
25.0
24.3
15.0
14.1
11.3
107-od
46177
20CK
necation
Figure 1: Training SiT-XL/2 on ImageNet 256x Adding REPA slashes FID early on, but its benefit fades and ultimately re- verses; dropping the alignment loss mid- training restores progress.
Simple Remedy: Holistic alignment, then release. Two observations motivate our remedy. First, the teacher's attention maps encode relational priors that are as valuable as its embeddings [31, 20]; guiding
only features leaves this structural knowledge untapped. Second, the alignment needs a stage-wise schedule: thick guidance early, zero guidance once gradients diverge. We therefore introduce HASTE (Holistic Alignment with Stage-wise Termination for Efficient training). During Phase I we distill both projected features and mid-layer attention maps from DINOv2 into the DiT, giving the student relational and semantic shorteuts. Once simple trigger (e.g., fixed iteration) is hit, enter Phase II: the alignment loss is disabled and training proceeds with the vanilla denoising objective. The recipe is two lines of code, no kernel changes.
Contribution Summary. Our findings refine the community's 's understanding of external representa- tion guidance: it is immensely helpful early, but must be let go for the generative model to focus on specific tasks. We outline our contributions follows.
Diagnosis. We identify a capacity mismatch that flips REPA from accelerator to brake and quantifyy it via gradient-direction similarity.
Method. We propose holistic (attention feature) alignment combined with a stage-wise termination switch that deactivates alignment when it starts to impede learning.
Results. On ImageNet 256x256 our schedule matches vanilla SiT-XL/2 in 50 epochs, amounting to a 28x speed-up, and reaches REPA's best score in 500 epochs. Gains replicate on COCO text-to-image generation task.
training process
early stop point
2 Method
early stage
stage
hybrid alignment denoising
denoising
Our framework, HASTE, couples two ingredients (see Figure 2): (i) Holistic alignment: a dual- channel distillation that supervises both projected features and attention maps; (ii) Stage-wise termi- nation: a single switch that turns the alignment loss off once it ceases to help. We first recap REPA and attention alignment, then describe how we marry them and when we shut them off.
Figure 2: Overview of our framework. Phase (left) distills both feature embeddings and attention maps from frozen, generreeree teacher (DINOv2) into mid-level layers of the student DIT. When a simple trigger fires, the alignment loss is disabled; Phase II (right) then continues training with pure denoising.

---

2.1 Preliminaries
Notation. Let x be a clean image, &; its noised version at timestep t, and h, the hidden state of Diffusion Transformer Go. A frozen, non-generative vision encoder & (DINOv2) produces patch embeddings y = (x) and self-attention matrices
Representation alignment (REPA). A small MLP Js projects h, into the encoder space. REPA [55] then aligns the projected state go (bx) with by maximizing token-wise cosine similarities:
N
CREPA (OP (0, -xxcett
Σ
sim (ylml. go (h,") (n,"))
N
This regularization is jointly optimized with the original denoising objective, to guide the more efficient training of diffusion transformers.
Attention alignment (ATTA). ATTA aims to transfer attention patterns from pre-trained teacher model student model guide the latter's training process [31]. For selected layers/heads (i, we minimize token-wise cross-entropy between teacher and student attention.
2.2 Early Stop of Representation Alignment
baadient-based autopsy reveals state evolution. Figure already hinted that REPA's benefit peaks early and tapers off. To pinpoint when the auxiliary loss flips from help to hindrance, we inspect the cosine similarity
= cos VoLai, VoLREPA) E 1--1,1)
computed on the 8th block SiT-XL/2 XL/2 (the alignment depth used by REPA) over 960 ImageNet images (see details in Appendix A. A positive Pr means the teacher pushes the student in roughly same direction as denoising; negative one means two losses actively fight.
Taking < 0.1 for example, Figure shows three distinct regimes:
1. Ignition (0-200 iterations): Pt starts with a relatively high level REPA adds power; diffusion transformer profits from the teacher's guidance representation learning.
2. Plateau (200 K-400 K. iterations): P; decreases to nearly orthogonal level objectives decouple; further REPA updates neither help nor hurt.
3. Conflict (> 400 K iterations): Pt exhibits negative values gradients oppose; REPA now erases detail the student tries to learn.
orthogonal
obtuse
The cross-over coincides with the iteration where Figure shows FID curves diverging, confirming that gradient geometry is a faithful early-warning signal.
early stage
intermediate stage
stage
dr(REPA)
diridencising)
Why does conflict arise? masacity-mismatch view. Once the student starts modeling the joint data distribution, it seeks high-frequency detail absent from the teacher's embeddings. A frozen encoder trained for invariant recognition discards such minutiae by design; forcing the student back into that dwwer-dimensional manifold yields destructive gradicnts. We see the same mismatch at the level of diffusson timesteps.
Figure
Cosine similarity between REPA and de- noising gradients. Acute + orthogonal - obtuse: the auxiliary signal turns from booster to brake.
Figure 4 plots P. versus the diffusion time index. For mid-noise steps (e.g., t = 0.5) where the image still blurry, gradients align. For late steps (t < 0.1) responsible for textures and are near-orthogonal from the start. This indicates that teacher guidance intrinsically global; when the denoiser must polish pixels, the encoder has little to teach.
early stage
internediate stage
late stage
1= 0.1
= 0.5
+ 0.9
fine grain
dinREPA) across SIT timesteps
dirddnndsinng
Figure 4: Gradient similarity as function of diffu- sion timestep At 0.11 (high-detail phase) the two losses already conflict even early in training.
3

---

We sharpen this claim by feeding the teacher low-frequency only versions of each image (Figure 5). Early FID improves almost identically to vanilla REPA, proving that the speed-up stems from coarse semanlic scaffolding; high-frequency cues are less relevant to REPA's benefit.
Take-away. REPA supplies valuable global context but obstructs local detail once the stu- dent matures. Hence, alignment should be transient for further improvement.
30.0
292
Vanilla SIT
ST REPA
of-dttwww Image
Fix: Stage-wise termination. Let r denote the termination iteration around which Pr ex- hibits low similarity and the alignment pro- vides limited benefit. We then discard the auxiliary alignment loss:
15.8
13,9
13.5
[0. [0.25] 1]
[0.5. 1]
[0.75, 1]
[(0,0) - [aair + Cr.
n TT
Figure 5: Replacing teacher inputs with low-pass images leaves REPA's early gain intact: evidence that the auxiliary loss transmits mainly global structure.
(2)
n n2T, T,
where Ln may itself be the holistic combo of feature (Section 2.1) and attention (Sec-
We train SiT-1// for 200K iterations.
tion 2.3) losses. fixed works reliably well, but the gradient rule adds robustness (Appendix A.2).
image
100K
400K
4M
DINOv2
cross antropy
coslnne simliarty?
0.004
Visile
0.0015
+0 16,95
-0.001
X selected token High
8th block of SIT-XL/2
12th block
50K
100K
200K
400K
50K
100K
200K
400K
low
iteration
Iteration
(a) Visualization of attention maps from DINOv2-B (h) Attention alignment (c)
Feature alignment
and RIT-XA/ different training iterations. progress with REPA alone. progress with ATTA alone.
Figure 6: Evaluating cross-effects between feature and attention alignment. (a) Attention map visualization of selected tokens for SiT-XL/2+REPA and DINOv2-B. (b) Alignment depth we track attention map cross-entropy between the 12th-layer of DINOv2-B B-B and the 5th-layer of SiT-B/2. (c) Attention maps from 3rd-5th layers of SiT-B/2 are aligned with those from 8th, 10th, and 12th layer of DINOv2-B. Since ATTA alone does not optimize the projector, we directly compute cosine similarity between the DINOv2-B features and the 5th-layer hidden states of SiT (without projection).
2.3 Holistic Alignment by Integration Attention
Why attending to attention? Compared with token embeddings, self-attention matrices reveal where a transformer routes information each layer- its "inference pathways" in the sense of Hoang et al. [20]. These pathways encode rich relational priors: object- grouping, long-range range symmetry, and background bockrroun--forgrround segregation emerge distinct heads in DINOv2, even though the model was trained without labels. Critically, such routing information is orthogonal to the static content captured by features: models can share identical patch embeddings but attend to them in entirely different patterns, leading to divergent downstream behavior.
Rationale:
Recent evidence echoes: Li et al. [31] show that distilling only attention maps from high-capacity teacher to randomly initialized ViT is more effective than transferring only embeddings, in recover- ing the teacher's linear probe precision on ImageNet. The asymmetry suggests that attention acts as structural prior: once the model taught how to look, it can relearn what to look at rapidly.
For diffusion transformers, they must integrate global spatial cues (layout, object boundaries) across hundreds of tokens for effective representation construction. While feature alignment (REPA) accelerates the learning process by injecting semantic anchors, the structural knowledge remains underexploited. Attention alignment targets the complementary regime: transfers the global routing template to DiT, thereby enabling precise spatial and global information guidance.

---

Motivational experiments. To disentangle the respective contributions of features and attention, we probe the two signals in isolation:
(i) Feature alignment only (REPA). Figure 6a shows that REPA gradually makes SiT heads resemble those of the teacher. However, the convergence of attention patterns is slow and incomplete (see cross-entropy trend in Figure 6b).
(ii) Attention alignment only (ATTA), Aligning attention maps alone can also pull the student's hidden features toward the teacher's ' embedding space (Figure 6c) and yields a training-spest boost on par with REPA (sce details Section 3.4).
Takeaway. REPA bootstraps semantics but leaves routing under-constrained: ATTA nails routing but still requires the conditional gates to be learned from scratch. Their complementary effects motivate combining both. For chosen set S of student-teacher layer pairs (Lg, fi) and the M heads,
CATTA = MS|
Σ
Σ H """" KT softmax(Q, K (oftma)Qf:="k(mmm))
(3)
(4..
where H. is token-wist cross-entropy.
Where and when to align attention? We distill teacher heads only into intermediate student blocks (e.g., SiT-XL/2 blocks 4-7). Two empirical observations justify this selective schedule:
(i) Shallow mismatch. Early DiT layers ingest Gaussian Gaussian-noisy latents; their representations are dominated by variance normalization and channel folding rather than semantics. Supervising those layers with pixel-space attention from clean-ing clean-image -imgge encoder is therefore off- off-manifold. In practice, forcing attention on too many shallow layers destabilizes the loss and raises FID,
(ii) Deep freedom. The ultimate objective of DiT is denoising for high-quality generation, rather than representation learning. The last blocks are responsible for translating high-level structure into precise generation update. Thus, these blocks should remain dedicated the denoising objective, unregularized.
Aligning mid-layers strikes the sweet spot: they are late enough that latents carry discernible semantics, yet early enough that constraining their routing gives downstream blocks a clean, well- organized feature tensor refine.
2.4 Final Recipe: HASTE
Where we align. Attention maps from the teacher are distilled into a range of mid-depth DiT blocks; features follow the original REPA setting-one projection single mid-layer. Neither the shallow noise processing blocks nor the final denoising blocks are regularized.
What align. During Phase I (iterations weight coefficients for balancing two regularizations.
7) we apply hybrid auxiliary loss. AR and Xa
LR An CREPA + AA CATTA
(4)
When we stop. At the single switch point chohonen as a fixed iteration or the gradient-angle trigger from $2t2-both terms in (4) are dropped and training proceeds with the vanilla denoising objective.
This three-line schedule constitutes HASTE: it retains REPA's semantic anchoring, adds ATTA's routing prior, and removes all auxiliary constraints once they turn counter-productive.
Experiments
3.1
Setup
Models and datasets. Following REPA [55], we conduct experiments on three diffusion transformers:
SiT [34], DiT [37], and MM-DiT [8]. ImageNet [4] and MS-COCO 2014 [32] datasets are used for class-to-image and text-to-image generation tasks, respectively. Moreover, we employ a pre-trained DINOv2-B [36] as the representation model to extract high-quality features and attention patterns.
5

---

Implementation details. We use a training batch size of 256 and SD-V VAE [40] for latent diffusion, and set A 0.5 following REPA ensure fair comparison. Additionally, we also adopt the SDE sampler with NFEs = 250 for image generation on SiT and DiT. We set λa = 0.5 as the weight of attention alignment. We use NVIDIA A100 and H100 compute workers.
Evaluation metrics. For ImageNet experiments, we sample 50K images to assess the performance, leveraging evaluation protocols provided by ADM [5] to measure FID [16], sFID [35], [42], and Precision and Recall [27]. For text-to-imag, tome generation, we follow the settings defined in [1].
3.2 Experiments on ImageNet 256 x 256
Setting. In this experiment, we set the termination point 100K iteration (around 20 epochs) for SiT-B/2 and 250K iteration (around 50 epochs) for large and xlarge size models. while all other settings remain at their default values.
Results without classifier-free guidance. As shown Table HASTE demonstrates significant acceleration performance, con- sistently outperforming REPA on both SiT- XL and DiT-XL. This validates the superi- ority stage-wise termination and holistic alignment. Notably, SiT-XL, HASTE achieves an FID of 8.39 with only 250K iterations (50 epochs), matching the perfor- mance of vanilla SiT-XL with 1400 epochs, representing 28x acceleration. Similarly, on DiT-XL, our approach surpasses the original DiT-XL. trained with 1400 epochs, using only 80 epochs.
method
epoch FID sFID
ISt
Prec.
Rec.
Withou Classifier-fre Guidance (CFG)
MaskDiT
1600
5.69
10.34
177.9
0.74
0.60
DiT
1400
9.62
6.85
121.5
0.67
0.67
SiT
1400
8.61
6.32
131.7
0.68
0.67
DIT+REPA
170
9.60
SiT+REPA
800
5.90
5.73
157.8
0.70
0.69
FasterDiT
400
7.91
5.45
131.3
0.67
0.69
MDT
1300
6.23
5.23
143.0
0.71
0.65
HAT+
80
9.33
5.74
114.3
0.69
0.64
50
8.39
4.90
119.6
0.70
0.65
SiT+HASTE
100
5.31
4.72
148.5
0.73
0.65
Results with classifier-free guidance. We also evaluate the generation performance SAT-XL+HASTE at different epochs with classifier-free guidance (CFG) [17] apply- ing guidance interval [28]. As shown in Table 1, HASTE outperforms most of the baselines in only 400 epochs, and can achieve a comparable FID score to REPA with 500 epochs, which proves that in later training stages, the denoising objective it- self is also able to lead diffusion transform- ers to satisfactory generation capability.
With Classifier-free Guidance (CFG)
MaskDiT
1600
2.28
5.67
276.6
0.80
0.61
DIT SiT
1400
2.27
4.60
278.2
0.83
0.51
1400
2.06
4.50
270.3
0.82
0.59
FasterDiT
400
2.03
4.63
264.0
0.81
0.60
MDT
1300
1.79
4.57
283.0
0.81
0.61
DiT+TREAD
740
1.69
4.73
292.7
0.81
0.63
MDTv2
1080
1.58
4.52
314.7
0.79
0.65
SiT+REPA
800
1.42
4.70
305.7
0.80
0.65
100
1.74
4.74
268.7
0.80
0.62
400
1.44
4.55
293.4
0.80
0.64
SIT+HASTE
Qualitatively comparison. We also pro- vide representative visualization results from SiT-XL/2 with REPA and HASTE in Figure 8, respectively. Our method achieves better semantic information and detail generation at early training stages.
500
1.42
4.49
299.5
0.80
0.65
600
1.41
4.51
296.9
0.80
0.65
Table 1: System-level comparison ImageNet 256 x 256. and + denote higher and lower values are better, respectively. Bold font denotes the best performance.
3.3 Text-to-Image Generation Experiment
Setting. To validate our approach in text-to-image generation tasks, we apply HASTE to MM-DIT [8], widely used architecture, and train it on the MS-COCO 2014 dataset [32] following REPA. In practice, we set termination point 7 = 200K for HASTE. Morcover, we only perform attention alignment with the QK" matrix generated from input image to avoid affecting the textual process.
Quantitative results. In Table 2, we compare our method with the original MM-DIT and MM- DIT+REPA using SDE sampler with NFEs 250. Results reflect that HASTE consistently out-

---

performs its counterparts and alignment termination leads to better performance, validating the generalizability of our holistic alignment and termination strategy in text-to-image generation tasks.
3.4 Ablation Studies
In this section, we conduct extensive experiments and comparisons across different SiT models on ImageNet 256,, to further support our analysis and claims in Section 2. We consistently use the SDE Euler-Maruyama sampler (NFEs 250) without classifier-free guidance.
Effectiveness of ATTA and termination. To validate the effectiveness of termination and Attention Alignment, we evaluate the performance of SiT-XL/2 with different methods applied before and after the termination point (50 epoch) and present the results in Table 3. Firstly, at both 40 and 100 epochs. we observe that using only At- tention Alignment can also obtain similar accelcration REPA. Moreover, the holistic alignment leads better performance at 40 epoch, which is consistent with our hypothesis in Section 2.3 that the two methods have com- plementary potentials.
model
iteration
term.
FID
MM-DiT
150K
5.26
+REPA
150K
4.16
+REPA
250K
4.28
+HASTE
150K
4.09
+HASTE
250K
4.10
+HASTE
250K
200K
4.06
Table FID results text-to-image generation on MS-COCO. HASTE con- sistently outperforms REPA to acceler- ate the training of MM-DiT.
However, the acceleration of such integration gets inferior to REPA alone at 100 epoch. We assume that consistently applying holistic aligmment leads reguaauulaii in later training stages. And the performance gets improved eventually with the termination strategy applied at 50 epoch.
Different termination iterations r. In this section, we analyze the impact of T across varying model sizes. First, conduct experiments in Table to further explore the effect of termination. The results reflect that stage-wise termination also leads to better generation quality on SiT-B/2 and SiT-L/2. For SiT-XL/2, interestingly, while 400 demonstrates lower FID 400K iteration, 7 250 model ultimately delivers superior performance when evaluated at 500K iteration.
As shown in Table 3 and Table 4, although holistic alignment achieves better performance 400K iteration, consistently regularizing the model leads to reduced performance. While termination at 400 K alleviates such trend, its performance at 500K iteration is still inferior that of 250 K. Therefore, we hypothesize that the acceleration effect gradually diminishes before 400K iteration, and the stage-wise termination, such at 250 K, can help alleviate the over-regularizattoo.
epoch REPA ATTA term.
FIDJ
sFID
1St
model
iteration
FID.,
sFID.
ISt
X
X
24.3
5.08
56.1
SiT-B/2
21.3
6.80
69.9
40
400K
X
10.7
5.02
103.9
+HASTE
100K
19.6
6.38
73.0
13.6
5.02
89.7
SiT-I/2
8.9
5.18
119.0
40
400K
O
5.04
108.8
+HASTE
250K
7.9
124.8
X
X
X
14.8
5.18
84.9
SiT-XL/2
5.5
4.74
144.4
100
400K
X
7.5
5.11
130.1
+HASTE
250K
7.3
5.05
128.7
8.5
5.00
120.7
8.1
5.20
126.1
SiT-XL/2
100
8.1
5.20
126.1
S00K
250K
5.3
4.72
148.5
+HASTE
4.72
148.5
400K
7.4
5.10
128.8
Table 3: Comparison of different methods applied to SiT-XL/2. and x denote methods applied or not, respectively. Results reflect that our termina- tion and holistic alignment strategies are effective.
Table 4: Comparison applying termination or not across different model sizes of SiT. 7 denotes termination point. find the termination strat- egy contributes to better performance eventually.
Taking SiT-XL/2 for example, we carefully assess the effect of different 7. We observe performance progresses slowly after 250K iteration (see Figure 7a). And the gradient cosine similarity between holistic alignment and denoising has shown negative values at late diffusion timesteps (see details in
7

---

Appendix A) Consequently, we consider termination near this threshold: results at 400K iteration Figure indicate that early stopping at 250K yields better performance.
training laration
100K
200K
400K
SIT-Rt STT-LLG GA Sir.a SIT-R PHRREAA
SITSAL 2A REEEN
SIT-KUP-HASTE
Figure 8: HASTE improves visual scaling. We compare images generated by SiT-XL/2+REPA and SiT-XL/2+HASTE (ours) at different training iterations. For both models, we use the same seed, noise, and sampling method with classifier- guidance scale of 4.0.
Different Attention Alignment loss weight A We evaluate the sensitivity of model to the at- tention alignment loss weight A in Equation with SiT-L/2 as an example.
-IIFID
-0- FID
7.5
74
As shown in Table 5, HASTE consistently im-
73
proves the performance of SiT-L/2 400K it- eration across different values of Ax, indicating that attention alignment provides relatively sta- ble benefits. We note that larger weights can lead to reduced performance. Therefore, we choose A 0.5 in our primary experiments.
250K 200K 270K 200K 280
220K 260K 270K 23CK 200K carly stop paint (7)
baration
(a) FID and sFID around 270K iter. SiT-XL/2 without termination.
(b) FID at 400K iter. with termination point around 270K.
Selection alignment layers. We try differ- ent transfer layers for HASTE SiT-L/2 in
Figure 7: Comparison of different termination point r on SiT-XI//2. We observe the training os- cillation after 250K iteration. Using 250 K leads to better performance 400K iteration.
Table For brevity, we denote layers from SiT and DINO as layer-S and layer-D, respectively. Additionally, we use [:ss and [-] to specify particular layer indices (counting from 0).
Firstly, we find that enough deeper layers should get involved for optimal performance. As shown in Table when choosing only two layers of each model for alignment, namely [10, 11]p and 7]s. the performance is inferior to choosing four layers. Additionally, results reflect that enough shallow layers should be left for processing the noisy inputs the latent space. We can observe that the distillation of [8, 9, 10, 1 to [4, 5, 6, 7]s achieves a better FID without including [6, 7 to [2, 3]s
model
Aa FIDJ
sFIDy
ISt
model
layer-D
layer-S
FIDJ
sFIDJ
ISt
0.5
7.9
5.08
124.8
[10, 11]p
7],
8.9
5.31
119.3
SiT-L/2
SIT-L/2
8.6
120.0
[8,9, 10, 11]p
5,
7]s
7.9
5.08
124,8
+HASTE
+HASTE
1.5
8.7
5.23
119.3
6,7,8,9 10, 11]p
[2,3,4,5,6.
7]s
8.3
5.12
121.3
3.0
9.0
5.34
116.9
Table Comparison HASTE with different choices of layers on SiT-L/2 at 400K iteration. While transferring attention maps for more deep layers provides greater benefits, we need to pre- serve enough shallow layers to process latent input.
Table 5: ATTA weight AA = 0.5
leads to better performance on SiT- L/2 at 400K iteration.

---

Our findings align with the observations of attention transfer on ViTs reported in [31]: transferring more attention maps from deeper layers provides greater benefits, and ViTs can learn low-level features well when guided on how to integrate these features into higher-level ones.
Related Work
4.1 Accelerating Training Diffusion Transformers
To accelerate the training of diffusion transformers, existing methods can be broadly classified into two categories: architectural modifications and representation enhancements.
Architecture modification. These methods focus on directly improving the efficiency of the model architecture. For example, SANA series [51, 52], DiG [58], and LiT [47] introduce Linear Attention [22, 50, to improve the efficiency of diffusion transformers. Additionally, methods like MaskDiT [56] and MDT [10, 11] introduce masked image modeling [14] to reduce the cost during training.
Representation incorporation. In contrast to architecture modifications, these methods do not require designing specialized structures and instead leverage external representations to achieve acceleration. For instance, REPA [55] observes the difficulty learning effective representations for diffusion models [43, 18, 44]. which hinders the training efficiency. To address this, REPA proposes to align the internal features of diffusion transformers with the output of pre-trained representation models, and significantly accelerates the training process.
Furthermore, recent works [53, 45, 30] have also achieved better results based on representation
methods. For example, U-REPA [45] improves REPA with a manifold alignment loss, and demon- strates its effectiveness on U-Nets [41]. External representations can also help enhance generation
and reconstruction capabilities of VAE, such as in VA-VAS [53] and E2E- VAE [30].
Unlike these methods, research focuses mainly on the diffusion transformer itself. We investigate the relationship between external representation guidance and the seff-improvement of diffusion transformers, and propose to remove the regularization at appropriate training stage.
4.2 Attention Transfer for Vision Transformers
The attention mechanism [46] has been shown to provide vision models, such as Vision Transformers (ViTs) [6], with strong adaptability and scalability across various tasks. While prior works [15, 13] have achieved improved downstream performance by leveraging entire pre-trained models, Li et al. [31] demonstrates that the attention patterns learned during pre-training sufficient for ViTs to learn high-quality representations from scratch, achieving performance comparable to fine-tuned models on downstream tasks. Consequently. attention distillation [31, 48, 39] has been proposed to transfer knowledge efficiently.
The transfer of attention maps has been extensively studied in Vision Transformers (ViTs), but remains underexplored in diffusion transformers. While recent work [57] applies attention distillation for characteristics transfer tasks using diffusion models, its explorations remain in the sampling process. Moreover, the relationship between attention mechanisms in ViTs and diffusion transformers requires further investigation. In this work, demonstrate that attention maps from pre-trained ViT can effectively guide the learning process of diffusion transformers.
Conclusion
In this paper, have proposed HASTE, simple but effective way to improve the training efficiency of diffusion transformers. Specifically, we identify the capacity mismatch and reveal that represen- tation alignment is not always beneficial throughout the training process. In addition, we analyze the stages when feature alignment is most effective and investigate the dilemma between external feature guidance and internal self-improvement of diffusion transformers. We prove that HASTE can significantly accelerate the training process of mainstream diffusion transformers, such as SiT and DIT. We hope our work will further reduce the cost for researchers to train diffusion transformers, and broaden the application of diffusion models in downstream tasks.

---

Limitations and future work. We mainly focus on diffusion transformers in latent space for image generation. Explorations of HASTE with pixel-level diffusion [5, 24]. or in video generation tasks [19] would be exciting directions for future work. Additionally, HASTE may also be incorporated with other methods [53, 30] on different model architectures [45].
Faster generative model training could potentially lower barriers for harmful content generation. As a technical contribution focused on training efficiency, our work does not directly address broader AI safety challenges, which we acknowledge as important future work.
It is meaningful to further explore attention alignment over cross cross-attentes attention maps between MM-DiT and pre-traine pre-trained modti-modal encoders. The development of more robust and efficient automatic termination triggers possibly guided by gradient-basee based signals remains open direction.
Acknowledgement. We sincerely appreciate Liang Zheng and Ziheng Qin for valuable discussions and feedback during this work. We acknowledge anonymous reviewers for insightful questions and discussions. Our work is sponsored by NUS startup grant (Presidential Young Professorship), Singapore MOE Tier-1 grant, BytelDnce grant, NUS ARTIC grant, Apple grant, Alibaba grant, Google Research and Google grant for TPU usage.
References
[1] Fan Bao, Shen Nie, Kaiwen Xue, Yue Cao, Chongxuan Li, Hang Su, and Jun Zhu. All are worth words: A vit backbone for diffusion models. In CVPR, 2023.
[2] Tim Brooks, Bill Peebles, Connor Holmes, Will DePue, Yufei Guo, Li Jing, David Schnurr, Joe Taylor, Troy Luhman, Eric Luhman, Clarence Ng, Ricky Wang, and Aditya Ramesh. Video generation models as world simulators. https://openai.com/eett www..reeeeeeee//iiieeeeeo weweeggeieeeenneeeeeeeee models- www--wordsssllattee 2024.
[3] Krzysztof Marcin Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreca Gane, Tamas Sarlos, Peter Hawkins, Jared Quincy Davis, Afroz Mohiuddin, Lukasz Kaiser, David Benjamin Belanger, Lucy J Colwell, and Adrian Weller. Rethinking attention with performers. In ICLR, 2021.
[4] Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In CVPR, 2009.
Prafulla Dhariwal and Alex Nichol. Diffusion models beat gans on image synthesis. In NeurlPS, 2021.
[6] Alexey Dosovitskiy. Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer. Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. In ICLR, 2021.
[7] Stefan Elfwing, Eiji Uchibe, and Kenji Doya. Sigmoid-weighte linear units for neural network function approximation in reinforcement learning. Neural Networks, 107:3-11, 2018. ISSN 0893-6080. doi: http:/// htw..r//10.1011/...eeee....2
URL htt.nttt//.w..ecieneeeeett sciencedirect........... wwcm//cciene/arttiiiiieeeeeeeeee article/ pi1/80993608017302976. Special issue on deep reinforcement learning.
[8] 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.
[9] Kunyu Feng, Yue Ma. Bingyuan Wang, Chenyang Qi, Haozhe Chen, Qifeng Chen, and Zeyu Wang. Dit4edit: Diffusion transformer for image editing. In AAAI, 2025.
[10] Shanghua Gao, Pan Zhou, Mingg- Mingg-Ming Cheng, and Shuicheng Yan. Masked diffusion transformer is strong image synthesizer. ICCV, 2023.
[11] Shanghua Gao, Pan Zhou, Ming-Ming Cheng, and Shuicheng Yan. Mdtv2: Masked diffusion transformer is a strong image synthesizer. arXiv preprint arXiv:2303. 143889 2024. URL https://arxxi http::///awvvrveaeeeeee ortg 2303.14388. 14389.
[12] Tiankai Hang. Shuyang Gu, Chen Li, Jianmin Bao, Dong Chen, Han Hu, Xin Geng, and Baining Guo. Efficient diffusion training via min-snr weighting strategy. In ICCV, 2023.
[13] Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick. Momentum contrast for unsupervised visual representation learning. In CVPR, 2020.
10

---

[14] Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross Girshick. Masked autoencoders scalable vision learners. In CVPR, 2022.
[15] Kaiming Hc, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross Girshick. Masked autoencoders are scalable vision learners. In CVPR, 2022.
[16] Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by a two time-scale update rule converge to local nash equilibrium. NeurlPS, 2017.
[17] Jonathan Ho and Tim Salimans. Classifier-fiee diffusion guidance. In NewrIPS Workshop, 2021.
[18] Jonathan Ho, Ajay Jain, Pieter Abbeel. Denoising diffusion probabilistic models. In NeurIPS, 2020.
[19] Jonathan Ho, Tim Salimans, Alexey Gritsenko, William Chan, Mohammad Norouzi, and David J Fleet. Video diffusion models. In NewrIPS, 2022.
[20] Scott Hoang, Minsik Cho, Thomas Merth, Atlas Wang. Mohammad Rastegari, and Devang Naik. compressed Ilms forget knowledge? an experimental study with practical implications. In NeurIPS Workshop, 2024.
[21] Tero Karras, Miika Aittala, Jaakko Lehtinen, Janne Hellsten, Timo Aila, and Samuli Laine. Analyzing and improving the training dynamics of diffusion models. In CVPR, 2024.
[22] Angelos Katharopoulos, Apoory Vyas, Nikolaos Pappas, and François Fleuret. Transformers are mns: fast autoregressive transformers with linear attention. In ICMI., 2020.
[23] Dongjun Kim. Seungjae Shin, Kyungwoo Song, Wanmo Kang, and I1-Chul Moon. Soft truncation: A universal training technique of score-based diffuusion model for high precision score estimation. In ICML. 2022.
[24] Diederik Kingma and Ruiqi Gao. augmentation. In NeurIPS, 2023.
Understanding diffusion objectives as the elbo with simple data
[25] Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In ICLR, 2015.
[26] Felix Krause, Timy Phan, Vincent Tao Hu, and Bjöm Ommer. Tread: Token routing for efficient architectur-agnostic diffusion training. arXiv preprint arXiv:2501.0470 2025.
[27] Tuomas Kynkäänniemi, Tero Karras, Samuli Laine, Jaakko Lehtinen. Timo Aila. Improved precision and recall metric for assessing generative models. NeurlPS, 2019.
[28] Tuomas Kynkäänniemi, Miika Aittala, Tero Karras, Samuli Laine, Timo Aila, and Jaakko Lehtinen. Applying guidance in limited interval improves sample and distribution quality in diffusion models. In NeurIPS, 2024.
[29] Black Forest Labs. Flux. https:///giee www/b//cc/////o/ooo//////....oaa forest- -awesst-lab///eb labslllblllfe fabs/flux, 2024.
[30] Xingjian Leng, Jaskirat Singh, Yunzhong Hou, Zhenehang Xing. Saining Xie, and Liang Zheng. Repa-e: Unlocking vae for end-to-end tuning with latent diffusion transformers. arXiv preprint arXiv:2504. 10483, 2025. URL https: /atee:///rreee /ttp///xvvvvor/////////e/.. .rg// abs abs/2504 2504. 10483.
[31] Alexander Cong Li, Yuandong Tian, Beidi Chen, Deepak Pathak, and Xinlei Chen. On the surprising effectiveness of attention transfer for vision transformers. NeurIPS. 2024.
[32] 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, 2014.
[33] Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. /CLR, 2017.
[34] Nanye Ma, Mark Goldstein, Michael S. Albergo, Nicholas M. Boffi, Eric Vanden-Eijnden, and Saining Xie. Sit: Exploring flow and diffusion-based generative models with scalable interpolant transformers. arXiv preprint arXiv:2401.08744 2024. URL https: //ar:i/////ii/////a/ www.xi.oog/aaae arg/abs/2444 08740.
[35] Charlie Nash, Jacob Menick, Sander Dieleman, and Peter Battaglia. Generating images with sparse representations. In ICML, 2021.

---

[36] Maxime Oquab, Timothée Darcet, Théo Moutakanni, Huy V. Vo. Mare Szafraniec, Vasil Khalidov, Pierre Fernandez, Daniel HAZIZA, Francisco Massa, Alaaeldin El-Nouby, Mido Assran, Nicolas Ballas, Wojciech Galuba, Russell Howes, YaoYao Huang, Shang- Wen Li, Ishan Misra, Michael Rabbat, Vasu Sharma, Gabriel Synnaeve, Hu Xu, Herve Jegou, Julien Mairal, Patrick Labatut, Armand Joulin, and Piotr Bojanowski. DINOv2: Learning robust visual features without supervision. TMLR, 2024, ISSN 2835-8856. URL https: / /ttp//ppeeeeeeeeeeeeeeeeeeeeeeee ew net/forum? ? F6.. Featured Certification.
[37] William Peebles and Saining Xie. Scalable diffusion models with transformers. In ICCV, 2023.
[38] Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, Gretchen Krueger, and llya Sutskever. Learning transferable visual models from natural language supervision. In ICML, 2021.
[39] Sucheng Ren, Fangyun Wei, Zheng Zhang, and Han Hu. Tinymim: An empirical study of distilling mim pre-trained models. In CVPR, 2023.
[40] Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm Ommer. High-resolution image synthesis with latent diffusion models. CVPR, 2022.
[41] Olaf Ronneberger. Philipp Fischer, and Thomas Brox, U-net: Convolutional networks for biomedical image segmentation. In MICCAI, 2015.
[42] Tim Salimans, Ian Goodfellow, Wojciech Zaremba, Vicki Cheung, Alec Radford, and Xi Chea. Improved techniques for training gans. In NeurlPS, 2016.
[43] Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In ICML, 2015.
[44] Yang Song, Jascha Sohl-Dickstein, Diederik Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In ICIR, 2021.
[45] Yuchuan Tian, Hanting Chen, Mengyu Zheng, Yuchen Liang, Chao Xu, and Yunhe Wang. U-repa: Aligning diffusion u-nets to vits arXiv preprint arXiv:2503. 2025. URL atttps://arxiy org/ aaaaab////00 /2003. 18414.
[46] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, L ukasz Kaiser, and Illia Polosukhin. Attention all you need. In NeurlPS, 2017.
[47] Jiahao Wang, Ning Kang, Lewei Yao, Mengzhao Chen, Chengyue Wu, Songyang Zhang, Shuchen Xue, Yong Liu, Taiqiang Wu, Xihui Liu, Kaipeng Zhang, Shifeng Zhang. Wenqi Shao, Zhenguo Li, and Ping Luo. Lit: Delving into a simplified linear diffusion transformer for image generation. arXiv preprint arXiv:2501. 2025, URL https: / arxiv. org/ ats///bs///// /201. 12976.
[48] Kai Wang, Fei Yang 0004, and Joost van de Weijer 0001. Attention distillation: supersuue vision transformer students need more guidance. BMVC, 2022.
[49] Kai Wang, Mingjia Shi, Yukun Zhou, Zekai Li, Zhihang Yuan, Yuzhang Shang, Xiaojiang Peng. Hanwang Zhang, and Yang You. A closer look at time steps is worthy of triple speed-up for diffusion model training. arXiv preprint 2rXiv 17403, 2024.
[50] Sinong Wang, Belinda Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprint 2tte:20006007668. 2020.
[51] Enze Xic, Junsong Chen, Junyu Chen, Han Cai, Haotian Tang, Yujun Lin, Zhekai Zhang, Muyang Li, Ligeng Zhu, Yao Lu, and Song Han. Sana: Efficient high-resolution image synthesis with linear diffusion transformer. arXiv preprint arXiv:2410.
2024. URL https: / //t////axvvvaa//////ee//aaaaa.aa org/abs//////oooooooooo..o aa///////24/0 241../111............1111
[52] Enze Xie, Junsong Chen, Yuyang Zhao, Jincheng Yu, Ligeng Zhu, Yujun Lin, Zhekai Zhang, Muyang Li, Junyu Chen, Han Cai, Bingchen I. Liu, Daquan Zhou, and Song Han. Sana 1.5: Efficient scaling of training-time and inference-time compute in linear diffusion transformer. arXiv preprint a5Xiv: 18427, 2025. URL https: / /ttp///axv/vv/////////////aa. org/ abs////s2/50 18427.
[53] Jingfeng Yao and Xinggang Wang. Reconstruction vs. generation: Taming optimization dilemma in latent diffusion models. In CVPR, 2025.
[54] Jingfeng Yao, Cheng Wang, Wenyu Liu, and Xinggang Wang. Fasterdit: Towards faster diffusion transformers training without architecture modification. In NewrlPS, 2024.
12

---

[55] Sihyun Yu, Sangkyung Kwak, Huiwon Jang. Jongheon Jeong, Jonathan Huang. Jinwoo Shin, and Saining Xie. Representation alignment for generation: Training diffusion transformers easier than you think. In ICLR, 2025.
[56] Hongkai Zheng, Weili Nie, Arash Vahdat, and Anima Anandkumar. Fast training of diffusion models with masked transformers. TMLR, 2024, ISSN 2835-8856. URL httpp::/ wtwp://rrriie..ee//rruuuuaaaaaaa www.:neeefofmmmmme vTBjBtGioE.
[57] Yang Zhou, Xu Gao, Zichong Chen, and Hui Huang. Attention distillation: unitied approach visual characteristics transfer. arXiv preprint arXiv:2502.22233 2025. URL atttp:://araee org/ abs/2502. 20235.
[58] Lianghui Zhu, Zilong Huang, Bencheng Liao, Jun Hao Liew, Hanshu Yan, Jiashi Feng. and Xinggang Wang. Dig: Scalable and efficient diffusion models with gated lincar attention. arXiv preprint arXiv:2405. 18428,
2024. URL https: / /arxivv..vv/v wwgg/bs//////e/eooee/...o... abs/2405 18428.
13

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The abstract and introduction have clearly stated our contributions, including the motivation and our acceleration strategy, which are consistent with the scope presented in the main body of the paper.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
See Conclusion.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (c.g.. independence assumptions, noiseless settings, model sell-specfiiatiinn asymptotic approximations only holding locally). ). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex, system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
14

---

Answer: [NA]
Justification:
The paper mainly focuses on empirical studies.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: See Method and Experiments.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
specific model and empirical evaluation, may
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily new algorithm, the paper should make it clear how to reproduce that algorithm. (b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with open-source dataset instructions for how construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
15

---

Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [No]
Justification: ImageNet and MS-COCO 2014 are public datasets. We will release our code after acceptance.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.. cc/
for more details.
While we encourage the release of code and data, we understand that this might be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https: / httppsscccppbccccccceecce geides/ Codeeubbm ssi aasonnooicyy for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments should state which ones are omitted from the script and why.
reproducible, they
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: See Experiments.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary appreciate the results and make sense of them.
The full details can provided either with the code, in appendix, or as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [No]
Justification: In practice, we find error bars of evaluation with 40K or 50K samples are small and omitted for simplification.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
16

---

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
should be clear whether the error bar is the standard deviation or the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables figures symmetric error bars that would yield results that are out range (c.g. negative emror rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification: See Experiments.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that NeurIPS Code of Ethics https: / /nwuriir..p.. oc cc/public/ wtw.tthiiiiiiillleeeeeeeeeeeeeee
didn't 't make it into the paper).
9. Code ethics
Question: Does the research conducted in the paper conform, in every respect, with the
Answer: [Yes]
Justification: Yes, our research conforms in every respect to the NeurIPS Code of Ethics.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: There is no socictal impact of the work performed.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
17

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility MI..
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: The paper poses no such risks.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Lieenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: We have cited all related papers.
Guidelines:
The answer NA means that the paper does existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com com/dataseta has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
18

---

If this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: The paper does not release new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/mmoel as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question:
Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
16. Declaration of LLM usage
19

---

Question: Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods this research? Note that if the LLM used only for writing, editing, or formatting purposes and does not impact the core methodology. scientific rigorousness, or originality the research, declaration is not required.
Answer: [NA]
Justification: The core method development in this research does not involve LLMs any important, original, or non-standard components.
Guidelines:
The answer NA
that the method development in this research does involve LLMs as any important, original, or non-standard components.
Please refer to our LLM policy (https : h/tp://ies...//eeeeeeeeee http///rrs..c//ooeeeeeeeeeeeeeee / for what should or should be described.
20

---

A Additional Results
A.1 Gradient Angle Diagnosis
We provide detailed results of cosine similarity between REPA [55] and denoising gradients. In Figure 9, we separately compute gradients of the feature alignment and the denoising objective for SiT-XL/2 [34] and compare the cosine similarity of their directions at different training iterations. Specifically, we randomly sample 960 images from the training dataset of ImageNet [4] for the comparison and take gradients of parameters in the eighth block of SiT-XL/2 for example (REPA sets the default alignment depth as 8).
iteration I 0.02
0.04 t= 0.06 t= 0.08
0.10
cosine similarity
0.25
170.1
100K
0.0070
0.0064
0.0327
0.0525
0.0692
tm0.5
190.9
200K
0.0350
0.0476 0.0434
0.0568 0.0628
iteraiio
300K
-0.0235
-0.0324 -0.0316
-0.0044
-0.0116
100K
200K
400K
4M
400K
-0.1236
-0. -0.1056 -0.1133
0.0232 0.0130
500K
0.0346
-0.0368 -0.0246 -0.0063 -0.0409
00.27
-0.29
600K
-0.1185
-0.0546 0.0645 -0.0039 0.0372
4.40
4M
-0,2065
-0.1279 -0,1928 -0,3621 -0.4942
Figure 9: Gradient cosine similarity between REPA and the denoising objective.
Table 7: Detailed cosine similarity results of the 8th block in SIT-X1/22 at / S 0.10.
We first observe relatively high cosine similarity, representing an acute angle between gradients of the two objectives. However, the similarity shows a decreasing trend as the training progresses, and the angle becomes nearly orthogonal at the intermediate stage (around 400K iteration). Furthermore, we find that the similarity becomes obviously negative at the final training stage, such as at 4M iteration, indicating that there might be some potential conflict between REPA and diffusion loss.
In addition to training iterations, we also find a feature alignment gap over different diffusion timesteps: As reported in [55]. well-trained DiT [37] or SiT exhibits a higher feature alignment at the intermediate diffusion timesteps, while the alignment is notably weaker at those closer to the data distribution, i.e., nearby the sampling results, such as t 0.1 for SiT. We observe a similar trend in our gradient similarity comparison. According to diffusion sampling properties, the initial steps starting from noise mainly contribute global fidelity, namely the basic outline of images, while the steps closer to the data are to refine microscopic details such as textures [23]. We hypothesize that the diffusion transformer eventually needs to refine its own representations for detail generation beyond learning directly from external features.
iteration
t= 0.02
t= 0.05
t= 0.07
t= 0.1
t= 0.2
t= 0.5
t= 0.9
100K
-0.0138
-0.0131
-0.0688
0.0129
0.0488
0.0541
-0,0093
200K
-0.0423
-0.0674
00.7199
-0.0491
0.0068
0.0801
0.0099
250K
-0.0323
-0.0597
-0.0598
-0.0599
-0.0264
0.0354
0.0419
260K
-0.0232
-0.0331
-0.0243
-0.0034
0.0436
0.0729
0.0065
270K
0.0029
0.0152
0.0113
0.0097
0.0419
0.0554
0.0233
280K
-0.0263
-0.0131
-0.0031
0.0011
0.0217
0.0455
-0.0176
290K
-0.0524
0.0199
0.0308
0.0532
0.0832
0.0550
0.0111
Table Detailed gradient cosine similarity results between holistic alignment and denoising objec- tives on the 8th block of SiT-XL/2 at different training iterations.
For our method, HASTE, also examine the gradient cosine similarity between holistic alignment and denoising. The similarity trend serves as a kind of reference for our termination strategy.
21

---

A.2 Gradient Angle Trigger
In Section 2.2, we introduce the gradient- -aniee rule for robustness and theoretical insight of HASTE. In this section, we further explain the role of gradient signals, and then provide a plain example of HASTE with automatic gradient- angle trigger.
Hybrid termination strategy. In practice, we find setting a fixed termination iteration T for HASTE actually performs reliably well in ImageNet and AS-COCO experiments. Gradient threshold serves as supplemental alternative, adding adaptability to different situations. Moreover, gradient analysis helps us to locate the termination point, As we state Section 3.4, we choose to stop around 250K iteration when we observe limited alignment benefits and relatively low gradient cosine similarities.
Gradient-angle trigger (plain implementation). We train SiT-B/2 on ImageNet 256 x 256 following the original settings of HASTE. We start evaluating the gradient cosine similarity between holistic alignment and denoising objectives every 10K iterations after 50K iteration, when the training process has been relatively stable. At each evaluation, we use selected timesteps & at various noise levels with 2048 images randomly sampled from ImageNet, and average the similarities. When the average similarity falls below threshold 6, we terminate the holistic alignment. In the plain example, we perform gradient evaluation on the 5th SiT block (counting from 1). We select 0.02, = 0.1, = 0.5 and = 0.8, and set 8 -005 As shown in Tables 9 and 10, our automatic trigger terminates the holistic alignment at 80K iteration, and achieves better performance at 100K iteration.
Limitations. In practice, we find it challenging to design such trigger for precise termination because of: (i) Additional overhead: gradient evaluation introduces extra computational costs, limiting the practical evaluation frequency. (ii) Hyperparameter setting: the automatic trigger is sensitive to hyperparameters, such as block choice, timestep selection, and the threshold.
Therefore, since monitoring model performance during training is a standard practice in modern deep learning workflows, it's more practical to integrate the periodic evaluation of both metrics and gradient cosine similarity to determine the optimal termination point.
iteration
timestep
method
iteration
encoder
FID4
sFID.
1St
50K
60K
70K
80K
0.02
-0.0174
0,0418
0.0872
-0.0032
SYT-B/2
100K
63.46
7.31
20.6
0.1
0.0126
-0.0767
-0.0148
-0.1155
+REPA
100K
DINOv2-B
49.50
7.00
27.5
0.5
0.0134
0.0106
-0.0848
-0.1272
+REPA
100K
CLTP-ViT-B
54.92
7.63
24.7
0.8
0.0034
0.0406
0.1025
-0.0212
+HASTE
100K
DINOv2-B
39.86
7.16
35.8
average
0.0030
0.0041
0.0225
-0.0668
+HASTE
100K
CLIP-VIT-B
48.74
7.92
28.7
Table 9: Gradient cosine similarity between Table 11: Results of holistic alignment on SiT-B/2 with aligmment and denoising objectives. Bold CLIP-VIT-B/IE external encoder. While HASTE font denotes the average value below thresh- with CLIP consistently accelerates training, using old. We remove alignment 80K iteration. DINOv2-B achieves better performance.
method
iteration term. FID4 sFID4
IST
method iteration term. FID sFID
IS1
SiT-B/2
100K
63.46
7.31
20.6
SiT-B/2
150K
52.71
7.06
26.2
+HASTE
100K
39.86
7.16
35.8
+HASTE
150K
38.36
7.23
38.2
+HASTE
100K
80K
38.69
6.88
36.4
+HASTE
150K
100K
36.13
6.95
40.6
Table 10: Results of SiT-B/2+HASTE with Table 12: Validation of termination strategy on HASTE gradient angle trigger. Alignment termination with CLIP-ViT-BA ViT-B/16. Alignment termination 100K at 80K iteration leads to better performance. iteration achieves better performance.
22

---

A.3 Experiments with CLIP
In this section, we take CLIP-VIT-B/16 [38] as an example to investigate the impact of different pre-trained encoders on ImageNet 256 x 256. Additionally, we propose a solution for performing alignment when there is a mismatch in patch numbers between SiT and the external encoder.
Setting, We use S/T-B/2 for fast validation following the original settings of HASTE. Since SiT- B/2 and CLIP-ViT-B/II ViT-B/16 possess different patch numbers (i.e., 16 x 16 and 14 x 14), we perform additional spatial interpolations to align the dimensions:
(i) Feature map interpolation. Given SiT-B/2 block hidden states of shape [16 16, 768]. instead of directly merging channels from 256 into 196, we first reshape the sequence to [16, 16, 768] to preserve the 2D spatial topology of image patches. Subsequently, we apply bilinear interpolation to downsample the feature from spatial dimensions [16, 16] [14, 14).
(ii) Attention map interpolation. Similarly, given SiT-B/2 block attention map of shape (h. 16 x 16, 16 x 16), we reshape it into spatial tensor [h, 16, 16, 16, 16), where the dimensions represent head number, query height, query width, key height and key width respectively. Subsequently, we sequentially apply bilinear interpolation to key and query spaces downsample the attention map from dimensions [16, 16, 16, 16) to [14, 14, 14, 14].
Results (w/o efg). As shown in Tables 11 and 12, our implementations of both REPA and HASTE with CLIP-ViT BTT-B11 can accelerate the training process of S/T-B/2. However, using DINOv2-B B achieves better acceleration performance, which aligns with the findings in [55]: when the diffusion transformer is aligned with a pre-trained encoder that offers more semantically meaningful representations, the model exhibits enhanced generation performance. At 150K iteration, the holistic alignment early- stopped at 100K iteration achieves better performance, validating our termination strategy.
Consequently, validate the generalizability of HASTE over different pre-trained encoders. More- over, apart from directly varying the size of input image to the external encoder, we provide an interpolation solution to address the patch number mismatch between SiT and the encoder.
A.4 ImageNet 512 512 Experiment
We conduct an additional experiment on ImageNet 512 x 512 to validate the scalability of HASTE to image resolution. We follow the original settings of HASTE except we resize the input image to DINOv2 from 512 x 512 to 448 x 448 resolution following REPA. In practice,
300K iteration. Specifically,
stable diffusion VAE [40] compress the input image to SiT during training process, instead of pre-computing the latent vectors due to storage limitations.
method
epoch
termination
FID
sFID
1ST
SiT-XL/2
600
2.62
4.18
252.2
+REPA
80
2.55
4.16
241.2
+REPA
100
2.36
4.16
254.2
+HASTE
80
2.78
4.34
209.5
+HASTE
80
60
2.49
4.20
231.4
+HASTE
100
60
2.34
4.23
253.4
Table 13: Evaluation results ImageNet 512 x 512. 1 and 4 denote higher and lower values better, respectively. Bold font denotes the best performance.
resperereeeeeee
We report the quantitative results in Table 13, which reflects that HASTE can effectively accelerates the training process with various image resolutions, also validating our alignment and termination strategy. We evaluate the performance of SiT-XL/2+HASTE with SDE sampler (NFEs = 250, cfg scale w 1.25) and do apply guidance interval [28].
23

---

A.5
Detailed Quantitative Results
We provide detailed evaluation results of HASTE on different SiT models in Table 14. All results are reported with the SDE Suler-Maruyan sampler (NFEs 250) and without classifier- free guidance.
model
#params
iteration
FID, [16]
sFID [35]
1st [42]
Prec.f [27]
Rec.1 [27]
SIT-B/2 [34]
130M
400K
33.0
6,46
43.7
0.53
0.63
+HASTE
130M
100K
39.9
7.16
35.8
0.52
0.61
+HASTE
130M
200K
25.7
6.66
57.0
0.59
0.62
+HASTE
130M
400K
19.6
6.38
73.0
0.62
0.64
SiT-L/2 [34]
458M
400K
18.8
5.29
72.0
0.64
0.64
+HASTE
458M
100K
19.6
5.70
67.9
0.64
0.63
+HASTE
458M
200K
12.1
5.28
96.1
0.68
0.64
+HASTE
458M
400K
8,9
5.18
118,9
0.69
0.66
SiT-XL/2 [34]
675M
7M
8.6
6.32
131.7
0.68
0.67
+HASTE
675M
100K
15.9
5.64
78.1
0.67
0.62
+HASTE
675M
200K
9.9
5.04
108.8
0.69
0.64
+HASTE
675M
250K
8.4
4,90
119.6
0.70
0.65
+HASTE
675M
400K
7.3
5.05
128.7
0.72
0.64
+HASTE
675M
500K
5.3
4.72
148.5
0.73
0.65
Table 14: Additional evaluation results on ImageNet 256 x 256. f and 4 denote higher and lower
values are better, respectively. Bold font denotes the best performance.
Additionally, we provide the results of SiT-XL/2+HASTE with different classifier-free guidance [17] scales and intervals [28].
model
#params
iteration
interval
CFG scale
FID
sFID
1ST
Prec.t
Rec.1
SIT-XL/2
675M
[0. 1]
1.50
2.06
4.50
270.3
0.82
0.59
+HASTE
675M
500K
[0, 1]
1.25
2.18
4.67
240.4
0.81
0.60
+HASTE
675M
500K
0.71
1.50
1.80
4.58
252.1
0.80
0.61
+HASTE
675M
500K
[0, 0.6]
1.825
1.74
4.74
268.7
0.80
0.62
+HASTE
675M
2M
[0, 0.7]
1.7
1.45
4.55
297.3
0.80
0.64
+HASTE
675M
2M
[0, 0.71
1.65
1.44
4.56
289.4
0.79
0.64
+HASTE
675M
[0, 0.71
1.675
1.44
4.55
293.7
0.80
0.64
+HASTE
675M
2.5M
[0, 0.7]
1.7
1.43
4.56
298.8
0.80
0.64
+HASTE
675M
2.5M
[0, 0.7]
1.65
1.43
4.57
290.7
0.80
0.64
+HASTE
675M
2.5M
[0, 0.721
1.65
1.42
4.49
299.5
0.80
0.65
+HASTE
675M
3M
[0, 0.71
1.67
1.41
4.51
296.9
0.80
0.65
Table 15: Evaluation results ImageNet 256 256 with different classifier-free guidance settings.
24

---

B
Additional Implementation Details.
SiT-B
SIT-L
SiT-XL
DIT-XL
Architecture
imput dim.
32x32x4
32x32x4
32x32x4
32x32x4
num. layers
12
24
28
28
hidden dim.
768
1024
1152
1152
num. heads
12
16
HASTE
An
0.5
0.5
0.5
0.5
Aa
0.5
0.5
0.5
0.5
alignment depth
5
8
student layers
[2, 3, 4]
[4, 5, 6, 7]
[4, 5, 6, 7]
[4, 5,
teacher model
DINOv2-B [36]
DINOv2-B [36]
DINOv2-B [36]
DINOv2-B [36]
teacher layers
[7. 9, 11]
[8. 9. 10. 11]
[8. 9, 10, I1]
[8, 9, 10, 11]
termination iter.
100 K
250 K
250 K
250 K
alignment hcads
0-11
0-11
0-11
0-11
Optimization
batch size
256
256
256
256
optimizer
AdamW [25, 33]
AdamW [25, 33]
AdamW [25, 33]
AdamW [2.,, 33]
Ir
0.0001
0.0001
0.0001
0.0001
(31. (2)
(0.9, 0,999)
(0.9, 0.999)
(0.9, 0.999)
(0.9, 0.999)
weight decay
Diffusion
objective
linear interpolants
linear interpolants
linear interpolants velocity Euler- Maruyama
improved DDPM noise and variance
prediction
velocity
velocity Buler- Maler-Maruyama
sampler
Euler Maruyama
Euler-Mann Maler-Maruyama
sampling steps
250
250
250
250
Table 16: Detailed training settings.
Further implementation details. For XL and L-sized models, we set the feature alignment depth to 8 following REPA, and extract the attention maps from layer [4, 5, 6. 7] (counting from 0) of diffusion transformers, to align with those from layer [8. 9, 10, 11] of DINOv2-B, According to [31], the performance almost saturates when transferring 12 out of 16 heads, and the student can also develop its own attention patterns for unused heads. Specifically. since the number of heads for DINOv2-B layer is only 12, conduct attention alignment partially over the first 12 heads of diffusion transformer layer. For B-sized models, the feature alignment depth is adjusted to 5, and we extract the attention maps from layer [2, 3, 4] to align with those from layer [7. 9, 11] of DINOv2-B.
We enable mixed-precision (fp16) for efficient training. For data pre-processing, we leverage the protocols provided in EDM2 [21] to pre-compute latent vectors from images with stable diffusion decoder to translate generated latent vectors into images. Following REPA [55], we also use three-layer MLP with SiLU activations (7) as projector of hidden states, For MM-DiT, we use CLIP [38] text model encode captions.
VAE [40]. Specifically, we use wwwslllttasddddde--aaee
25

---

C HASTE Training Loop
For clarity, we provide pseudocode in Algorithm for the HASTE training loop with fixed iteration as the trigger to further illustrate our two-phase schedule in addition to Figure 2.
Algorithm 1 HASTE Training Loop
Require: Diffusion model Gg. Pre-trained encoder E
Require: Training dataset 2), Termination point
Require: Loss weights AR. AA Maximum iterations
Ensure: Trained diffusion model Go
1: for = to do
Phase I: Holistic Alignment Sample training batch
2:
Sample x, i, e from D, (O, (), I)
3:
It a4t Bie
4:
he < Hidden(Go, Go,
D DiT hidden states DiT attention maps D. Encoder features Encoder attention maps
5:
Ag < ALln(Go,
6:
+ F(xx)
7:
AE + Attn(E,x)
8:
Caift + ( (12, t), target)
9:
CREPA + CosSimm(/ (h+)
10:
CATTA + CrossEntropy (Ac. Ag)
L C Luil ARCREPA AACATIA
b Alignment and denoising Update parameters
12:
+ 0 - VoL
13:
end for
for
iori=ro + N do
Phase II: Pure Denoising Sample training batch
15:
from D) ((, (0, It - a B,E
Sample
16:
17:
L G t), target)
Denoising only
18:
0+0- - VoL
Update parameters
19:
end for
D Additional Visualizations
Figure 10: Uncurated generation results of SiT-XL/2+AASTE, We use classifier-free guidance
with w 4.0. Class label "loggerhead turtle" (33).
26

---

Figure 11: Uncurated generation results of SiT-XL/++ASTE. We use classifier-fre free guidance
with w 4.0. Class label = "macaw" (88).
Figure 12: Uncurated generation results of SiT-XL/2+HASTE. We use classifier-free guidance
with w = 4.0. Class label = "golden retriever" (207).
27

---

Figure 13: Uncurated generation results of SiT-XL/2+AASTE. We use classifier-fre guidance
with w 4.0. Class label = "arctic wolf" (270).
Figure 14: Uncurated generation results of SiT-XL/2+HASTE. We use classifier-free guidance
with w = 4.0. Class label = "red panda" (387).
28

---

Figure 15: Uncurated generation results of SAT-XL/2+AASTE. We use classifier-free guidance with w 4.0. Class label = "panda" (388).
Figure 16: Uncurated generation results of SiT-XL/2+HASTE. We use classifier-free guidance
with w 4.0. Class label = "acoustic guitar" (402).
29

---

Figure 17: Uncurated generation results of SiT-XL/2+HAST. We use classifier- free guidance
with w 4.0. Class label = "balloon" (417).
Figure 18: Uncurated generation results of SiT-XL/2+HASTE. We use classifier-free guidance with w 4.0. Class label = "baseball" (429).
30

---

Figure 19: Uncurated generation results of SiT-XL/2+ASSTE We use classifier free guidance with w 4.0. Class label = "dog sled" (537).
WITTEE
Figure 20: Uncurated generation results of SiT-XL/2+ASSE. We use Massifier-free guidance with w 4.0. Class label = "fire truck" (555).
31

---

Figure 21: Uncurated generation results of SiT-XL/+HASTE. We use classifier-free guidance with w 4.0. Class label = "laptop" (620).
Figure 22: Uncurated generation results of SiT-XL/2+HASTE. We use classifier-free guidance
with w = 4.0. Class label = "space shuttle" (812).
32

---

Figure 23: Uncurated generation results of SiT-XL/2+HSSE. We use classifier free guidance with w 4.0. Class label = "cheeseburger" (933).
Figure 24: Uncurated generation results of SiT-XL/2+HASTE. We use classifier-free guidance
with w 4.0. Class label = "cliff drop-off" (972).
33

---

Figure 25: Uncurated generation results of SiT-XL/2+HASTE We use classifier-free guidance with w 4.0. Class label = "coral reef" (973).
Figure 26: Uncurated generation results of SiT-XL/++AAST, We use classifier-free guidance
with w 4.0. Class label = "volcano" (980).
34


Reviews:
[{"review_id": "dFM9S2Mtx8", "reviewer": "Reviewer_KpE1", "summary": "This paper addresses the slow training of Diffusion Transformers (DiTs) by improving Representation Alignment (REPA). The authors observe that REPA accelerates early training but hinders performance later due to a \"capacity mismatch\" between the generative student model and the non-generative teacher (e.g., DINOv2). They propose HASTE (Holistic Alignment with Stage-wise Termination), a two-phase approach, which achieves up to 28× faster training on ImageNet and improves text-to-image generation on MS-COCO, matching or surpassing REPA without architectural changes. The method combines gradient-angle analysis to diagnose alignment conflicts and stage-wise termination to avoid over-regularization.", "questions": "Please refer to the weakness part", "limitations": "Please refer to the weakness part", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "Strength:\n- The authors identify the capacity mismatch between generative and non-generative models, supported by gradient-angle analysis showing how REPA shifts from helpful to harmful .\n- Integrating attention maps (relational priors) with feature alignment (semantic anchors) provides a more comprehensive guidance mechanism .\n- Extensive experiments on ImageNet and MS-COCO demonstrate significant speed-ups and performance gains, with ablation studies supporting the effectiveness of each component .\n\nWeakness:\n- The paper primarily uses fixed iterations for termination (e.g., τ=100K or 250K) but mentions gradient-angle thresholds without detailing how they are practically implemented or optimized .\n- While attention maps are shown to improve training, the paper lacks a theoretical framework explaining why mid-layer attention distillation is optimal, especially for diffusion models .\n- On MS-COCO, HASTE does not apply termination due to \"limited iterations,\" leaving unanswered questions about its effectiveness in text-guided generation over longer runs .\n- Combining feature and attention alignment may introduce additional training costs, which the paper does not explicitly quantify (e.g., memory or FLOPs) .\n- The choice of mid-layer alignment (e.g., SiT blocks 4–7) is empirical, but the paper does not systematically explore how layer depth affects performance .\n- Methods like TREAD [23] or FasterDiT [49] are cited but not directly compared in all experiments, limiting the scope of the performance claims .\n- The paper uses DINOv2 as the teacher, but it does not analyze how different pre-trained encoders (e.g., CLIP or Swin Transformer) affect HASTE’s performance .\n- While FID improvements are highlighted, the paper does not address scenarios where HASTE might underperform, such as when generating highly complex or rare visual concepts .", "was_revised": true, "final_justification": "Thanks for the authors' reply. I keep my original rating.", "rebuttal": "Thanks reviewer KpE1 for constructive comments. We are delighted to discuss the questions below and sincerely hope these answers can address the concerns.\n\n> **Q1:** Implementation of gradient-angle threshold.\n\n**A1:** Thanks for your question. Here’s a clarification and an updated revision for transparency.\n\n- **Dual termination strategy:** We mainly use fixed iteration $\\tau$ due to its simplicity and effectiveness (*e.g.*, ImageNet/MS-COCO). The gradient-angle threshold adds theoretical benefits. Gradient analysis (Lines 283–287) shows alignment gains drop around **250K iterations**.\n\n- **Gradient-angle trigger (plain implementation):** We train SiT-B/2 on ImageNet 256 $\\times$ 256 following the original settings of HASTE using a gradient-angle trigger.\n\n    - **Process:** \n        - After 50K iterations, every 10K, we check the gradient cosine similarity between alignment and denoising objectives.\n        \n        - At each check, we use selected timesteps $t$ at various noise levels on 2048 random ImageNet images; average the resulting similarities. If average $<$ threshold $\\delta$, we stop alignment.\n        \n        - Example uses 5th SiT block, timesteps $t=0.02$, $0.1$, $0.5$, $0.8$, and $\\delta=-0.05$.\n\n    - **Results:** \n        - **Cosine similarity:**\n        \n        |timestep\\iteration|50K|60K|70K|80K|\n        |-|-|-|-|-|\n        |0.02|-0.0174|0.0418|0.0872|-0.0032|\n        |0.1|0.0126|-0.0767|-0.0148|-0.1155|\n        |0.5|0.0134|0.0106|-0.0848|-0.1272|\n        |0.8|0.0034|0.0406|0.1025|-0.0212|\n        |avg|0.0030|0.0041|0.0225|**-0.0668(terminate)**|\n      \n        - **Performance:**\n      \n        |method|iter|terminate|FID$\\downarrow$|sFID$\\downarrow$|IS$\\uparrow$|\n        |-|-|-|-|-|-|\n        |SiT|100K|-|63.46|7.31|20.6|\n        |SiT+HASTE|100K|-|39.86|7.16|35.8|\n        |SiT+HASTE|100K|80K|**38.69**|**6.88**|**36.4**|\n    \n    - **Key point:** The automatic trigger ends holistic alignment at 80K and improves results at 100K.\n    \n    - **Limitations:** \n        - **Overhead:** Gradient checks add computation, limiting check frequency.\n        - **Hyperparameters:** Trigger is sensitive to hyperparameters (block, timestep, threshold).\n\n    - **Conclusion:** Since monitoring model performance during training is a standard practice in modern deep learning workflows, **it's more practical to integrate the periodic evaluation of both metrics and gradient cosine similarity to determine the optimal termination point**.\n\n\n> **Q2&Q5:** Exploration of alignment layer depth and theoretical framework for mid-layer attention distillation.\n\n**A2&A5:** Thanks for your question. Below are concise explorations of alignment depth and our choice of mid-layer for distillation.\n\n- **Alignment depth:**\n    - **Setup:** Train SiT-B/2 on ImageNet 256$\\times$256 with HASTE for 100K iters. **Vary alignment depth; no termination.**\n    \n    - **Results:**\n        |method|DINO|SiT|Depth|FID$\\downarrow$|sFID$\\downarrow$|IS$\\uparrow$|\n        |-|-|-|-|-|-|-|\n        |SiT|[7,9,11]|[2,3,4]|5|**39.86**|**7.16**|**35.8**|\n        |SiT+HASTE|[7,9,11]|[4,5,6]|7|40.59|7.27|35.6|\n        |SiT+HASTE|[7,9,11]|[6,7,8]|9|43.37|7.75|33.5|\n    \n    - **Conclusion:** Depth **5** yields best performance.\n\n- **Mid-layer distill choice:**\n    \n    - Alignment depth set per REPA (e.g., 8 for SiT-L/XL), to observe ATTA for HASTE.\n    \n    - As in Lines 171-183, mid-layer choice is due to:\n        - **Shallow mismatch:** Early layers should handle noisy input. (SiT-L/2: layer 4-7 > layer 2-7 in Table 6)\n        - **Deep freedom:** Deeper layers focus on denoising. Aligning too deep hurts results (REPA: depth >8 reduces performance). So, use layer 4-7 for alignment.\n\n- **Limits & future work:** While empirical layer choice works, theory may explain more. We plan to study attention in diffusion transformers further.\n\n\n> **Q3:** Validation of termination strategy in HASTE on MS-COCO.\n\n**A3:** To address your concern, we ran a T2I test on MS-COCO.\n\n- **Setup**: Using HASTE/REPA settings, we trained MMDiT with both from 150K-250K iterations, comparing MMDiT+HASTE with/without termination at $\\tau$=200K. At 250K, cfg set to 3.5 for eval; Sampler: SDE Euler-Maruyama, NFEs=250.\n\n- **Results**:\n  \n    |method|iter|term|FID$\\downarrow$|\n    |-|-|-|-|\n    |MMDiT+REPA|250K|-|4.28|\n    |MMDiT+HASTE|250K|-|4.10|\n    |MMDiT+HASTE|250K|200K|**4.06**|   \n    \n- **Conclusion:** Holistic alignment beats REPA for MMDiT; termination strategy further improves T2I.\n\n\n> **Q4:** Quantification of additional training costs introduced by holistic alignment.\n\n**A4:**\n\nThank you for the question. We evaluate the extra overhead of HASTE on SiT-XL/2 introduced by external vision model DINOv2-B:\n\n- **Extra memory overhead:** The peak CUDA memory with a local batchsize 32:\n\n    |HASTE|REPA|\n    |-|-|\n    | 30666.84 MB | 27818.01 MB |\n\n- **DINOv2-B forward overhead:** During training, we compute DINOv2‑B representations online following REPA. The forward pass of DINOv2‑B incurs approximately 23.19 GFLOPs for the base‑size model at 224×224 input resolution without XFormers.\n\n> **Q6**: Comparison with TREAD and FasterDiT.\n\n**A6:** Thanks for your comment. Here’s a detailed comparison with TREAD[1] and FasterDiT[2].\n\n- **ImageNet 256$\\times$256**\n  **Since TREAD does not report SiT+TREAD+cfg, we use their best DiT+TREAD result.**\n\n    - **Results:**\n\n      - w/o cfg\n\n        |method|epoch|FID$\\downarrow$|\n        |-|-|-|\n        |FasterDiT|400|7.91|\n        |SiT+HASTE|80|7.31|\n        |SiT+TREAD|80|**4.89**|\n\n      - w/ cfg\n\n        |method|epoch|FID$\\downarrow$|\n        |-|-|-|\n        |FasterDiT|400|2.03|\n        |DiT+TREAD|740|1.69|\n        |SiT+HASTE|400|**1.44**|\n\n   - **Analysis:** TREAD gets lower FID w/o cfg, but **HASTE is better w/ cfg**. TREAD notes, “Restricting info in training can hurt CFG performance.” Our HASTE gains a lot with cfg.\n\n     - HASTE at 600 epochs even outperforms REPA’s best FID:\n\n       |method|epoch|FID$\\downarrow$|sFID$\\downarrow$|IS$\\uparrow$|\n       |-|-|-|-|-|\n       |SiT+REPA|800|1.42|4.70|**305.7**|\n       |SiT+HASTE|600|**1.41**|**4.51**|296.2|\n\n   - **Conclusion:** As CFG is key for most gen tasks, HASTE is **more practical** for acceleration.\n\n- **MS-COCO 256$\\times$256**\n  We did not compare with FasterDiT or TREAD in T2I because:\n\n    - FasterDiT lacks T2I results/public code.\n    \n    - Our model/setup differs from TREAD:\n\n      - TREAD uses smaller DiT-B, HASTE uses larger MMDiT.\n\n      - TREAD uses FID-10K, HASTE uses FID-50K.\n\n\n> **Q7: Exploration of different pre-trained encoders.**\n\n**A7:** Thanks for your advice. To address this, we further explore using CLIP(ViT-B/16) as the pre-trained encoder:\n\n- **Setting:** We use SiT-B/2 for fast val as in HASTE's original setup. Since SiT-B/2 and CLIP-ViT-B/16 have different patch numbers (16x16 vs. 14x14), we interpolate spatial dims for alignment:\n\n  - **Feature map interp**: SiT hidden states `[bs,16x16,768]` are reshaped to `[bs,16,16,768]`, then downsampled with bilinear interp from `[16,16]` to `[14,14]`.\n\n  - **Attn map interp**: SiT attention maps `[bs,H,16x16,16x16]` are reshaped to `[bs,H,16,16,16,16]` and interpolated on key/query spaces.\n\n- **Results (w/o cfg):**\n\n  - **HASTE vs REPA at 100K:**\n\n    |method|iter|encoder|FID↓|sFID↓|IS↑|\n    |-|-|-|-|-|-|\n    |SiT|100K|-|63.46|7.31|20.6|\n    |SiT+REPA|100K|DINOv2-B|49.50|7.00|27.5|\n    |SiT+REPA|100K|CLIP-ViT-B/16|54.92|7.63|24.7|\n    |SiT+HASTE|100K|DINOv2-B|39.86|7.16|35.8|\n    |SiT+HASTE|100K|CLIP-ViT-B/16|48.74|7.92|28.7|\n\n  - **HASTE termination at 150K:**\n\n    |method|iter|encoder|term|FID↓|sFID↓|IS↑|\n    |-|-|-|-|-|-|-|\n    |SiT|150K|-|-|52.71|7.06|26.2|\n    |SiT+HASTE|150K|CLIP-ViT-B/16|-|38.36|7.23|38.2|\n    |SiT+HASTE|150K|CLIP-ViT-B/16|100K|36.13|6.95|40.5|\n\n- **Analysis:** At 100K, both REPA and HASTE with CLIP-ViT-B/16 speed up SiT-B/2 training. DINOv2-B gives better gains, matching [REPA] that **semantically meaningful pre-trained encoders** enhance gen performance. At 150K, earlier termination with HASTE further lowers FID.\n\n- **Conclusion:** HASTE with CLIP-ViT-B/16 shows its generalizability across different encoders. We handle patch mismatch by interpolation, not just via resizing. \n\nIn sum, Swin Transformer[3]'s window attn/patch merging means its features differ from those in diffusion transformers, so we use ViT-based encoders (e.g., DINOv2) for easier analysis and implementation; experiments focus on these.\n\n\n> **Q8: While FID improvements are highlighted, the paper does not address scenarios where HASTE might underperform, such as when generating highly complex or rare visual concepts.**\n\n**A8:** Thanks for your comment. First, we clarify that our work mainly focuses on:\n\n- relation between vision representative models & diffusion models,\n\n- and acceleration Methods for generative diffusion training,\n\nnot on solutions for spec. hard cases. So, we test mainly on widely-used gen. datasets (e.g., ImageNet), and note HASTE may underperform in some conditions.\n\nWe also evaluate HASTE in more complex cases:\n\n- **Text-to-image generation:** In our paper, we test HASTE on MS-COCO 256x256, where MMDiT captures complex semantics with ~80K train imgs.\n\n- **High-resolution image generation:** We also test HASTE on ImageNet 512x512. We set $\\tau=300K$, cfg scale 1.25.\n\n  |method|iters|term.|FID$\\downarrow$|sFID$\\downarrow$|IS$\\uparrow$|\n  |-|-|-|-|-|-|\n  |SiT+REPA|400K|-|2.55|**4.16**|241.2|\n  |SiT+REPA|500K|-|2.36|**4.16**|254.2|\n  |SiT+HASTE|400K|300K|2.49|4.20|231.4|\n  |SiT+HASTE|500K|300K|**2.34**|4.23|253.4|\n\nThanks for your insightful question, we will explore rare visual concept generation in future work.\n\n\n\n[1] TREAD: Token Routing for Efficient Architecture-agnostic Diffusion Training. In arXiv.\n\n[2] FasterDiT: Towards Faster Diffusion Transformers Training without Architecture Modification. In NeurIPS 2024.\n\n[3] Swin Transformer: Hierarchical Vision Transformer using Shifted Windows. In ICCV 2021."}, {"review_id": "O8SKrtaa1I", "reviewer": "Reviewer_fiaM", "summary": "This paper proposes improvements to REPA, a method for accelerating the training of diffusion models. Two key enhancements are introduced:\n\n- **Termination** $\\tau$: Instead of applying REPA throughout the entire training, it is stopped partway through.\n\n- **ATTA**: Beyond aligning features, an additional alignment loss encourages the student model’s attention maps to resemble those of the vision encoder.\n\nThese two strategies yield improvements over vanilla REPA in terms of both training efficiency and performance.", "questions": "- Could you clarify precisely what you mean by “capacity mismatch”? How does this relate to the divergence between gradient signals?\n\n- In ATTA, the CLIP attention quality seems degraded by attention sinks. Would it be possible to use register tokens to remove these sinks from the attention maps before computing ATTA? Recent work has proposed applying register tokens without additional training [1]\n\n[1]: Vision Transformers Don't Need Trained Registers (https://arxiv.org/abs/2506.08010)", "limitations": "The paper clearly discusses limitations related to model architecture, pixel-level diffusion, and video generation, which is appreciated.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "**Strengths**\n\n- The writing is very clear and accessible, making the paper easy to follow.\n\n- The proposed solution is simple yet effective, achieving substantial gains.\n\n- The ablation studies clearly demonstrate how each component contributes to improving REPA, with Table 3 particularly compelling.\n\n**Weaknesses**\n\n- The approach heavily relies on heuristics. For example, while the authors attribute the necessity of termination to a “capacity mismatch,” the evidence in Section 2.2 does not strongly support this explanation. It might have been more straightforward to describe this simply as gradient signal interference.\n\n- Additionally, I find Figures 3 and 4 somewhat disappointing and potentially misleading. In Appendix A, the cosine similarity plots show that the gradient signals only begin to point in opposite directions after much longer training (4M iterations), rather than within the training regime actually used (500K iterations).", "was_revised": true, "final_justification": "I believe this paper suggests a simple but effective strategy for improving REPA. However, because the REPA is just a single regularization technique for training the diffusion model, I think the impact is limited. Therefore I suggest BA.", "rebuttal": "Thanks reviewer fiaM for adoring our work. We are delighted to discuss the questions below.\n\n> **W1&Q1:** Relationship between \"capacity mismatch\" and \"gradient divergence\".\nThanks for your insightful question. We first provide a further detailed clarification of \"capacity mismatch\", and then explain how it connects with \"gradient divergence\".\n\n- **Capacity Mismatch:** We refer \"capacity mismatch\" to two **fundamental** differences between vision representation models (e.g., DINOv2) and diffusion models.\n    - **Task mismatch:** Diffusion models are designed to generate high-quality samples through denoising, while representation models are aimed at capturing semantic information and constructing good representations. Therefore, the two kinds of models possess fundamentally different objectives.\n    - **Distribution mismatch:** Diffusion models learn complex, noisy distributions across different noise levels. In contrast, representation models learn only through the clean marginal data distribution. In the training process, external representations help diffusion models to get a coarse and fast grasp of the generative task but fail to further improve the fine-grained generation capacity, especially at diffusion timesteps with lower noise levels.\n- **Relationship**: \n    - **Statement concern:** While we acknowledge that gradient interference is a more direct description, we would like to clarify that capacity mismatch serves as the cause and provides deeper insight, which helps to explain the source of training constraint. We also follow that insight to design the low-frequency filter experiment in Figure 5, Section 2.2.\n    - **Experimental evidence:** In Section 2.2, we leverage gradient analysis for diagnosis and observe the Ignition-Plateau-Conflict trends, which are consistent with the shift of REPA from accelerator to brake. Additionally, our low-frequency filter experiment further verifies the existence of capacity mismatch.\n    - **Mismatch and divergence:** When diffusion models seek high-frequency detail to further improve the generation quality, the capacity mismatch leads to the constraint of improvement, which is obvious at diffusion timesteps at lower noise levels (e.g., 0-0.1). At these timesteps, we observe evident gradient divergence.\n\n> **W2:** Mismatch between gradient direction figures and cosine similarity plots.\n\nThanks for your careful observation. We understand the concern about potential inconsistency between the main text figures and the appendix data. Please allow us to clarify this important point.\n\n- **Timestep-dependent gradient evolution:**\n\n  As reported in Appendix A.1, we observe that gradient cosine similarities across different diffusion timesteps exhibit similar overall trends—starting relatively high, gradually decreasing, and eventually becoming negative. However, **the rate of this evolution varies significantly across timesteps**:\n\n  - For t ∈ (0, 0.1] (fine detail generation), we observe negative cosine within 500K iterations.\n  - For larger timesteps (coarser features), similar behavior appears much later.\n  - Figures 3 and 4 are **schematic illustrations** intended to demonstrate this general trend, with line 98 explicitly noting that Figure 3 uses t = 0.1 as an example.\n\n- **Balancing alignment benefits and conflicts**\n\n  Our method fundamentally seeks to balance two competing effects:\n\n  - **Benefits**: Alignment provides valuable guidance for low-frequency, structural features\n  - **Costs**: Alignment eventually hinders the learning of high-frequency details\n\n  The key is to **terminate before the harm to fine-grained details outweighs the benefits to coarse-grained structure**. Notably, our actual termination point (250K iterations for HASTE) is chosen well before gradient similarities become fully negative across all timesteps. This reflects a principled trade-off rather than waiting for complete gradient opposition. Additionally, our experimental results with HASTE strongly validate this termination strategy.\n\n> **Q2:** Using register tokens to remove attention sinks in CLIP.\n\nThanks for your advice. First, we would like to clarify that our original submission doesn't include experiments of ATTA with CLIP attention. Nevertheless, we can provide the experimental results of HASTE with CLIP-ViT-B/16.\n\n- **Setting:** We use SiT-B/2 for fast validation following the original settings of HASTE. Since SiT-B/2 and CLIP-ViT-B/16 possess different patch numbers (i.e., 16x16 and 14x14), we perform additional spatial interpolations to align the dimensions.\n\n- **Results (w/o cfg):**\n\n  - Comparing HASTE with REPA at 100K iteration:\n\n    | method    | iteration | encoder       | FID↓      | sFID↓    | IS↑      |\n    | - | - | - | - | - | - |\n    | SiT       | 100K      | -             | 63.46     | 7.31     | 20.6     |\n    | SiT+REPA  | 100K      | DINOv2-B      | 49.50     | **7.00** | 27.5     |\n    | SiT+REPA  | 100K      | CLIP-ViT-B/16 | 54.92     | 7.63     | 24.7     |\n    | SiT+HASTE | 100K      | DINOv2-B      | **39.86** | 7.16     | **35.8** |\n    | SiT+HASTE | 100K      | CLIP-ViT-B/16 | 48.74     | 7.92     | 28.7     |\n\n  - Validating the termination strategy of HASTE:\n\n    | method    | iteration | encoder       | termination | FID↓      | sFID↓    | IS↑      |\n    | - | - | - | - | - | - | - |\n    | SiT       | 150K      | -             | -           | 52.71     | 7.06     | 26.2     |\n    | SiT+HASTE | 150K      | CLIP-ViT-B/16 | -           | 38.36     | 7.23     | 38.2     |\n    | SiT+HASTE | 150K      | CLIP-ViT-B/16 | 100K        | **36.13** | **6.95** | **40.5** |\n\n- **Register tokens:** We sincerely appreciate the insightful paper you referred to. We think it's a promising direction to leverage the register tokens without additional training to improve the attention quality of pre-trained encoders. We will discuss the possibility introduced by the proposed method in **Limitations and future work**,  Section 5."}, {"review_id": "Vgky3nwlzz", "reviewer": "Reviewer_qGU6", "summary": "This paper identifies a critical failure mode of Representation Alignment (REPA) in accelerating DiT training: as the DiT matures, REPA’s fixed teacher becomes a constraint. The authors diagnose this as a capacity mismatch via gradient analysis and propose HASTE, which combines holistic alignment (features + attention) with stage-wise termination of the alignment loss. This enables rapid early convergence and avoids late-stage degradation. Experiments on ImageNet show up to 28x acceleration over vanilla training.", "questions": "- Can you provide experiments validating gradient-angle-based automatic triggers?\n- Could you test termination in text-to-image training to strengthen claims of generality?\n- Do you have guidelines for τ selection without exhaustive tuning?\n- for clarity, would it be possible to include pseudocode for the HASTE training loop?\n- The ablation studies show that attention alignment (ATTA) by itself provides a significant acceleration, similar to REPA. In data-scarce or computationally constrained scenarios, would using ATTA alone be a more efficient and sufficient strategy than the full HASTE recipe?", "limitations": "only thing is that the paper lacks discussion of risks from faster generative model training", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 2, "originality": 3, "strengths_and_weaknesses": "**Strength**\n\n- Insightful analysis of when REPA becomes counterproductive.\n- HASTE dramatically accelerates DiT training on ImageNet\n- Well-written \n- The method is simple and implementable.\n\n**Weakness*\n- Optimal t requires manual tuning. While automatic triggers was suggested but not tested experimentally.\n- Termination not tested on MS-COCO text-to-image tasks, undermining general claims of generalizability.\n- Paper lacks discussion of risks from faster generative model training", "was_revised": true, "final_justification": "My biggest concern was that the termination strategy wasn't tested on text-to-image models . The MS-COCO experiment successfully showed that the termination strategy works there too, which is a crucial addition as a result I'm happy to increase my score.", "rebuttal": "Thanks for the insightful and motivational comments from reviewer qGU6. We discuss the questions as follows.\n\n> **Q1:** Validating gradient-angle-based automatic trigger.\n\n**A1:** Thanks for the question. To address this concern, we provide a plain implementation as an example for an automatic gradient-angle trigger as follows.\n\n- **Setting:** We train SiT-B/2 on ImageNet 256 $\\times$ 256 following the original settings of HASTE using a gradient-angle trigger.\n- **Implementation:** \n    - We start evaluating the gradient cosine similarity between holistic alignment and denoising objectives every 10K iterations from 50K iterations, when the training process has been relatively stable. \n    - In each round of evaluation, we compute the similarity separately on selected diffusion timesteps $t$ at different noise levels with 2048 images randomly sampled from ImageNet, and take the average value. When the average similarity falls below a threshold $\\delta$, we terminate the holistic alignment.\n    - In the plain example, we perform gradient evaluation on the 5th SiT block (counting from 1). We select $t=0.02$, $t=0.1$, $t=0.5$ and $t=0.8$, and set $\\delta=-0.05$.\n- **Results:** \n    - **Gradient cosine similarity:**\n\n        |timestep\\iteration|50K|60K|70K|80K|\n        |-|-|-|-|-\n        |0.02|-0.0174|0.0418|0.0872|-0.0032|\n        |0.1|0.0126|-0.0767|-0.0148|-0.1155|\n        |0.5|0.0134|0.0106|-0.0848|-0.1272|\n        |0.8|0.0034|0.0406|0.1025|-0.0212|\n        |average|0.0030|0.0041|0.0225|**-0.0668(terminate)**|\n    - **Performance:**\n        |method|iteration|termination|FID$\\downarrow$|sFID$\\downarrow$|IS$\\uparrow$|\n        |-|-|-|-|-|-|\n        |SiT|100K|-|63.46|7.31|20.6|\n        |SiT+HASTE|100K|-|39.86|7.16|35.8|\n        |SiT+HASTE|100K|80K|**38.69**|**6.88**|**36.4**|\n\n- **Findings:** Our automatic trigger terminates the holistic alignment at 80K iterations, and achieves better performance at 100K iterations.\n\n- **Limitations:** Though we validate the automatic gradient-angle trigger on SiT-B/2, we find it challenging to design such a trigger for precise termination, because of:\n    - **Additional overhead:** Gradient evaluation introduces extra computational cost, limiting the practical evaluation frequency.\n    - **Hyperparameter setting:** The automatic trigger is sensitive to hyperparameters, such as block choice, timestep selection, and the threshold.\n\n- **Conclusion:** Given that monitoring model performance during training is a standard practice in modern deep learning workflows, we think **it's more practical to integrate the periodic evaluation of both metrics and gradient cosine similarity to determine the optimal termination point**.\n\n\n> **Q2:** Validation of termination strategy in HASTE on MS-COCO.\n\n**A2:** To address this concern, we conduct the following text-to-image (T2I) generation experiment on MS-COCO.\n\n- **Setting**: Following the original settings of HASTE and REPA on MS-COCO, we continue to train MMDiT with both methods from 150K to 250K iterations, and compare the performance of MMDiT+HASTE with and without termination $\\tau$=200K. As training progresses, we raise the classifier-free guidance (cfg) [1] scale to 3.5 for all evaluations at 250K iterations to steer the sampling process within the data manifold. We consistently use the SDE Euler-Maruyama sampler with NFEs=250.\n- **Results**:\n    |method|iteration|termination|FID$\\downarrow$|\n    |-|-|-|-|\n    |MMDiT+REPA|250K|-|4.28|\n    |MMDiT+HASTE|250K|-|4.10|\n    |MMDiT+HASTE|250K|200K|**4.06**|\n- **Conclusion:** As expected, while our holistic alignment outperforms REPA for MMDiT, the termination strategy also leads to better performance in the T2I task. \n\n> **Q3:** Guidelines for τ selection.\n\n**A3:** Thanks for the question. We first explain our selection strategy in practice, and then give guidelines for τ selection.\n- **\"Exhaustive\" tuning is not necessary:**\n    - **Gradient analysis helps:** In practice, we don't really exhaustively tune $\\tau$, because the gradient analysis helps us to locate the termination point. As we state in Lines 283-287, we choose to stop around 250K iterations when we observe **limited performance improvement** and **relatively low gradient cosine similarities**.\n    - **Tolerance of $\\tau:$** Additionally, we find HASTE is relatively robust to $\\tau$. In Figure 7b, we show that varying $\\tau$ from 250K to 290K doesn't lead to a significant difference in the FID of SiT-XL/2 at 400K iteration.\n- **Guidelines for τ selection:** As we have mentioned above, the periodic evaluation of metrics and gradient cosine similarity helps to locate the approximate termination period. Additionally, termination within such a period works reliably well for HASTE and is more practical than locating the precise point.\n\n> **Q4:** Pseudocode for HASTE training loop.\n\nThanks for the advice. We will add it to our latest version.\n\n> **Q5:** Using ATTA alone in data-scarce or computationally constrained scenarios.\n\n**A5:** Thank you for the insightful advice. In the following, we provide experiments on MS-COCO to explore the performance of ATTA, HASTE, and REPA in the given scenario.\n\n- **Setting:** To simulate data-scarce or computationally constrained conditions, we sample **20K** images ($<\\frac{1}{4}$ original size) from MS-COCO training dataset, and set the batchsize of **64** ($\\frac{1}{4}$ default batchsize), and train MMDiT for only **50K** iterations. We do not apply termination for HASTE.\n- **Results:**\n    |method|iteration|FID|\n    |-|-|-|\n    |MMDiT|50K|49.03|\n    |MMDiT+REPA|50K|32.59|\n    |MMDiT+ATTA|50K|43.17|\n    |MMDiT+HASTE|50K|**31.39**|\n- **Findings:** Under the specific setting, all three methods provide considerable acceleration to the training process. HASTE works best and REPA performs better than ATTA.\n- **Analysis:** As we mentioned in **Lines 167-169**, ATTA nails routing but lacks supervision over condition processing, which is even more evident in MMDiT architecture: we only perform attention alignment with the $QK^T$ matrix generated from the input image to avoid affecting textual processing. In contrast, REPA performs alignment over the hidden states of MMDiT, acting on both vision and textual processing. Therefore, combining feature and attention alignment provides more holistic guidance for training.\n- **Conclusion:** In data-scarce or computationally constrained scenarios, ATTA can accelerate the training process, but not as effectively as REPA and HASTE.\n\n> **Q6:** Discussion of risks from faster generative model training.\n\n**A6:** Thanks for the concern. We acknowledge that faster generative model training could potentially **lower barriers for harmful content generation** and **accelerate deepfake production**.\n\nWhile we recognize these concerns, we believe HASTE's risk profile is relatively contained for several reasons:\n- **Research-focused contribution:** Our method primarily benefits legitimate research by **reducing computational barriers**, enabling smaller labs and researchers to participate in generative modeling research.\n- **Safeguards applicable:** Standard safety measures, such as content filtering, ethical guidelines, and responsible disclosure, remain **compatible** with HASTE-trained models.\n\n**Limitations and future work:** As a technical contribution focused on training efficiency, our work doesn't directly address broader AI safety challenges, which we acknowledge as important future work. We are going to include safety concerns in Section 5 in the revision."}, {"review_id": "b7RkDNRs2F", "reviewer": "Reviewer_FRQ1", "summary": "This paper identifies limitations in REPA, a regularization method aimed at accelerating diffusion transformer training, and proposes simple improvements. The authors show that REPA is only helpful in the early stages of training and may become detrimental later, advocating for an early stopping mechanism that disables REPA in the middle of the training. Furthermore, they introduce a holistic alignment, which aligns not only the features of DINO and the diffusion transformer (as in REPA) but also their attention maps. The paper empirically demonstrates the benefits of both REPA early stopping and attention alignment through ImageNet generation experiments.", "questions": "In Table 4, it's unclear why the performance of SiT-XL/2 at 500K iterations is worse than at 400K. This inconsistency is not addressed.", "limitations": "In the limitations section, the paper should explicitly state that the use of external pre-trained models like DINO requires additional pre-training computation and introduces extra computational overhead during diffusion model training.", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "## Strengths\n\n- Clearly identifies and analyzes the limitations of REPA, offering valuable insights for the generative AI community.\n\n- Proposes attention alignment (ATTA), a novel idea not yet explored in diffusion model literature.\n\n- Demonstrates faster convergence compared to DiT, SiT, and REPA baselines.\n\n## Weaknesses\n\n- The paper is poorly written, with missing explanations in several figures. For instance, Figure 5 does not clarify what [0, 1], [0.25, 1], etc., mean, nor how low-frequency information is filtered. In Figure 6 (b and c), the lack of baselines makes it difficult to interpret the quality of the results.\n\n- While the paper shows that REPA mainly helps with learning low-frequency semantic information early in training, it overlooks a simpler baseline: timestep sampling biased toward high noise levels, as in EDM [A, B] and Stable Diffusion 3 [C]. If learning low-frequency semantics early is beneficial, allocating more training computation to larger timesteps (which carry low-frequency information) may serve as a stronger and more straightforward baseline than using external vision encoders.\n\n- While Line 190 mentions the use of a gradient-angle trigger as an alternative to fixed-iteration early stopping, it appears that this trigger is not actually used in the experiments.\n\n[A] Karras et al., Elucidating the Design Space of Diffusion-Based Generative Models, NeurIPS 2022.\n\n[B] Karras et al., Analyzing and Improving the Training Dynamics of Diffusion Models, CVPR 2024.\n\n[C] Esser et al., Scaling Rectified Flow Transformers for High-Resolution Image Synthesis, ICML 2024.", "was_revised": true, "final_justification": "This paper shows that the REPA regularization, which is effective for training diffusion models, is only beneficial in the early iterations of training, and proposes a simple method to address this limitation. \n\nHowever, the “gradient trigger” component of the method is not actually used in practice, yet it occupies a substantial portion of the method description, which may confuse readers. The authors explain that it can be used for analysis purposes, but in the current manuscript, it is unclear whether the gradient trigger is merely for analysis or is an effective part of the method. This gives the paper an immature impression in this regard. \n\nTherefore, I lower my score to 3. Nonetheless, the paper’s core observation is meaningful, and I will not oppose the accept recommendation.", "rebuttal": "Thanks for your comments. We discuss the questions as follows.\n\n> **Q1:** Missing explanations in several figures.\n\nThank you for your comment. We acknowledge that the illustrations of Figures 5,6 can be clearer. Therefore, we provide detailed explanations of them.\n\n- **Figure 5:** The time interval [0, 0.25] stands for the diffusion timesteps where the input to pre-trained encoders is replaced with low-frequecy part of the image. For image filtering, we convert the image to frequency domain using the Fourier transform, and filter out the high-frequency components, then invertly convert it back to an image.\n- **Figure 6:** We compute the feature and attention alignment progress mainly to explain the complementary effect of both methods. We will add baselines to (b) and (c) in the latest version.\n\n> **Q2:** Will allocating more training computation to larger timesteps with higher noise levels serve as a better baseline?\n\nThanks for your insightful question. We provide comparision of REPA, HASTE, and log-normal timestep sampling on SiT-B/2.\n\n| method         | iteration | FID$\\downarrow$ | sFID$\\downarrow$ | IS$\\uparrow$ |\n| - | - | - | - | - |\n| SiT+log-normal | 50K       | 92.27           | **11.11**        | 14.76        |\n| SiT+REPA       | 50K       | 78.20           | 11.71            | 17.1         |\n| SiT+HASTE      | 50K       | **71.32**       | 13.85            | **19.26**    |\n\nWhile log-normal has been proved effective on models like EDM, we find acceleration methods based on external representations still achieve better performance. \n\n> **Q3:** The use of gradient-angle trigger.\n\nThanks for your important question. In the following, we first explain the role of gradient signals in our method, and then provide a plain example of automatic gradient-angle trigger.\n- **Dual termination strategy clarification:** Our method employs a hybrid approach where fixed iterations serve as the primary and practical termination strategy, while the gradient-angle rule provides theoretical insight and robustness across different datasets. Specifically:\n    - In practice, we find setting a fixed termination iteration $\\tau$ for HASTE actually performs reliably well in ImageNet and MS-COCO experiments. Gradient-angle threshold serves as a supplemental alternative, adding adaptability to different scenerios.\n    - The gradient analysis helps us to locate the termination point. As we state in Lines 283-287, we choose to stop around 250K iteration when we observe limited alignment benefits and relatively low gradient cosine similarities.\n- **Plain implementation of automatic gradient-angle trigger:** We train SiT-B/2 on ImageNet 256 $\\times$ 256 following original settings of HASTE using a gradient-angle trigger.\n    - **Implementation:** \n        - We start evaluating the gradient cosine similarity between holistic alignment and denoising objectives every 10K iterations from 50K iteration, when the training process has been relatively stable. \n        - In each round of evaluation, we compute the similarity separately on selected diffusion timesteps $t$ at different noise levels with 2048 images randomly sampled from ImageNet, and take the average value. When the average similarity falls below a threshold $\\delta$, we terminate the holistic alignment.\n        - In the plain example, we perform gradient evaluation on the 5th SiT block (counting from 1). We select $t=0.02$, $t=0.1$, $t=0.5$ and $t=0.8$, and set $\\delta=-0.05$.\n    - **Results:** \n        - **Gradient cosine similarity:**\n        \n        |timestep\\iteration|50K|60K|70K|80K|\n        |-|-|-|-|-\n        |0.02|-0.0174|0.0418|0.0872|-0.0032|\n        |0.1|0.0126|-0.0767|-0.0148|-0.1155|\n        |0.5|0.0134|0.0106|-0.0848|-0.1272|\n        |0.8|0.0034|0.0406|0.1025|-0.0212|\n        |average|0.0030|0.0041|0.0225|**-0.0668(terminate)**|\n        - **Performance:**\n\n        |method|iteration|termination|FID$\\downarrow$|sFID$\\downarrow$|IS$\\uparrow$|\n        |-|-|-|-|-|-|\n        |SiT|100K|-|63.46|7.31|20.6|\n        |SiT+HASTE|100K|-|39.86|7.16|35.8|\n        |SiT+HASTE|100K|80K|**38.69**|**6.88**|**36.4**|\n        \n    - **Findings:** Our automatic trigger terminates the holistic alignment at 80K iteration, and achieves better performance at 100K iteration.\n    - **Limitations:** Though we validate the automatic gradient-angle trigger on SiT-B/2, we find it challenging to design such trigger for precise termination, because of:\n        - **Additional overhead:** Gradient evaluation introduces extra computational cost, limiting the practical evaluation frequency.\n        - **Hyperparameter setting:** The automatic trigger is sensitive to hyperparameters, such as block choice, timestep selection, and the threshold.\n    - **Conclusion:** Given that monitoring model performance during training is a standard practice in modern deep learning workflows, we think **it's more practical to integrate the periodic evaluation of both metrics and gradient cosine similarity to determine the optimal termination point**.\n> **Q4:** Inconsistent performance of SiT-XL/2 at 400K and 500K iteration.\n\nThanks for your question. As we mentioned in Line 277-281, we attribute the inconsistency to over-regularization introduced by continuously using holistic alignment. Actually, our termination strategy at 250K iteration helps to alleviate such a trend and demonstrates better performance at 500K iteration.\n\n> **Q5:** The use of external pre-trained models introduces extra training cost.\n\nThank you for the question. We evaluate the extra training overhead of HASTE on SiT-XL/2 introduced by external vision model DINOv2-B:\n\n- **Extra memory overhead:** The peak CUDA memory with local batchsize 32:\n\n\t| HASTE       | REPA        |\n\t| -| - |\n\t| 30666.84 MB | 27818.01 MB |\n\n- **DINOv2-B forward overhead:** During training, we compute DINOv2‑B representations online following REPA. The forward pass of DINOv2‑B incurs approximately 23.19 GFLOPs for the base‑size model at 224×224 input resolution without XFormers.\n\nWe will add these evaluations in our latest version.\n\n> **Formatting Concern:** \"The paper would benefit from a figure that illustrates the REPA and ATTA regularization mechanisms.\"\n\nThanks for your advice. We will add the illustrations in the latest version."}]

Author rebuttal:

Meta-review: This paper identifies the issues of current representation alignment for diffusion models with empirical analysis. Then the paper proposes an early stopping strategy for representation alignment. Additionally, the representation alignment is extended to hybrid attention alignment, and the alignment is only applied to certain blocks.

Reviewers commented positively on the identification of the capacity mismatch, the associated empirical studies, the enhanced alignment integrating attention, and extensive experiments.

Reviewers commented negatively on the naive strategy for termination using fixed iterations, the method design is mainly based on empirical and heuristic approaches, and the limited impact. Particularly, the gradient-angle trigger is only for analysis, not used as part of the method.

Since the insightful empirical analysis is the main contribution of this paper, not using the gradient-angle trigger does not undermine the main contribution. Therefore, it is recommended to accept. But the authors do need to carefully revise the related wording in the paper to avoid misunderstanding.
Decision: Accept (poster)