Jet Expansions of Residual Computation

Yihong Chen, Xiangxiang Xu, Yao Lu, Pontus Stenetorp, Luca Franceschi

Keywords: Interpretability, residual networks, transformers, LLMs, expansions, Taylor series, logit lens

TL;DR: We introduce a novel framework that expands residual networks using jets, revealing their internal computational paths and enabling various applications of data-freee interpretability.

Abstract:

We introduce a framework for expanding residual networks using \textit{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 the logit lens,
reveals a (super-)exponential path structure in the network depth and opens up several applications.
These include the extraction of $n$-gram statistics from a transformer large language model, and the definition of data-free toxicity scores.
Our approach enables data-free analysis of residual networks for model interpretation, development, and evaluation.

Primary Area: interpretability and explainable AI

Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.

Submission Guidelines: I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2025/AuthorGuide.

Anonymous Url: I certify that there is no URL (e.g., github page) that could be used to find authors’ identity.

No Acknowledgement Section: I certify that there is no acknowledgement section in this submission for double blind review.

Submission Number: 11117