Decomposing LLM Computation with Jets

Overview

This project introduces a systematic framework called JET Expansion, which uses jet operators to expand residual computation into interpretable paths. Unlike other methods, JET Expansion requires no data or additional training, offering a direct way to analyze the computational behavior of models like transformers, where the residual computation plays a vital role.

• Provides insights into the inner workings of large language models with jet lens.
• Reformat language models as symbolic databases with jet n-grams for model diffing
• Enables data-free analysis of component roles, pretraining dynamic, and finetuning effects with applications like toxicity analysis.

Visualizations

Technical Details

The method relies on recursively applying jet expansions, which generalize Taylor series, providing a way to disentangle different computational paths in deep learning models like transformers. It does not require additional data, training, or sampling. Our paper offers more details about the method.

Use Cases

Below are examples of how JET Expansion can be used to gain insights into popular models:

• Model Diff Diff language models using jet n-grams.
• Trace Pretraining Dynamic Trace the evolution during pretraining.
• Data-Free Toxicity Indexing Measure the toxicity embedded in models before and after RLHF fine-tuning.
• Jet Lens for Understanding Inner Workings

BibTeX

@inproceedings{chen2026decomposing,
  title={Decomposing LLM Computation with Jets},
  author={Chen, Yihong and Xu, Xiangxiang and Stenetorp, Pontus and Riedel, Sebastian and Franceschi, Luca},
  booktitle={ICLR 2026, The Fourteenth International Conference on Learning Representations.},
  year={2026},
  organization={https://openreview.net/forum?id=u6JLh0BO5h}
}
        

Contact

For more information or inquiries, please reach out via email.

If you have questions or want to contribute, please visit our GitHub Repository.