PaperlayerView exact revision on arXiv

Source-linked AI summary

A Reproducibility Study of Partial Residual Ablations in Pre-LN Transformers

Pratikkumar Babariya

arXiv:2608.14689v1cs.LG

TL;DR

The roles of attention and feed-forward residual pathways in Pre-LN transformers remain unclear when studied independently. This study performs controlled partial-residual ablations at 10M and 124M scales, finding robust attention-residual collapse and reproducible 10M recovery after removing the feed-forward residual.

  • Problem

    The independent roles of attention and feed-forward residual pathways in transformer architectures remain poorly understood.

  • Method

    The study compares four residual configurations across 10M and 124M Pre-LN GPT-style transformers using controlled implementations and deterministic multi-seed experiments.

  • Results

    AttnOnly reaches 1.580 ± 0.003 validation loss across eight 10M seeds, while FFNOnly collapses to 3.350 ± 0.002; the 124M pattern remains unconfirmed.

  • Takeaways & Limitations

    The 10M asymmetry is confirmed, while cross-position routing remains a working hypothesis and the 124M behavior requires further controlled testing.

  • Takeaways & Limitations

    Results depend on implementation and evaluation choices, with fixed validation batches required to eliminate evaluation noise as a confound.

Abstract

from arXiv · show

Residual connections are a fundamental component of transformer architectures, yet the roles of the attention and feed-forward residual pathways remain poorly understood when considered independently. This paper presents a reproducibility study of partial residual ablations in Pre-LN GPT-style transformers trained at two scales (10M and 124M parameters). I compare four architectural configurations by selectively removing the attention residual connection, the feed-forward residual connection, or both. Across all experiments, removing the attention residual (FFNOnly) consistently causes deterministic collapse to the No-Residual performance floor. In contrast, removing the feed-forward residual (AttnOnly) exhibits a reproducible recovery effect at 10M scale under a controlled 8-seed deterministic study, while its behavior at 124M remains unresolved because of substantial seed variance. During the investigation, I identified and corrected an experimental measurement confound in runtime gain scaling and document both the failed intermediate reproduction and the subsequent controlled replication. Based on the empirical results, I propose a cross-position routing hypothesis to explain the observed asymmetry while explicitly distinguishing confirmed findings from unresolved questions. To support reproducibility, I release the complete source code, experiment configurations, checkpoints, training logs, and all experimental results, including intermediate non-reproducing runs.

1 Introduction

This study examines how independently removing attention and feed-forward residual pathways affects Pre-LN transformers, motivated by residual skips’ roles in gradient flow and residual learning. Initial nanoGPT experiments suggested joint collapse, but a controlled reimplementation revealed an asymmetric 10M-scale outcome and unresolved cross-environment and scale-dependent behavior.

  • Motivation: Residual skips provide a gradient highway and let each layer learn residual transformations instead of full mappings.The residual Jacobian contains an identity term that guarantees gradient flow regardless of sublayer behavior.
  • Ablation design: Four configurations selectively retain or remove attention and FFN identity paths: Full Residual, AttnOnly, FFNOnly, and No Residual.The configurations are defined by setting the corresponding residual flags to False.
  • Initial observation: Initial nanoGPT experiments showed all three partial configurations reaching validation loss ≈3.35 at 10M and ≈7.4 at 124M, matching No Residual.This was treated as a starting observation because the collapse could reflect architecture or implementation and evaluation methodology.
  • Controlled replication: At 10M, a controlled 8-seed deterministic sweep found AttnOnly mean 1.580 ± 0.003 versus FFNOnly mean 3.350 ± 0.002, a 1.770 validation-loss-unit gap.No seed approached the collapse floor, while an intermediate 3-seed run produced 2.591 ± 0.544 and remains an unresolved cross-environment discrepancy.
  • Scale dependence: At 124M, FFNOnly/AttnOnly = 1.56× at seed 1337, while two additional AttnOnly seeds showed large variance.This scale-dependent behavior motivates a controlled multi-seed investigation rather than a definitive conclusion.

2 Background

Residual connections support both gradient propagation and the shared residual stream in transformers, while initialization scale is coupled to the identity path’s stabilizing role during early training.

  • Identity shortcuts support deep-network training by ensuring gradient magnitude at layer ℓ is at least that at layer ℓ+1.
  • In transformers, each sublayer reads from and writes to a shared residual-stream vector that accumulates information across blocks.Removing a residual connection disrupts this compositional structure, not merely the optimization landscape.
  • Standard initialization with std = 0.02 assumes the identity path dominates early training, linking initialization scale to residual connections.The falsification experiment tests whether adjusting initialization can rescue partial-residual configurations after this assumption breaks.

3 Experimental Setup

Experiments compare four residual-connection configurations at 10M and 124M scales using two implementations and fixed validation data in ResidualGPT runs. The datasets and tokenization differ by scale, while configuration hyperparameters remain constant except for active residual connections.

  • Experimental Setup: Two implementations support the experiments: nanoGPT for the initial joint-collapse observation and ResidualGPT for controlled evaluation with fixed validation batches.ResidualGPT uses weight tying, flash attention, and validation batches created once with seed 99,991.
  • Experimental Setup: The 10M experiments use TinyShakespeare with character-level tokenization, vocab=65, train=1.0M tokens, and val=111K.
  • Experimental Setup: The 124M experiments use OpenWebText with BPE tokenization via tiktoken, vocab=50,257, and ∼113M tokens from 100,000 documents.
  • Experimental Setup: Four configurations are tested: Full Residual, No Residual, AttnOnly, and FFNOnly, with all hyperparameters held constant except active residual connections.The configurations selectively retain or remove the attention and feed-forward residual skips.

4 The Initial Observation: Joint Collapse

At 10M scale, all three partial-residual configurations jointly collapse to the No-Residual performance floor, with validation loss ≈3.35 and a 2.27× degradation versus Full Residual. The same plateauing pattern and comparable degradation ratios appear at 124M, while internal gradient behavior differs sharply between partial and full residual models.

  • 10M joint collapse: Validation loss ≈3.35 marks symmetric collapse of all three partial configurations, a 2.27× degradation relative to Full Residual and indistinguishable from No Residual.This is the joint-collapse observation at 10M scale on TinyShakespeare with seed 1337.
  • Cross-scale replication: At 124M, degradation ratios of 2.13× to 2.18× remain consistent with the 10M ratio of 2.27×, confirming robust joint collapse across scale, dataset, and tokenization scheme.Full Residual descends to low loss, whereas all partial configurations plateau immediately after an initial drop at both scales.
  • Internal dynamics: Layer 0 gradient norm collapses to exactly 0.000 from step 300 onward for every partial configuration, while Full Residual maintains ≈0.114 throughout.The identity term provides a direct gradient path to early layers regardless of sublayer behavior, producing the gradient highway effect.
  • Internal dynamics: Despite identical validation losses, the four configurations exhibit qualitatively different internal behavior, including significantly divergent hidden-state norm growth ratios.The contrast between external performance and internal dynamics motivates examining residual pathways beyond validation loss alone.

5 Challenging the Observation

The study challenges the initial recovery observation by identifying runtime gain scaling and noisy validation as measurement confounds. Controlled diagnostics and reproduction results distinguish the original AttnOnly signal from subsequent non-recovery while confirming FFNOnly’s outcome.

  • Initialization sensitivity: Runtime gain scaling was a confound because AdamW compensated for constant multiplicative reparameterization within ≈200 steps, making tested gain values functionally equivalent.The initial g = 2.0 result appeared to recover AttnOnly to best val 1.54, but this was identified as invalid before reporting.
  • Mechanistic diagnostics: Layer 0 gradient norms collapsed to 0.000 at step 300 for all partial configurations, while Full Residual sustained ≈0.114.Representation drift was 14.0× for AttnOnly, 5.4× for FFNOnly, and 1.4× for Full Residual.
  • Reproduction results: The originally recorded AttnOnly value of 1.600 was followed by a non-recovering reproduction at 2.591, whereas FFNOnly reproduced closely at 3.349.The training-outcome comparison was reported for ResidualGPT rather than nanoGPT.
  • Evaluation methodology: Fixed validation batches were introduced to eliminate the high evaluation variance caused by randomly sampled validation batches in nanoGPT.ResidualGPT reused identical samples across all four configurations, with batches created at startup using seed 99,991.

6 The Central Observation: A Confirmed 10M Asymmetry and an Unresolved 124M Scaling Behavior

The 10M asymmetry is confirmed under controlled deterministic evaluation: AttnOnly avoids the collapse floor while FFNOnly reaches it. At 124M, AttnOnly remains suggestive but unresolved because substantial seed variance obscures the qualitative distinction.

  • 10M scale: At 10M, AttnOnly achieved 1.580 ± 0.003 mean val loss across 8 seeds, while FFNOnly reached 3.350 ± 0.002, a 1.770-unit gap.No AttnOnly seed approached the FFNOnly collapse floor under the controlled deterministic sweep.
  • 10M scale: The originally recorded 10M FFNOnly/AttnOnly ratio was 2.09×, although an intermediate reproduction failed to recover the original AttnOnly statistic.A subsequent forced-determinism reproduction showed substantial seed spread, including one seed at the collapse floor, before the controlled 8-seed sweep confirmed the effect.
  • 10M scale: FFNOnly collapse was essentially deterministic, with mean 3.350 ± 0.002 versus the NoResidual anchor at 3.356.FFNOnly and No Residual differed by only 0.007 loss units, within measurement noise.
  • 124M scale: At 124M, AttnOnly values were 4.899, 7.586, and 5.967 across three seeds, while FFNOnly reached 7.666 and reproduced closely.Seed 42 was close to the FFNOnly/No Residual floor, so the qualitative distinction is not clearly established at this scale.
  • 124M scale: The 124M AttnOnly direction remains suggestive but unconfirmed because high seed variance prevents a firm qualitative distinction from the collapse floor.The cross-scale summary treats the 10M direction as confirmed across 8 seeds but the 124M behavior as unresolved.

7 Mechanistic Hypothesis

The proposed cross-position routing hypothesis explains why AttnOnly recovers despite the same Layer 0 gradient starvation affecting FFNOnly, while a mixer test remains inconclusive because AttnOnly itself failed to reproduce in that environment.

  • Empirical asymmetry: Both AttnOnly and FFNOnly show identical Layer 0 gradient starvation, yet the deterministic 8-seed sweep confirms recovery for AttnOnly but not FFNOnly.Layer 0 gradient norm is 0.000 from step 300 onward; an intermediate 3-seed reproduction on different hardware did not recover the distinction.
  • Cross-position routing: The FFN’s pointwise transformation cannot restore missing information through cross-position interactions, explaining why FFNOnly approximates No Residual.The reported gaps are 0.007 at 10M and 0.023 at 124M.
  • Cross-position routing: Attention can learn αtt ≈1, providing implicit cross-position routing that acts as a learnable substitute for the unavailable FFN identity path.Attention aggregates representations from preceding positions using learned content similarity, unlike the pointwise FFN.
  • Hypothesis evidence: AttnOnly hidden norm grows 14.03× versus 5.37× for FFNOnly, while the originally recorded asymmetry decreases at larger scale but did not reproduce.The originally recorded gaps are +50.1% at 124M and +6.5% at 10M; the scale trend remains unresolved.
  • Hypothesis test: A position-blind MLP mixer converged to 2.487, but the result cannot confirm or falsify routing because this environment’s AttnOnly baseline was 2.591 ± 0.544.The originally recorded AttnOnly value was 1.580, while the FFNOnly floor was 3.349.

8 Discussion

The discussion confirms AttnOnly recovery at 10M under controlled determinism, while 124M behavior, the asymmetry’s mechanism, and an intermediate reproduction discrepancy remain unresolved. It also documents a corrected runtime-gain experiment, universal Layer 0 gradient starvation, and a cross-position routing hypothesis.

  • Confirmed and unresolved findings: AttnOnly achieved mean 1.580 ± 0.003 across eight deterministic 10M seeds, with no seed approaching collapse and a 1.770 val loss unit gap from FFNOnly.AttnOnly behavior at 124M remains unresolved.
  • Confirmed and unresolved findings: All partial configurations exhibited zero gradient norm from step 300 onward at both 10M and 124M scales in nanoGPT.The discussion identifies this Layer 0 gradient starvation as universal.
  • Reproducibility limitations: Fixed validation batches explain the nanoGPT/ResidualGPT AttnOnly discrepancy, although this explanation is not formally proven and implementation differences remain individually unablated.The proposed explanation is therefore a documented limitation rather than a demonstrated causal result.
  • Mechanistic interpretation: Cross-position routing is proposed to explain AttnOnly recovery but not FFNOnly, consistent with diagnostics yet unproven and untested in alternative architectures.The mechanism behind the asymmetry remains an unresolved question.
  • Methodological correction: The runtime-gain investigation corrected an optimizer-reparameterization confound after apparent recovery showed all gain values were equivalent, then rebuilt the experiment from scratch.The corrected result revealed a deeper reproducibility gap.
  • Reproduction investigation: Three non-deterministic seed-1337 runs produced 1.97, 1.94, and 2.06, while deterministic seeds 1337/42/123 produced mean 2.591 and std 0.544.Forced deterministic algorithms eliminated run-to-run variance but did not close the reproduction gap.

9 Conclusion

The study confirms a robust asymmetry at 10M: removing the attention residual causes deterministic collapse, whereas removing the feed-forward residual produces reproducible recovery under controlled conditions. The 124M AttnOnly pattern and the proposed cross-position routing mechanism remain unresolved.

  • Confirmed findings: FFNOnly causes complete, deterministic collapse across every seed and tested environment.The reported standard deviation was 0.0009.
  • Open questions: At 124M, three AttnOnly seeds produced mean 6.151 with standard deviation 1.105, leaving the pattern unresolved.A controlled multi-seed sweep remains the most direct next experiment.
  • Confirmed findings: At 10M, AttnOnly recovery is reproducible and seed-stable under controlled conditions, with performance of 1.600 ± 0.015.An intermediate reproduction on different hardware instead recorded 2.591 ± 0.544.
  • Open questions: Cross-position routing as a substitute identity path remains a working hypothesis whose falsifiable prediction could not be conclusively tested under the anomalous intermediate environment.The prediction should be rerun under the clean sweep conditions.
Loading 2608.14689v1…