ICLR2026
Decomposing LLM Computation with Jets
Yihong Chen, Xiangxiang Xu, Pontus Stenetorp, Sebastian Riedel, Luca Franceschi
摘要
Large language models are becoming general knowledge engines for diverse applications. However, their computations are deeply entangled, resisting modularization which complicates interpretability, auditing, and long-term maintenance. We introduce JET EXPANSIONS, a framework for expanding recursive residual computational graphs using jet operators that generalize truncated Taylor series. Our method systematically decomposes language models into explicit input-tooutput computational paths and complementary remainders. This functional decomposition provides a principled, knife-like operator for cutting through entanglement in LLMs, enabling scalable model inspection. We demonstrate how JET EXPANSIONS ground and subsume the popular interpretability technique Logit Lens, reveal an ensemble of an exponential number of paths analytically verify prior research, and support several interpretability applications, including sketching a transformer language model with n-gram statistics extracted from its computations and indexing model toxicity levels without curated benchmarks.