2. Related Work

REPA established the basic claim that diffusion transformers benefit from aligning intermediate hidden states to pretrained visual features. iREPA strengthened that result by replacing the MLP projector with a convolutional projection and spatial normalization, arguing that spatial structure matters more than global semantic strength alone.

HASTE is the main bridge to the later stages of this project. It suggests that alignment helps early yet can plateau or become restrictive later, which turns schedule design into a first-class question rather than a minor tuning detail.

Related directions such as SARA, RAE, and U-REPA broaden the representation-alignment design space. The open problem is no longer whether external guidance can help, but which parts of that guidance are worth paying for. This project addresses that question through a staged study of layer choice and schedule design.

4. Stage 1: Baseline Confirmation and Layer Ablation

Stage 1 establishes the aligned anchor used by the later schedule studies. It first checks that a faithful small-scale implementation preserves the expected ordering of Vanilla, REPA, and iREPA, then asks the first allocation question: which internal layer should receive alignment?

4.1 Baseline Confirmation

The baseline runs provide a clean signal. Vanilla is a valid control, while both aligned variants separate from it by the middle of training. The main Stage 1 result is therefore not just final ranking, but faster convergence under alignment.

Figure 1. Training-side dynamics across Vanilla, REPA, and iREPA. Vanilla remains a clean control, REPA maintains active alignment, and iREPA reaches the lowest final denoising loss.

Figure 2. Checkpoint evaluation curves. The main signal is not only the final ranking but the faster mid-stage improvement of aligned training relative to Vanilla.

The mid-stage checkpoints make the separation easiest to read. At 30k, REPA already leads Vanilla on all three metrics, while iREPA is competitive on FID and KID. By 50k, both aligned variants outperform Vanilla across the board, with iREPA taking the strongest FID/KID position. This establishes a credible aligned reference trajectory for the later schedule studies.

4.2 Layer Ablation for iREPA

We next vary the aligned block while keeping the rest of the iREPA recipe fixed. The practical question is whether earlier layers improve mid-stage acceleration enough to challenge the standard Layer 8 anchor.

Figure 3. iREPA layer-ablation checkpoint curves. Earlier layers are competitive in the faster-convergence regime, while Layer 8 remains the strongest final endpoint.

The iREPA sweep is tight, but the pattern is clear enough: earlier layers are competitive in the fast-convergence regime, while Layer 8 remains the most stable final anchor. The practical takeaway is that iREPA benefits from earlier-layer acceleration without forcing a change to the baseline configuration inherited by Stages 2 and 3.

4.3 Layer Ablation for REPA

We repeat the same sweep for REPA to test whether the early-layer effect is specific to iREPA or a broader property of representation-aligned training.

Figure 4. REPA layer-ablation checkpoint curves. Layer 4 is best very early, reinforcing the view that earlier layers can help reach a better regime faster.

REPA shows the sharper early-layer effect. Layer 4 is best at the earliest checkpoint, and the separation is cleaner than in iREPA. By the mid-stage checkpoints the curves regroup, which suggests that the main gain of earlier-layer alignment is faster entry into a good regime rather than a radically different final endpoint.

4.4 Stage 1 Takeaway

Stage 1 establishes a usable anchor for the rest of the project. Alignment is beneficial, the strongest signal is the mid-stage convergence advantage, and earlier layers can help the model reach that regime faster. With that anchor in place, the next question is no longer whether alignment works, but how it should be scheduled.

6. Stage 3: Parameterized Schedule Study

Stage 2 shows that hard cutoffs can approximate always-on alignment at low cost, but they are blunt instruments — the schedule is fixed before training begins and cannot respond to how the optimization is actually unfolding. That raises a natural follow-up: can we design a parameterized schedule that is expressive enough to study alignment decay, while also being a credible candidate for improvement? A smooth decay gives us a continuous knob to study whether quality differences are structural — the schedule is redirecting the optimization path — or simply an efficiency problem — the model is on the same path but progressing more slowly. Stage 3 tests this with a single comparison: cosine annealing versus always-on iREPA over a 250k horizon.

6.1 Long-Horizon Cosine Schedule at 250k

Schedule Definition

Let the training step be \(k\), the total number of steps be \(T\), the denoising loss be \(\mathcal{L}_{\text{denoise}}\), and the projection loss be \(\mathcal{L}_{\text{proj}}\). The base objective is

$$ \mathcal{L}(k) = \mathcal{L}_{\text{denoise}}(k) + \lambda_{\text{proj}}(k)\,\mathcal{L}_{\text{proj}}(k). $$

For the fixed schedule, the projection multiplier remains constant:

$$ \lambda_{\text{proj}}^{\text{fixed}}(k) = 1. $$

For the cosine schedule implemented in the loss function, define normalized progress \(p = k / T\), start fraction \(s\), end fraction \(e\), and normalized in-window coordinate \(x = (p - s)/(e - s)\). Then

$$ \lambda_{\text{proj}}^{\text{cos}}(k) = \begin{cases} \frac{1}{2}\left(1 + \cos(\pi x)\right), & s \le p < e, \\ 0, & \text{otherwise.} \end{cases} $$

In the cosine run, the interval spans the full training horizon, so \(s = 0\) and \(e = 1\). At initialization, \(p = 0\), hence \(x = 0\) and \(\lambda_{\text{proj}}^{\text{cos}}(0) = 1\), matching the fixed iREPA weight exactly at step zero and decaying continuously to zero by step \(T\).

Training Diagnostics

Figures 10 and 13 show the projection-loss decay and fixed-seed sample evolution across the 250k horizon — the cosine weight tracks the fixed schedule closely early on before diverging smoothly as the multiplier decays.

Figure 13. Stage 3 synchronized training timeline. Scrub or autoplay to inspect fixed-seed sample evolution across the 250k horizon for fixed iREPA and cosine iREPA.

Metric Comparison

Cosine annealing opens with an early FID and IS deficit — even a gentle decay relaxes projection pressure before fidelity has consolidated. The gap narrows through the mid-horizon as the diminishing-returns dynamic from Stage 2 takes hold. The decisive inflection is between 200k and 250k: by the final checkpoint cosine overtakes fixed iREPA on FID, KID, and IS simultaneously, a benefit that only materialises when the horizon is long enough for late-stage relaxation to matter.

Pareto Frontier View

Both schedules make joint progress toward the upper-right (preferred) region through 100k. After that they diverge — and what is particularly interesting here is not just that cosine ends better, but how the fixed schedule behaves. iREPA_full backtracks by 200k, regressing on both IS and log(1/FID) before a partial late recovery. This is not a plateau but a genuine reversal, suggesting that sustaining full projection pressure deep into training may gradually push the model away from a jointly preferred outcome. Scalar endpoints would smooth over this — the final numbers might still look reasonable — but the trajectory makes the instability visible.

Cosine annealing appears to avoid this. By gradually relaxing projection pressure, it seems to prevent the model from becoming over-constrained in the way the fixed schedule does. The result is a more stable path — moving forward, or at least holding, on at least one axis per checkpoint. By 250k it reaches a Pareto-superior point on both IS and log(1/FID), with the gain seeming to come from sidestepping the late-stage regression rather than from simply training longer.

Residual and Frequency Views

6.2 Residual, Frequency Analysis and Takeaway

The signed luminance residual (Figure 14) shows divergence at object edges and contour transitions — structural rather than a global brightness shift. Two frequency regimes stand out in the FFT view (Figure 15). At coarsest global scales, fixed iREPA carries somewhat more power through the mid-to-late phase. At intermediate scales capturing contours and fine detail, cosine progressively opens capacity as the projection weight decays — consistent with the Pareto interpretation that fixed alignment persistently anchors coarse structure at the cost of finer scales.

The early FID cost of cosine annealing is transient. By 250k it reaches a Pareto-superior outcome on FID, KID, and IS — with gains localised to intermediate spatial frequencies where fixed alignment appears to constrain capacity. Alignment is an allocation problem: the question is not whether to use it, but how its weight should evolve.

8. Conclusion & Path Forward

Across three stages, Inspect-REPA traces alignment from a validated baseline through early-stop ablations to a long-horizon cosine schedule. Stage 1 confirms that iREPA accelerates convergence and that earlier layers help the model reach a good regime faster. Stage 2 shows that projection pressure is front-loaded: removing it late (stop-75) costs almost nothing, while removing it too early causes measurable degradation — and schedule timing changes not just the amount of quality achieved, but the order in which fidelity and diversity are acquired. Stage 3 shows that cosine annealing avoids the late-stage Pareto backtracking observed in always-on iREPA, reaching a jointly superior outcome on FID, KID, and IS by 250k.

The main conclusion is not that less alignment is always better, but that alignment should be treated as a resource to allocate. The optimization trajectory through fidelity-diversity space is a more informative object than scalar endpoints alone — and the natural next step is moving from fixed schedules to adaptive policies that respond to the model's observed training state.