# Jet Expansions of Residual Computation

arXiv: https://arxiv.org/abs/2410.06024 · HTML: https://mlsanity.com/paper/2410.06024/

## Key facts

- Cited by 4 papers in the map (data through 14 Oct 2026).
- Builds on 39 of its 57 references that are in the map.
- First posted on arXiv 8 Oct 2024.

## Builds on (most cited first, 30 of 39)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [LLaMA: Open and Efficient Foundation Language Models](https://mlsanity.com/paper/2302.13971/) (2023)
- [Llama 2: Open Foundation and Fine-Tuned Chat Models](https://mlsanity.com/paper/2307.09288/) (2023)
- [A Unified Approach to Interpreting Model Predictions](https://mlsanity.com/paper/1705.07874/) (2017)
- [Finetuned Language Models Are Zero-Shot Learners](https://mlsanity.com/paper/2109.01652/) (2021)
- [Training a Helpful and Harmless Assistant with Reinforcement Learning from Human Feedback](https://mlsanity.com/paper/2204.05862/) (2022)
- [Axiomatic Attribution for Deep Networks](https://mlsanity.com/paper/1703.01365/) (2017)
- [Locating and Editing Factual Associations in GPT](https://mlsanity.com/paper/2202.05262/) (2022)
- [Transformer Feed-Forward Layers Are Key-Value Memories](https://mlsanity.com/paper/2012.14913/) (2020)
- [RealToxicityPrompts: Evaluating Neural Toxic Degeneration in Language Models](https://mlsanity.com/paper/2009.11462/) (2020)
- [Interpretability in the Wild: a Circuit for Indirect Object Identification in GPT-2 small](https://mlsanity.com/paper/2211.00593/) (2022)
- [Toy Models of Superposition](https://mlsanity.com/paper/2209.10652/) (2022)
- [Towards Automated Circuit Discovery for Mechanistic Interpretability](https://mlsanity.com/paper/2304.14997/) (2023)
- [Eliciting Latent Predictions from Transformers with the Tuned Lens](https://mlsanity.com/paper/2303.08112/) (2023)
- [OLMo: Accelerating the Science of Language Models](https://mlsanity.com/paper/2402.00838/) (2024)
- [Transformer Feed-Forward Layers Build Predictions by Promoting Concepts in the Vocabulary Space](https://mlsanity.com/paper/2203.14680/) (2022)
- [ToxiGen: A Large-Scale Machine-Generated Dataset for Adversarial and Implicit Hate Speech Detection](https://mlsanity.com/paper/2203.09509/) (2022)
- [Mechanistic Interpretability for AI Safety -- A Review](https://mlsanity.com/paper/2404.14082/) (2024)
- [Highway Networks](https://mlsanity.com/paper/1505.00387/) (2015)
- [Depth-Adaptive Transformer](https://mlsanity.com/paper/1910.10073/) (2019)
- [Localizing Model Behavior with Path Patching](https://mlsanity.com/paper/2304.05969/) (2023)
- [Infini-gram: Scaling Unbounded n-gram Language Models to a Trillion Tokens](https://mlsanity.com/paper/2401.17377/) (2024)
- [Language Models Implement Simple Word2Vec-style Vector Arithmetic](https://mlsanity.com/paper/2305.16130/) (2023)
- [An Interpretability Illusion for BERT](https://mlsanity.com/paper/2104.07143/) (2021)
- [Information Flow Routes: Automatically Interpreting Language Models at Scale](https://mlsanity.com/paper/2403.00824/) (2024)
- [Neurons in Large Language Models: Dead, N-gram, Positional](https://mlsanity.com/paper/2309.04827/) (2023)
- [A Primer on the Inner Workings of Transformer-based Language Models](https://mlsanity.com/paper/2405.00208/) (2024)
- [Residual Connections Encourage Iterative Inference](https://mlsanity.com/paper/1710.04773/) (2017)
- [Conditional Deep Learning for Energy-Efficient and Enhanced Pattern Recognition](https://mlsanity.com/paper/1509.08971/) (2015)

## Built on by (most cited first, 4 of 4)

- [LLM Unlearning via Neural Activation Redirection](https://mlsanity.com/paper/2502.07218/) (2025)
- [Scaling Embedding Layers in Language Models](https://mlsanity.com/paper/2502.01637/) (2025)
- [Evolution of Concepts in Language Model Pre-Training](https://mlsanity.com/paper/2509.17196/) (2025)
- [DDEQs: Distributional Deep Equilibrium Models through Wasserstein Gradient Flows](https://mlsanity.com/paper/2503.01140/) (2025)

## Abstract (arXiv)

We introduce a framework for expanding residual computational graphs using jets, operators that generalize truncated Taylor series. Our method provides a systematic approach to disentangle contributions of different computational paths to model predictions. In contrast to existing techniques such as distillation, probing, or early decoding, our expansions rely solely on the model itself and requires no data, training, or sampling from the model. We demonstrate how our framework grounds and subsumes logit lens, reveals a (super-)exponential path structure in the recursive residual depth and opens up several applications. These include sketching a transformer large language model with n-gram statistics extracted from its computations, and indexing the models' levels of toxicity knowledge. Our approach enables data-free analysis of residual computation for model interpretability, development, and evaluation.