PaperlayerView exact revision on arXiv

Source-linked AI summary

Training Trajectories Determine Circuit Removability in Annealable Soft-Prior Transformers

Zonglin Yang, Ziming Zhao, Wei Tang, Xunyu Jiang, Yihong Liu, Tailin Chen, Zifu Yu, Jiayu Liu

arXiv:2609.10287v1cs.LGcs.NE

TL;DR

The paper asks whether soft positional priors leave behind functional retrieval circuits after removal, rather than remaining permanent dependencies. It tests gated priors under alternative training trajectories and finds that smooth fade-to-zero training preserves removability on small discrete retrieval tasks, with important task and length-generalization boundaries.

  • Problem

    The central problem is distinguishing internally implemented retrieval from continued dependence on an external positional bias after the prior is removed.

  • Method

    An annealable soft-prior Transformer varies learned, faded, clamped, or zeroed attention biases during training and evaluation, comparing trajectories with matched model and training conditions.

  • Results

    Smooth fade-to-zero training preserves prior-free retrieval on associative recall and Markov induction, whereas unforced, forced-zero, hard-switch, and post hoc paths fail to reproduce the same effect.

  • Takeaways & Limitations

    For small discrete retrieval circuits, removability depends on the training trajectory rather than only the final architecture.

  • Takeaways & Limitations

    The positive trajectory results are limited to synthetic discrete retrieval tasks with Transformers no larger than L=4 and d=128, while length extrapolation remains weak and schedule scans are limited.

Abstract

from arXiv · show

Soft positional priors can help small Transformers learn retrieval circuits, but it is unclear whether the resulting circuits remain functional once the prior is removed. We test this with an annealable soft-prior Transformer whose attention biases can be learned, faded, or zeroed during training and evaluation. On associative recall, unforced models perform well with the prior active ($0.772 \pm 0.020$) but collapse at zero gate ($0.095 \pm 0.009$). Smooth fade-to-zero training preserves high zero-gate accuracy ($0.734 \pm 0.028$), whereas forced-zero training, hard switching, and post hoc continuation fail to recover the same effect. The pattern also appears on Markov induction. Linear regression ICL provides a boundary case because zero-gate training can learn that task directly. Mechanistic traces show that circuit consolidation occurs after the gate reaches zero, even though the responsible heads vary across seeds. These results suggest that circuit removability in small discrete retrieval tasks depends on the training trajectory, not just the final architecture.

1 Introduction

The paper asks whether attention priors remain necessary after helping retrieval circuits form, and tests this by varying and removing a gated soft prior during training. Smooth fade-to-zero training preserves prior-free retrieval, while abrupt or post hoc removal does not reproduce the effect.

  • Motivation: Standard evaluation cannot distinguish an internal retrieval circuit from permanent reliance on an external positional bias.Removing the bias at evaluation tests which mechanism the model uses.
  • Approach: The annealable Transformer lets attention-head position and content biases be learned, clamped, or scheduled during training.The protocol distinguishes prior-active performance from prior-free circuit removability.
  • Results: Fade-to-zero training preserves high zero-gate accuracy on associative recall, whereas unforced, forced-zero, hard-switch, and post hoc paths do not reproduce the same effect.The model begins with the prior active, smoothly reduces it to zero, and continues training without the prior.
  • Results: The same removability pattern appears on Markov induction, while linear regression ICL forms a boundary case in which zero-gate training can learn the task directly.The discrete retrieval tasks therefore differ from linear regression ICL in whether prior-aided formation is needed.
  • Mechanism: Mechanistic traces show stable consolidation timing after the gate reaches zero despite variation in which heads carry retrieval across seeds.Schedule sweeps and three-seed traces support smooth removal while showing that head identity is not fixed.

2 Related Work

The paper relates soft-prior annealing to positional inductive biases, in-context learning circuits, continuation methods, and removable network structure. Its distinction is to keep the final architecture fixed while varying only the attention-prior gate trajectory and testing circuit survival after removal.

  • Position and attention priors: Prior work uses positional representations, adaptive position functions, and locality biases to shape Transformer attention or improve extrapolation.These approaches generally treat the prior as part of the final architecture.
  • In-context learning and circuits: Linear regression ICL and discrete retrieval tasks expose in-context learning algorithms and interpretable induction or successor circuits.Their circuit structure makes it possible to test whether a prior remains causally necessary.
  • Continuation paths and soft targets: The fade-to-zero schedule resembles continuation methods and curricula by changing the optimization path from an easier condition toward a harder one.Unlike distillation or data curricula, the auxiliary signal enters as a prior added to attention logits.
  • Removable and reparameterized structure: Methods such as pruning, layer dropping, and structural reparameterization study usefulness after components are removed or graphs are simplified.This paper instead holds the final architecture fixed and changes only the attention-prior trajectory.

3 Method

The method adds gated positional and content biases to attention and compares training trajectories that retain, schedule, or remove these priors. It frames removability as continuation along a parameter branch, with smooth gate changes intended to preserve circuit formation.

  • Annealable soft-prior attention: Each attention head receives gated positional and content bias terms before softmax.The positional prior uses clipped relative position, while the content prior uses learned token embeddings.
  • Gate modes: Gate modes include freely learned gates, fixed gates, and manually scheduled gates.Masked gate logits remain near initialization, so scheduled runs are evaluated at their final scheduled value.
  • Training trajectories: The main fade600→0 path holds g=0.5, linearly fades to zero, then trains with the prior absent.Hard-switch and post hoc controls use comparable endpoints or zero-gate budgets but alter the training trajectory.
  • Continuation view: Continuation is plausible when successive branch displacements remain within the optimizer’s basin, whereas hard switches create a larger perturbation.The method motivates smoothness through local minimizer branches and curvature-sensitive displacement.
  • Continuation view: Post hoc continuation can remain in a prior-dependent basin because zero-gate gradients may poorly project onto the prior-free circuit subspace.The hard-switch and post hoc experiments separately test these proposed failure modes.
  • Tasks and models: The study evaluates associative recall, Markov induction, and linear regression ICL with default and larger Transformer backbones.The trajectory claim targets the two discrete retrieval tasks, while linear regression ICL serves as a boundary case.

4 Experimental Setup

The experiments compare matched positional and attention baselines across retrieval and regression tasks, using controlled training budgets and schedule-related evaluations. Figures emphasize static performance and difficulty variation in Markov induction.

  • Baselines and metrics: Static comparisons include vanilla, NoPE, RoPE, ALiBi, learned RPE, DAPE, CoPE, FIRE, Kerple, GPSA, and soft-prior variants.All methods use matched backbone dimensions, optimizer, and training budget.
  • Baselines and metrics: Accuracy is reported for associative recall and Markov induction, while mean squared error is reported for linear regression ICL.Markov induction varies repeat probability across p ∈ {0.2, 0.5, 0.8}.
  • Evaluation protocol: Main trajectory results average five seeds for fade-window and causal controls, and three seeds for capacity, shape, task-boundary, and mechanism experiments.The archive also contains schedule-shape, linear-trajectory, and mechanism-trace runs.
  • Figures: Figure 2 compares static baselines across position and attention variants before removability testing.Figure 3 compares Markov induction accuracy across repeat probabilities, with the hardest setting separating methods.

5 Static Baselines

Soft-prior baselines perform strongly enough to support removability tests, with positional information providing the useful associative-recall signal. Markov induction shows a similar advantage, especially in its hardest setting.

  • Associative recall: 0.794 ± 0.022 associative-recall accuracy makes soft_bias the highest-mean result among thirteen baselines.soft_pos_only is close at 0.783, while soft_content_only is near chance at 0.067.
  • Associative recall: Position, rather than content, provides the useful soft-prior signal on associative recall.The content-only variant is near chance despite the strong soft_bias result.
  • Markov induction: Soft_bias is among the strongest Markov-induction methods, particularly at repeat probability p=0.2.These static results establish a meaningful baseline for testing whether the trained model still needs the prior.
  • Removability test: The causal-intervention figure concerns whether positional and content gates remain necessary after training.Its reported comparison is developed in the gate-intervention section.

6 Gate Interventions

Causal gate interventions show that the normally trained associative-recall model relies on its positional prior at inference, while removing the content prior is nearly neutral.

  • Causal intervention: The intervention evaluates content, position, or both gates at inference after a normally trained soft-prior model.The figure identifies positional dependence while treating the content gate as inert on this task.
  • Causal intervention: 0.791 versus 0.794 normal accuracy shows that removing the content gate is nearly neutral.The content gate is retained as an architectural control.
  • Causal intervention: 0.073 accuracy after position-gate removal demonstrates a collapse from the normally trained model.Removing both gates gives 0.081 accuracy.
  • Interpretation: Under unforced training, the position prior becomes part of the computation rather than merely improving normal accuracy.Normal accuracy alone therefore cannot establish circuit removability.

7 Training Path Determines Removability

Associative-recall removability depends on how the positional prior is removed during training: smooth fading preserves zero-gate performance, whereas abrupt or post hoc removal does not.

  • 7 Training Path Determines Removability: The fade600→0 path preserves high zero-gate accuracy in both base and larger backbones, while unforced and forced_zero training fail.Scheduled runs evaluate at the gate value reached at the end of training.
  • 7 Training Path Determines Removability: Zero-gate accuracy remains stable across fade starts 400–800, with values of 0.748, 0.745, 0.734, 0.751, and 0.754.Unforced and forced-zero baselines remain near 0.09.
  • 7 Training Path Determines Removability: Linear, cosine, and convex schedules achieve 0.749–0.764 zero-gate accuracy, whereas a two-step deletion reaches only 0.212.A concave schedule is unstable across seeds.
  • 7 Training Path Determines Removability: Abrupt switching reaches only 0.293 ± 0.128, and post hoc continuation from an unforced checkpoint reaches 0.178 ± 0.035.These controls match the zero-gate training budget but do not reproduce smooth fading.
  • 7 Training Path Determines Removability: The results locate removability in the training path: the prior must shape circuit formation while the gate decreases gradually enough for internal attention to take over.This interpretation follows the matched-path comparison and causal controls.

8 Boundary Across Tasks

The trajectory effect transfers to Markov induction but is task-dependent: linear regression ICL can be learned directly at zero gate, and removability does not improve long-length extrapolation.

  • 8 Boundary Across Tasks: Fade600→0 retains 0.830, 0.920, and 0.969 zero-gate accuracy at repeat probabilities p=0.2, 0.5, and 0.8.Unforced training collapses more strongly, especially at p=0.2.
  • 8 Boundary Across Tasks: Forced_zero reaches 0.0722 loss on linear regression ICL, while fade600→0 reaches 0.0646 and unforced zero-gate loss degrades to 0.1296.This makes linear regression a boundary case where zero-gate training can learn directly.
  • 8 Boundary Across Tasks: The fade advantage does not hold for every ICL task because linear regression can form a solution without the prior.The trajectory claim is therefore strongest for discrete retrieval circuits.
  • 8 Boundary Across Tasks: At T=120, soft_bias reaches only 0.263 accuracy, so removability does not imply improved length extrapolation.The paper treats this as a limitation of the recovered prior-free circuit.

9 Mechanistic Diagnostics

Mechanistic traces show delayed consolidation after the gate reaches zero, while the heads carrying the relevant signals vary across seeds.

  • 9 Mechanistic Diagnostics: Ablation importance first crosses threshold at step 1,200 and successor-value mass at step 1,500 in all three seeds, after fading ends around step 1,050.The traces measure both successor routing and zero-gate ablation importance.
  • 9 Mechanistic Diagnostics: The head carrying successor mass changes by seed, but consolidation timing remains stable across seeds.This pattern is consistent with head permutation symmetry rather than a universal head identity.
  • 9 Mechanistic Diagnostics: For a representative seed, L0H0 becomes critical after step 1,200, while L1H3 carries the largest successor-value mass after step 1,500.The figure identifies timing rather than a universal responsible head.

10 Limitations and Conclusion

The evidence supports trajectory-dependent removability for small discrete retrieval circuits, but its scope is limited by task, model scale, and weak length extrapolation.

  • 10 Limitations and Conclusion: Positive trajectory results come from synthetic discrete retrieval tasks using Transformers no larger than L=4 and d=128.The paper does not establish the same pattern beyond this setting.
  • 10 Limitations and Conclusion: Length extrapolation remains weak, and the schedule-shape scan supports smoothness over abrupt deletion without covering all possible annealing paths.The content gate was mostly inert, so the evidence concerns positional rather than content priors.
  • 10 Limitations and Conclusion: Linear regression ICL is a boundary case because zero-gate training can learn the task directly, unlike the discrete retrieval setting.This limits the claim that gradual fading is universally necessary for removability.
  • 10 Limitations and Conclusion: Mechanistic traces place consolidation after prior removal even though the responsible heads vary by seed.The conclusion concerns timing of consolidation rather than a fixed head identity.
Loading 2609.10287v1…