Training paths shape circuit removal in small transformer models
Training Trajectories Determine Circuit Removability in Annealable Soft-Prior Transformers
Machine LearningNeural and Evolutionary Computing
Summary
Small Transformer models can use helpful hints during training to learn how to retrieve information, but it’s unclear if these tricks stay useful when removed later. The authors tested a method where the model slowly loses these hints during training, which helped keep its skills strong even without the hints. However, if the hints were just switched off suddenly or removed at the end, the performance dropped. This shows that how a model is trained over time matters as much as its final design.
What this means in practice
- •For machine learning engineers: Design training schedules that gradually remove helpful biases to maintain model retrieval performance after initial guidance is faded.
- •For natural language processing teams: Improve small Transformer models for tasks requiring memory recall by adjusting training trajectories rather than relying on final architecture tweaks alone.
Authors
Zonglin Yang, Ziming Zhao, Wei Tang, Xunyu Jiang, Yihong Liu, Tailin Chen, Zifu Yu, Jiayu Liu
Abstract
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.