Elamur Research⌕
☰
home / papers / 2025 / Tracr-Injection: Distilling Algorithms into Pre-trained Language Models

Tracr-Injection: Distilling Algorithms into Pre-trained Language Models

Paper · 2025 · Tomás Vergara-Browne and Álvaro Soto

Summary

Transformers can in principle run many algorithms that are expressible in RASP, but pre-trained language models rarely learn them from ordinary unsupervised text. This paper introduces tracr-injection, which closes part of that gap by transferring compiled RASP programs into an existing model. Each algorithm is written in RASP and compiled with tracr into a small causal transformer. GPT-2-large is then fine-tuned with three objectives: task cross-entropy, a layer-wise cosine loss that aligns a linear projection of its residual stream with the compiled model's activations, and a KL term on Fineweb text to limit forgetting. An extra intervention swaps the out-of-subspace part of the hidden state between batch examples, which pushes the model to depend on the injected subspace. The authors test three tasks: Shuffle-Dyck recognition, counting characters, and integer addition. The projected subspace can be decoded directly into RASP variables. Out-of-distribution robustness improves clearly for Shuffle-Dyck but not consistently for the other tasks. The approach matters because it plants known "ground truth" mechanisms, which could serve as benchmarks for interpretability and circuit-discovery methods.

Key points

  • Method: The training objective is L = α·L_CE + β·L_ALG + γ·L_KL, with all three weights set to 1. L_ALG averages (1 − cosine similarity) between the mapped GPT-2 hidden states and the tracr-compiled hidden states across the compiled model's layers.
  • Baseline: The comparison model, "Finetuning + KL", drops the alignment loss and the batch-roll intervention. Both approaches reach high in-distribution accuracy and preserve general abilities, measured with LAMBADA and perplexity on Tiny-Shakespeare.
  • Shuffle-Dyck o.o.d. results: Tracr-injection stays at about 96.6–99.1% on almost-balanced strings, 3× longer inputs, and an extra parenthesis family. The baseline drops to roughly 75–89% on the same tests.
  • Addition and counting o.o.d. results: There is no consistent advantage. For example, both methods score 0.0% on decimals, and on X-only counting the baseline scores 33.3% versus 21.9% for tracr-injection. This suggests the benefit depends on the specific algorithm.
  • Resources and limitations: The code is released at github.com/tvergara/tracr-injection. Limitations are that only three simple tasks were tested, only GPT-2 was used because of compute constraints, and generalization fails on inputs far from the training distribution.

Our summary of the paper, in our own words — not its text. The work belongs to its authors; read it at the links above.

 Fields, topics and concepts

📚 Shelves 4The field of science this work belongs to (OpenAlex fields).Computer Science · 19,503 worksArtificial Intelligence · 18,805 worksArts and Humanities · 17,959 worksLanguage and Linguistics · 17,789 works

🧭 Topics 1Research subjects and languages the work is about.Language models & pretraining · 4,145 works

Cite

Cite as Tomás Vergara-Browne and Álvaro Soto. 2025. Tracr-Injection: Distilling Algorithms into Pre-trained Language Models. arXiv:2505.10719.

The paper itself, not this page. BibTeX from the library's record of the work. DOI 10.48550/arxiv.2505.10719 · arXiv 2505.10719

@misc{vergarabrowne2025tracr,
  title = {{Tracr-Injection: Distilling Algorithms into Pre-trained Language Models}},
  author = {Tomás Vergara-Browne and Álvaro Soto},
  year = {2025},
  doi = {10.48550/arxiv.2505.10719},
  eprint = {2505.10719},
  archivePrefix = {arXiv},
  url = {https://doi.org/10.48550/arxiv.2505.10719}
}

References

36 references, 13 of them with a card of ours. ↗ leads to the work, «find» searches for one we have no id for yet.

  1. Yoshua Bengio, Sören Mindermann, Daniel Privitera, Tamay Besiroglu, Rishi Bommasani, Stephen Casper, Yejin Choi, Philip Fox, Ben Garfinkel, Danielle Goldfarb, et al. 2025. International ai safety report. arXiv preprint arXiv:2501.17805. find
  2. Adithya Bhaskar, Alexander Wettig, Dan Friedman, and Danqi Chen. 2024. Finding transformer circuits with edge pruning. Advances in Neural Information Processing Systems, 37:18506–18534. find
  3. Bill Dolan and Chris Brockett. 2005. Automatically constructing a corpus of sentential paraphrases. In Third international workshop on paraphrasing (IWP2005).
  4. Javid Ebrahimi, Dhruv Gelda, and Wei Zhang. 2020. How can self-attention networks recognize dyck-n languages? arXiv preprint arXiv:2010.04303.
  5. Nelson Elhage, Neel Nanda, Catherine Olsson, Tom Henighan, Nicholas Joseph, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Nova DasSarma, Dawn Drain, Deep Ganguli, Zac Hatfield-Dodds, Danny Hernandez, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Ja… find
  6. Javier Ferrando, Gabriele Sarti, Arianna Bisazza, and Marta R Costa-jussà. 2024. A primer on the inner workings of transformer-based language models. arXiv preprint arXiv:2405.00208.
  7. Atticus Geiger, Duligur Ibeling, Amir Zur, Maheep Chaudhary, Sonakshi Chauhan, Jing Huang, Aryaman Arora, Zhengxuan Wu, Noah Goodman, Christopher Potts, et al. 2023. Causal abstraction: A theoretical foundation for mechanistic interpretability. arXiv preprint arXiv:2301.04709. find
  8. Atticus Geiger, Hanson Lu, Thomas Icard, and Christopher Potts. 2021. Causal abstractions of neural networks. Advances in Neural Information Processing Systems, 34:9574–9586. find
  9. Atticus Geiger, Zhengxuan Wu, Hanson Lu, Josh Rozner, Elisa Kreiss, Thomas Icard, Noah Goodman, and Christopher Potts. 2022. Inducing causal structure for interpretable neural networks. In International Conference on Machine Learning, pages 7324–7338. PMLR. find
  10. Geoffrey Hinton. 2015. Distilling the knowledge in a neural network. arXiv preprint arXiv:1503.02531.
  11. Andrej Karpathy Blog. 2015. The unreasonable effectiveness of recurrent neural networks. URL: http://karpathy.github. io/2015/05/21/rnn-effectiveness/dated May, 21:31. ↗
  12. Xiaoqi Jiao, Yichun Yin, Lifeng Shang, Xin Jiang, Xiao Chen, Linlin Li, Fang Wang, and Qun Liu. 2019. Tinybert: Distilling bert for natural language understanding. arXiv preprint arXiv:1909.10351.
  13. Stephen Casper, Carson Ezell, Charlotte Siegmann, Noam Kolt, Taylor Lynn Curtis, Benjamin Bucknall, Andreas Haupt, Kevin Wei, Jérémy Scheurer, Marius Hobbhahn, et al. 2024. Black-box access is insufficient for rigorous ai audits. In Proceedings of the 2024 ACM Conference on Fairness, Accountability,… find
  14. James Kirkpatrick, Razvan Pascanu, Neil Rabinowitz, Joel Veness, Guillaume Desjardins, Andrei A Rusu, Kieran Milan, John Quan, Tiago Ramalho, Agnieszka Grabska-Barwinska, et al. 2017. Overcoming catastrophic forgetting in neural networks. Proceedings of the national academy of sciences, 114(13):3521… find
  15. Arthur Conmy, Augustine Mavor-Parker, Aengus Lynch, Stefan Heimersheim, and Adrià Garriga-Alonso. 2023. Towards automated circuit discovery for mechanistic interpretability. Advances in Neural Information Processing Systems, 36:16318–16352. find
  16. David Lindner, János Kramár, Sebastian Farquhar, Matthew Rahtz, Tom McGrath, and Vladimir Mikulik. 2024. Tracr: Compiled transformers as a laboratory for interpretability. Advances in Neural Information Processing Systems, 36. find
  17. Samuel Marks, Can Rager, Eric J Michaud, Yonatan Belinkov, David Bau, and Aaron Mueller. 2024. Sparse feature circuits: Discovering and editing interpretable causal graphs in language models. arXiv preprint arXiv:2403.19647. find
  18. William Merrill and Ashish Sabharwal. 2023. A logic for expressing log-precision transformers. Advances in neural information processing systems, 36:52453– 52463. find
  19. Marius Mosbach, Vagrant Gautam, Tomás Vergara-Browne, Dietrich Klakow, and Mor Geva. 2024. From insights to actions: The impact of interpretability and analysis research on nlp. arXiv preprint arXiv:2406.12618. find
  20. Denis Paperno, Germán Kruszewski, Angeliki Lazaridou, Quan Ngoc Pham, Raffaella Bernardi, Sandro Pezzelle, Marco Baroni, Gemma Boleda, and Raquel Fernández. 2016. The lambada dataset: Word prediction requiring a broad discourse context. arXiv preprint arXiv:1606.06031.
  21. Guilherme Penedo, Hynek Kydlíček, Anton Lozhkov, Margaret Mitchell, Colin Raffel, Leandro Von Werra, Thomas Wolf, et al. 2024. The fineweb datasets: Decanting the web for the finest text data at scale. arXiv preprint arXiv:2406.17557.
  22. Jing Qian, Hong Wang, Zekun Li, Shiyang Li, and Xifeng Yan. 2022. Limitations of language models in arithmetic and symbolic induction. arXiv preprint arXiv:2208.05051. find
  23. Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. 2019. Language models are unsupervised multitask learners. OpenAI blog, 1(8):9. find
  24. Adriana Romero, Nicolas Ballas, Samira Ebrahimi Kahou, Antoine Chassang, Carlo Gatta, and Yoshua Bengio. 2014. Fitnets: Hints for thin deep nets. arXiv preprint arXiv:1412.6550. find
  25. Peter Shaw, James Cohan, Jacob Eisenstein, Kenton Lee, Jonathan Berant, and Kristina Toutanova. Alta: Compiler-based analysis of transformers. In The First Workshop on System-2 Reasoning at Scale, NeurIPS’24. find
  26. Paul Smolensky, Roland Fernandez, Zhenghao Herbert Zhou, Mattia Opper, and Jianfeng Gao. 2024. Mechanisms of symbol processing for in-context learning in transformer networks. arXiv preprint arXiv:2410.17498. — [arXiv:2410.17498](https://arxiv.org/abs/2410.17498) find
  27. Richard Socher, Alex Perelygin, Jean Wu, Jason Chuang, Christopher D Manning, Andrew Y Ng, and Christopher Potts. 2013. Recursive deep models for semantic compositionality over a sentiment treebank. In Proceedings of the 2013 conference on empirical methods in natural language processing, pages 1631…
  28. Aaquib Syed, Can Rager, and Arthur Conmy. Attribution patching outperforms automated circuit discovery. In The 7th BlackboxNLP Workshop.
  29. A Vaswani. 2017. Attention is all you need. Advances in Neural Information Processing Systems.
  30. Alex Wang. 2018. Glue: A multi-task benchmark and analysis platform for natural language understanding. arXiv preprint arXiv:1804.07461.
  31. Gail Weiss, Yoav Goldberg, and Eran Yahav. 2021. Thinking like transformers. In International Conference on Machine Learning, pages 11080–11090. PMLR. find
  32. Yixuan Weng, Minjun Zhu, Fei Xia, Bin Li, Shizhu He, Kang Liu, and Jun Zhao. 2023. Mastering symbolic operations: Augmenting language models with compiled neural networks. arXiv preprint arXiv:2304.01665. — [arXiv:2304.01665](https://arxiv.org/abs/2304.01665) find
  33. Andy Yang and David Chiang. 2024. Counting like transformers: Compiling temporal counting logic into softmax transformers. arXiv preprint arXiv:2404.04393. — [arXiv:2404.04393](https://arxiv.org/abs/2404.04393) find
  34. Gilad Yehudai, Haim Kaplan, Asma Ghandeharioun, Mor Geva, and Amir Globerson. 2024. When can transformers count to n? arXiv preprint arXiv:2407.15160. — [arXiv:2407.15160](https://arxiv.org/abs/2407.15160) find
  35. Zhehao Zhang, Jiaao Chen, and Diyi Yang. 2024. Darg: Dynamic evaluation of large language models via adaptive reasoning graph. arXiv preprint arXiv:2406.17271.
  36. Hattie Zhou, Arwen Bradley, Etai Littwin, Noam Razin, Omid Saremi, Josh Susskind, Samy Bengio, and Preetum Nakkiran. 2023. What algorithms can transformers learn? a study in length generalization. arXiv preprint arXiv:2310.16028.

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/vergara-browne-2025-tracr-injection-distilling-algorithms/