Source-linked AI summary

Mean Mode Screaming: Mean--Variance Split Residuals for 1000-Layer Diffusion Transformers

Pengqi Lu

arXiv:2605.06169v1cs.LGcs.CV

TL;DR

Very deep DiTs can undergo silent mean-dominated collapse, triggered by Mean Mode Screaming and reinforced by suppressed attention-logit gradients after value homogenization. The paper audits this mechanism and introduces MV-Split Residuals, which remove collapse in matched 400-layer evaluation and remain trainable in a separate 1000-layer run.

  • Problem

    Ultra-deep DiTs exhibit a structural collapse state that homogenizes token representations and suppresses centered variation, limiting reliable depth scaling.

  • Method

    The paper decomposes residual-writer gradients into mean-coherent and centered components, analyzes Softmax-null-space suppression, and proposes MV-Split Residuals with separately gained centered updates and leaky mean replacement.

  • Results

    MV-Split removes collapse events in matched 400-layer evaluation and remains stably trainable in a separate 1000-layer run.

  • Takeaways & Limitations

    Separating centered and mean residual control provides the evaluated design with a stable frontier at 400 layers and trainability at 1000 layers.

  • Takeaways & Limitations

    The main comparison assumes zero-initialized residual writers; standard initialization reaches the same collapse progressively across depth rather than through the delayed MMS event.

Abstract

from arXiv · show

Scaling Diffusion Transformers (DiTs) to hundreds of layers introduces a structural vulnerability: networks can enter a silent, mean-dominated collapse state that homogenizes token representations and suppresses centered variation. Through mechanistic auditing, we isolate the trigger event of this collapse as Mean Mode Screaming (MMS). MMS can occur even when training appears stable, with a mean-coherent backward shock on residual writers that opens deep residual branches and drives the network into a mean-dominated state. We show this behavior is driven by an exact decomposition of these gradients into mean-coherent and centered components, compounded by the structural suppression of attention-logit gradients through the null space of the Softmax Jacobian once values homogenize. To address this, we propose Mean-Variance Split (MV-Split) Residuals, which combine a separately gained centered residual update with a leaky trunk-mean replacement. On a 400-layer single-stream DiT, MV-Split prevents the divergent collapse that crashes the un-stabilized baseline; it tracks close to the baseline's pre-crash trajectory while remaining substantially better than token-isotropic gating methods such as LayerScale across the full schedule. Finally, we present a 1000-layer DiT as a scale-validation run at boundary scales, establishing that the architecture remains stably trainable at extreme depth.

1 Introduction

Ultra-deep DiTs can enter a mean-dominated collapse state after apparently stable training, with token homogenization and suppressed centered variation. The paper identifies MMS as the abrupt gradient-driven entry event and proposes MV-Split Residuals to preserve centered updates while controlling the mean path.

  • Ultra-deep DiTs can diverge after thousands of stable steps without NaNs or obvious forward saturation.
  • Mean Mode Screaming denotes the abrupt entry event marked by a mean-coherent gradient spike, residual branch opening, and Q/K gradient suppression.
  • Mean-coherent and centered gradients expose a geometric asymmetry: aligned tokens can produce coherent mean updates while centered propagation becomes contractive.
  • Isotropic residual gates stabilize training but also damp centered signals responsible for spatially varying feature learning.
  • MV-Split combines a separately gained centered residual update with a leaky trunk-mean replacement to stabilize training without equally shrinking centered updates.
  • In matched 400-layer evaluation, MV-Split removes collapse events and converges faster than LayerScale, while a separate 1000-layer run remains stably trainable.

2 Preliminaries

The paper uses a stripped-down single-stream DiT and Rectified Flow training setup to study deep residual propagation and collapse dynamics. Its main runs use zero-initialized residual writers, while standard initialization provides a contrasting collapse trajectory.

  • The backbone is a single-stream Post-Norm residual chain designed to keep deep residual propagation central to signal and gradient transport.
  • Residual writers WO and W2 are zero-initialized in the main runs, while internal branch parameters retain standard initialization.
  • Standard initialization does not avoid mean-dominated collapse; instead, collapse appears progressively across depth rather than through the delayed MMS writer-opening spike.
  • The main training objective is Rectified Flow, using a linear interpolation between VAE latents and Gaussian noise.

3 Failure Dynamics: Mean-Dominated Collapse

The failure sequence begins with a mean-coherent gradient shock and proceeds through residual branch opening to mean-dominated token homogenization. Attention preserves pure means but can contract centered variation, leaving residual branches responsible for replenishing it.

  • The audited divergence sequence is a mean-coherent gradient shock, residual branch opening, mean-dominated forward collapse, and Q/K gradient suppression.
  • Row-stochastic attention preserves the sequence-mean component while centered variation is governed by the projected operator PAP.
  • Centered variation can leak into the output mean through the term JAPX for general inputs.
  • When ∥PAP∥2 is below one, attention is contractive on the centered subspace, so residual branches must replenish centered variation.
  • In the 400-layer baseline, the mean-coherent spike coincides with Q/K gradient collapse, branch opening, exploding ρT, and token homogenization across depth.

4 Mechanism

The mechanism combines exact mean/centered gradient decomposition with coherent accumulation under token alignment. Once values homogenize, the Softmax null space suppresses Q/K learning while residual-writer gradients remain active, locking in collapse.

  • Gradient decomposition: Token-wise linear-map gradients decompose exactly into mean-coherent and centered terms because cross-terms vanish under sequence summation.
  • Gradient decomposition: The mean-coherent term scales as T∥¯δ∥∥¯y∥ and becomes O(T) when sequence means no longer cancel.
  • Alignment amplification: Under token and adjoint alignment, signed off-diagonal terms stop canceling and the gradient enters its O(T) coherent-amplification regime.
  • Alignment amplification: The alignment amplification law distinguishes heterogeneous tokens, where A≈1, from aligned deep-layer representations that approach κ→1.
  • Q/K gradient extinction: When residual streams become mean-dominated, value homogenization makes the Softmax Jacobian eliminate the constant attention-weight-gradient component.
  • Q/K gradient extinction: With approximate value homogeneity, Q/K learning is strongly suppressed while residual-writer gradients remain nonzero because they are not removed by the Softmax null space.

5 Method: MV-Split Residuals

MV-Split Residuals route mean and centered token components through separately gained updates, damping mean accumulation while preserving centered variation. This mode-selective design contrasts with isotropic residual gates that suppress both components together.

  • Design: MV-Split decouples the rank-one mean-coherent gradient update from the centered residual update using orthogonal token-space projectors.The method replaces the standard Post-Norm merge with a subspace-routed residual merge.
  • Forward dynamics: The centered subspace receives a standard residual update with gain β, while the mean subspace uses a leaky trunk-mean replacement.For 0 < α_d ≤1, each layer contracts the trunk mean by 1 −α_d before adding a fresh correction.
  • Backward dynamics: Centered and mean-coherent gradients receive independent gains, so a small α damps mean accumulation without proportionally shrinking local centered branch gradients.The same small mean gain also shrinks the mean-coherent writer component of the gradient.
  • Comparison: LayerScale and ReZero apply token-isotropic residual gains that jointly suppress mean-coherent and centered updates.This shared suppression stabilizes training but can slow convergence by damping centered signals needed for token-varying feature learning.

6 Experiments

Experiments identify signed-cancellation loss as the writer-gradient trigger for MMS and evaluate MV-Split under stability-constrained quality comparisons. MV-Split avoids collapse in 400-layer runs, preserves a higher bounded centered-gradient regime than LayerScale, and remains trainable at 1000 layers.

  • 6.1 Testing the Alignment-Amplification Law: At t⋆=3400, active Attn_WO and FFN_W2 layers lie near the gradient-amplification saturation envelope as signed cancellation disappears.The shared behavior across attention and FFN writers supports a writer-interface explanation rather than an attention-specific one.
  • 6.1 Testing the Alignment-Amplification Law: A −1 ≈167 corresponds to approximately 13× writer-gradient norm amplification relative to the independent-token baseline.These measurements support MMS as the regime where the mean-coherent update approaches coherent O(T) scaling.
  • 6.2 MV-Split Shifts the Stability-Constrained Quality Frontier: MV-Split shifts the controlled 400-layer stability-constrained frontier by preserving early convergence speed while avoiding baseline collapse.Its advantage over LayerScale is already substantial by 20k–30k steps and persists through the added 40k/50k checkpoints.
  • 6.2 MV-Split Shifts the Stability-Constrained Quality Frontier: MV-Split operates in a higher bounded gradient band than LayerScale while avoiding the spikes of un-stabilized runs.The comparison treats a run as usable only if it remains non-divergent over the measured horizon.
  • 6.2 MV-Split Shifts the Stability-Constrained Quality Frontier: The separate 1000-layer MV-Split run remains stable over the measured horizon and serves as scale validation rather than a matched 400-layer frontier point.Its separate training and post-training pipeline prevents direct use as a matched comparison.
  • 6.3 Writer-Gradient Mode Decomposition: LayerScale bounds the mean-coherent writer component by shrinking the centered component as well, whereas MV-Split keeps the centered component in a higher stable band.The mode decomposition indicates that MV-Split stability is not explained by uniformly smaller gradients.
  • 6.3 Writer-Gradient Mode Decomposition: The paper defers infrastructure-level optimizations for ultra-deep training and uses the token mean as an implicit global timestep carrier.A linear probe reports near-perfect R2 for predicting t across depth, motivating gain-limiting rather than strictly projecting out the mean.

7 Related Work

Related work addresses deep-Transformer instability through normalization, residual scaling, connectivity, and analyses of attention or representation collapse. MMS is positioned as a distinct diagnostic centered on mean-coherent writer gradients and their interaction with token homogenization.

  • Diffusion Transformers: DiT and related token-based diffusion backbones establish depth, token count, long skips, multimodal mixing, and rectified-flow generation as relevant design dimensions.The paper situates its single-stream DiT within this broader diffusion-transformer lineage.
  • Training stability: Prior deep-Transformer stabilizers modify normalization placement, residual scaling, or residual connectivity, including ReZero, LayerScale, DeepNorm, and Keel.These approaches target training instability through architectural control of residual dynamics and gradient geometry.
  • Collapse and spikes: MMS differs from loss-spike, Jacobian, attention-growth, and representation-collapse studies by diagnosing a mean-coherent writer-gradient event.The paper connects this event to token homogenization and Q/K-gradient suppression while distinguishing its diagnostic object from related collapse analyses.

8 Conclusion

The paper identifies MMS as a mean-dominated collapse mechanism that limits depth scaling in DiTs and proposes MV-Split to address it. Matched 400-layer evaluations and a separate 1000-layer run support stable training with this design.

  • Mean Mode Screaming is the abrupt writer-gradient event accompanying entry into a mean-dominated collapse state in ultra-deep DiTs.The collapse homogenizes token representations and suppresses centered variation.
  • The collapse mechanism combines an O(T) mean-coherent writer update with insufficient maintenance of the centered subspace.The paper frames this as an imbalance between mean-coherent and centered paths.
  • MV-Split removes collapse events and achieves the best stable frontier among the evaluated matched-backbone 400-layer stabilizer protocols.The comparison is under the paper's stability constraint.
  • A separate 1000-layer run shows that the same MV-Split design remains trainable at that depth.This run is presented as scale validation at boundary scales.

B Standard Initialization Enters the Same Mean-Dominated State

Standard initialization reaches the same mean-dominated endpoint as zero-writer training, but with an early depth-wise collapse front rather than a delayed MMS spike. The shared failure is insufficient centered-subspace maintenance alongside mean-coherent writer dominance.

  • A 128-layer DiT with Gaussian-initialized residual writers enters a mean-dominated regime from the beginning of training.The loss quickly reaches a high plateau, unlike the delayed failure pattern of zero-writer runs.
  • High token similarity forms a depth-wise collapse front, while centered retention and branch-side centered replenishment remain small in deep layers.Nonzero RowDiv rules out literal collapse to identical attention rows.
  • In collapsed deep layers, the mean-coherent writer-gradient component Gmean dominates the centered component Gctr by several orders of magnitude for both WO and W2.The same endpoint appears in standard-init and zero-writer runs: suppressed token variation with mean-dominated writer updates.
  • The token-wise linear-map gradient decomposes exactly into a rank-1 mean-coherent component and a centered component because the cross-terms vanish.The decomposition is an algebraic identity, not an approximation.
  • When values homogenize, the attention-logit gradient vanishes through the Softmax Jacobian null space, while residual-writer gradients bypass that null space and remain mean-coherent.This selectively suppresses Q/K learning while leaving the output-writer pathway active.
  • MV-Split separately gains the centered path and contracts the trunk mean, interrupting the alignment-driven failure mechanism without tying centered replenishment to the mean gain.LayerScale instead applies a shared token-independent gain to both paths and leaves the carried trunk mean untouched.

F Step-Level Gradient Trace for Failure Attribution

The trace pipeline attributes failure by combining distributed gradient localization, cross-rank data/loss checks, and mean-versus-centered writer-gradient audits. It identifies a global update event whose escalation shifts activity toward residual writers, while branch-level control experiments show that both residual branches require MV-Split.

  • Top-K parameter-family grouping localizes large gradients but does not identify whether the responsible token-space mode is mean-coherent or centered.That distinction requires the separate writer-gradient decomposition audit.
  • The trace triggers on a global-norm threshold, ranks parameter families, audits cross-rank exclusions, and records the dominant family for mode analysis.The exclusions check loss agreement, output-gradient statistics, and NaN/Inf in stored parameters.
  • At the representative escalation step, residual output interfaces dominate the largest gradients, with Attn_WO carrying most top-K squared-norm mass and FFN_W2 also appearing.Earlier warning snapshots are mixed across embeddings, projections, FFN inputs, and residual outputs.
  • At the spike, writer updates are amplified in Gmean while Gctr shows no comparable increase.This connects parameter-family localization to the paper's mean/centered mechanism.
  • In the attention-only 1000-layer control, the training still spikes at Step 7415 and the largest gradients move to the unprotected FFN branch.FFN_W2 accounts for 14 of the top 15 contributors, while Attn_WO disappears from the top-K.
  • ∼0.81 to ∼2.12: per-rank loss rises uniformly across four steps while σ remains in [0.011, 0.023].The tight cross-rank agreement is consistent with a global update event rather than a single-rank or single-batch fault.
  • Attention-only residual control is insufficient; both attention and FFN residual branches require the mean/centered split.The result is branch-level rather than weight-level.

J Methods we try but failed to prevent MMS

The tested controls did not remove mean-dominated failure because they failed to jointly control token-mean and centered residual components. MV-Split instead applies subspace-specific gains and a forward leaky mean replacement, while several reported evaluations remain calibration-only or future work.

  • Failed interventions: None of the tested interventions removed the mean-dominated failure; some additionally degraded optimization.The controls targeted token means, attention mixing, attention-output gating, scalar gradient norms, and optimizer-side update geometry.
  • Failed interventions: Hard centering removes the token mean and degraded optimization while discarding useful global information, including image-level context and an implicit timestep signal.Attention-matrix modifications also do not protect the FFN branch or residual merge.
  • Failed interventions: Global gradient clipping reduces step length but cannot rotate a mean-coherent writer update into the centered subspace.When the mean component dominates, scalar clipping also suppresses the already-small centered update; MV-Split instead applies different gains to the two components.
  • Failed interventions: Muon reshapes singular values after token gradients are summed, preserving the direction of a mean-coherent update rather than implementing a token-space split.It also lacks MV-Split’s forward leaky mean replacement.
  • Evaluation scope: The 1000-layer text-conditioned evaluation is calibration-only, not a controlled comparison or claim of state-of-the-art performance.The checkpoint uses smaller and differently sourced data, a shorter schedule, and a simpler post-training pipeline.
  • Limitations and future work: Exact MMS onset prediction remains separate from architectural stabilization, and the mechanism’s attention-specific arguments do not directly transfer beyond Softmax attention.Testing convolutional, hybrid, and state-space diffusion backbones is identified as future work.

M More Visual Results

The appendix presents additional visual samples from the 1000-layer MV-Split DiT, including text-conditioned and class-labeled ImageNet-style generations under a fixed Euler-sampling setup.

  • Text-to-image samples: Additional uncurated samples demonstrate the breadth and fidelity of the 1000-layer MV-Split DiT across diverse semantic categories.Images are generated at 256 × 256 resolution with Euler sampling, 35 NFE steps, and classifier-free guidance scale w = 2.0.
  • Text-to-image samples: Text-conditioned samples use natural-language captions from the ImageNet-2012 validation set.The captions describe scene content in 10–25 words and are generated by a modern large language model.
  • Class-conditional samples: Class-labeled visual results cover animal, object, vehicle, clothing, building, and landscape categories.The displayed classes include alligator lizard, scorpion, jacamar, dogs, household objects, vehicles, and Alp.
Loading 2605.06169v1…