Jet Expansions of Residual Computation

Authors: Yihong Chen, Xiangxiang Xu, Yao Lu, Pontus Stenetorp, Luca Franceschi
Primary area: interpretability and explainable AI
Keywords: Interpretability, residual networks, transformers, LLMs, expansions, Taylor series, logit lens

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.

Decision: reject
Reviewer scores: 6, 6, 5, 6

Reviews and discussion:
<thread>
creation date:
1730300215007

thread_type:
official_review

<thread_object>
creator: reviewer_xfnq

<content>
summary:
This paper utilizes a convex combination of Taylor expansions to rewrite residual networks up to a nonlinear residual term. Crucially, these expansions are such that the contributions from different combinations of network subunits can be studied separately. This results in a data-independent interpretation tool for understanding black-box residual networks. The authors use this developed tool to study the functionality of subunits, pretraining dynamics, and finetuning in the context of language models.

strengths:
The proposed approach is principled, intuitive and unifies certain prior works. As shown by experiments, it can identify the linguistic functionality of various computational subunits in language models.

weaknesses:
Overall, the approximation quality of the jet expansions is not guaranteed, and hence faithfulness to the actual network and its behavior is unclear. This is acknowledged by the authors, and the approximation quality does not necessarily improve with scaling k (as seen from the experiments). Therefore, experimental explorations with jet expansions are only indicative without any confidence.

questions:
In Figure 2, the truthfulness of jet logits decay when $k=2$. Could you comment about scaling with respect to $k$? Do you expect it to improve the quality of the expansions?

Can you explain how to obtain bi-gram and skip-n-gram expansions? Overall, I think these two paragraphs (L259 "Jet bi-grams and skip-n-grams statistics") could be a bit more detailed for comprehension of the reader.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730474307716

thread_type:
official_review

<thread_object>
creator: reviewer_cns2

<content>
summary:
The authors rewrite a model with skip connections as a composition of many Taylor series, evaluated at various points that are related to intermediate activations, promising data-independent global interpretability. The approach is demonstrated on LLMs and conceptually compared to related methods such as LogitLens.

strengths:
- The mathematical exposition is clear.
- The authors acknowledge the limitation of their method in capturing the nonlinear model exactly.
- The applicability of the proposed method for evaluating models globally is interesting and promising. In particular, the model diffing experiments provide the potential of useful metrics for assessing the effectiveness of a specific fine-tuning method, the rate of model improvement and saturation, and the potential for certain emergent properties from n-gram statistics.

weaknesses:
- The first four sections contain clear mathematical expressions. The remaining sections do not use any of these notations which makes it very hard to digest what the figures are measuring in the context of the proposed method. I’d encourage the authors to improve the clarity of the figures and the captions. At the moment, they are very unclear.
- One way to address the first weakness would be to add a new section, between the theoretical section and the empirical section, which explains in the greatest detail possible what exactly is going to be measured empirically.
- The authors claim that their method provides global interpretability but it is unclear how their method is able to provide insights without evaluating the Jacobians at specific points.
- Is the proof of Lemma 1 a novel contribution or is it a well-known result?
- The proof of Lemma 1 (in the Appendix) is not clear.
- Before Equation 9, you state “x_{empty set} = eta”. The notation $x_set$ wasn’t previously defined and is therefore unclear.
- Could the authors expand on the algorithm bubble? At the moment, the steps are not very clear.
- It is unclear from the paper how the jet expansion relates to n-grams.
- The relation between LogitLens and the proposed method should be made more explicit.
- Superposition is mentioned multiple times throughout the paper. The relation with the jet expansion is not clear.

questions:
- "filtering out all paths that involve self-attention modules" — why is this necessary or reasonable?
- Figure 2
    - How is the top table related to the bottom figures?
    - Is it necessary to put all this information in a single figure?
    - Why should we measure the “cosine similarities between original and jet logits of joint (left) and iterative (right) lenses”? What do we learn from measuring it?
    - How do we see evidence for superposition or neuron polysemy in this figure?
    - Is a simpler method like LogitLens capable of identifying similar patterns?  Is there some simpler baseline you could compare your method to?
- Tables 1 and 2
    - "∆ logit after intervention" — what is the exact definition? unclear what this means.
    - What's the order of the expansion?
    - Is a simpler method like LogitLens capable of identifying similar patterns? Is there some simpler baseline you could compare your method to?
- "One-to-one bi-grams like" and “Many-to-many bi-grams” — unclear what does this mean.
- Figure 4
    - What’s the definition of a "hit ratio"?
    - What’s the definition of "total mass"?
    - How do I see double descent or grokking in this figure?
- What’s the definition of “diffing jet bi-grams”?
- "small change in mass" — what’s the definition of “mass”?
</content>
</thread_object>

</thread>
<thread>
creation date:
1731019154305

thread_type:
official_review

<thread_object>
creator: reviewer_nz3b

<content>
summary:
The authors develop a method that expands a residual network into the sum of an exponential number of jets (roughly, taylor expansions of different components). Each term is a "path" information takes as it traverses the network. They apply this to interperability,

strengths:
The method itself is very interesting; thinking of a network as a sum-of-paths is very natural and the jet formalism seems to capture it in a nice way.

The authors show that this generalizes  prior work such as the logit lens.

Since the top network architectures are residual, this is applicable to the most common model types.

Interpereting parts of the network is an important and relevant topic.

weaknesses:
- The primary issue is confusing exposition and incomplete details. These make the contributions difficult to asses. I enumerated specifics in the "questions" section. 
- The exponential expansion factor means that to analyze a model like LLaMa 405b, one would have 2^118 terms which seems a bit unwieldy
- Presumably, you need ways of computing k-th order jets for network components (the authors don't seem to discuss this), which makes implementation difficult.

questions:
- Paragraph after eq (9): what is gamma_3? it is not defined or used in the figure. This seems important. Is it a hypothetical 3rd block? 
- Lemma 1: Unclear where w comes from. I have 3 hypotheses: (1) w is arbitrary and changes which member of the equivalence class you have (2) there is a specific value of w needed ("there exists a w...") (3) there's a typo and the LHS should be a weighted sum of x. This needs to be explicit.
- What is the point of tthe if l < L in Algorithm 1? Isn't this always true? In the else, how can there be a gamma_{L+1} if there are L layers? where does w come from? 

line 264: "For example, bi-grams statistics related to Pq (z2|z1, . . . ) can be computed
by evaluating bi-gram paths, which we can obtain by expanding the LLM with Algorithm 2 and
filtering out all paths that involve self-attention modules."

If you filter out all self-attention paths, aren't the bigrams z1, z2 independent? this needs to be true because self-attention is the only mechanism for a transformer to route information amongst tokens.

How are you optimizing the jet weights? This seems to be very important as you show the weights in Figure 2, but you only briefly mention that it is "done cheaply" without any details. Do you need specific datasets? do you use SGD? How do you actually compute the residual in a way that can be optimized?

How do you represent (computationally) and evaluate the jets? The paper presents them abstractly, which is fine, but in order to compute anything you need to be able to evaluate the k-th order jet for an MLP or self-attention layer.


Minor comments:

In eq (7), I found the notation `J^k f(x) = f(x) + ...` to be a bit confusing; suggest `(J^k f) x = y \mapsto ...`
</content>
</thread_object>

</thread>
<thread>
creation date:
1731348598355

thread_type:
official_review

<thread_object>
creator: reviewer_s4js

<content>
summary:
The paper introduces a method for expanding residual networks, such as Transformers, using jets—operators that generalize truncated Taylor series. This approach aims to disentangle the contributions of individual computational paths to model predictions. The authors claim that their method subsumes the Logit Lens and demonstrate its ability to extract n-gram statistics from intermediate model layers, also enabling a data-free approach to detoxification.

strengths:
Well-written and theoretically well-developed, with content that is thorough yet not overwhelming.

weaknesses:
I find the theoretical foundation to be solid, but my main concerns lie with the experimental approach.

The experiments may be overly empirical and lack statistical rigor. For instance, in Section 5.1, only a handful of jet paths corresponding to specific linguistic functions are selected to demonstrate intervention effects. A more systematic approach is needed to demonstrate that the jet lens is more effective than the logit lens.

While the paper primarily offers an analytical framework, it lacks actionable insights for model steering. For example, could it be demonstrated that the bi-gram statistics can be leveraged to guide more efficient pre-training or improve RLHF techniques for reducing toxicity?

questions:
I have mostly raised them in the weakness section.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732964834094

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Revised PDF

comment:
Dear reviewers,

We would like to thank you all for your time and effort reviewing our paper. We believe your advice has been very helpful in helping us improve the readability, organization, and presentation of this work. Major changes in the main paper are highlighted in red. In the new revision, we have made the following changes:

1. We added a new experiment, where we measure the runtime scaling w.r.t to the jet order $k$.
2. We added a description on using JVP to compute the higher-order jets
3. We introduced the jet weights more clearly in Lemma 1 and described how we optimized them in more detail in Remark 2 and Sec 5.1 “Setup”
4. In  Remark 3, we clarify that the approximation quality depends on both the jet order and the choice of jet centers. This should make it clear that scaling k only does not necessarily result in lower approximation error.
5. We added a section (Section 5.1, “Setup”) that bridges the theory and the experiments, explaining how we compute jet bi- and tri-grams and introducing metrics that we compute in the experiments. In addition, we added more details and formal definition on these in Appendix C and D, respectively.
6. We added details on jet bigram diffing in appendix E.
7. We revise and clarify experimental comparison with LogitLens, decoupled ex Figure 2 and clarified exposition around line 350.

For your convenience, we uploaded a pdf diff in the supplementary material (although some parts that are highlighted as differences have not actually changed). 

We thank all reviewers for the suggestions and modified the manuscript to address the concerns. 
Our main contribution is a general theoretical framework to identify the contributions of different residual paths, which most reviewers also pointed out. Meanwhile, we would like to know if the reviewers have any further comments.
</content>
</thread_object>

</thread>


Metareview:
metareview:
In this work, the authors introduce a framework for expanding residual networks into linear residual networks using Taylor-type expansions (jets). These jets encode path information through the computational graph which can then be interpreted to better understand higher-order behavior during inference. All reviewers expressed mixed opinions of the work, acknowledging the report is generally well-written (at the sentence level), and clearly shows a substantial amount of thought and development. The general idea is attractive, but the report falters with the details. Much of the "weaknesses" raised seem to stem from a general confusion around many of the developed concepts, including how the jets are computed, how they provide insights, what some of the terminology means, what the baselines are actually capable of and how they are subsumed, how the jets relate to n-grams, and what the jet weights are. There are many other missed details here, but given the novelty of the approach and its potential for a new form of mechanistic interpretability, a clear exposition should receive glowing reviews. 

I would recommend that the authors reconsider how the document is structured. Perhaps rather than focusing on the nature of the theoretical framework and broad discussion of its implications, I suggest considering the report from the angle of a practitioner and place practical motivations front and center: first, how does one interpret a network using current approaches (e.g. LogitLens); then, what is a general model for interpretability and how might it be used (linear residual networks), what is the ideal toolkit etc.; and finally, introduce the jets as a natural generalization. Provide examples of computational complexity. All implementation details, including the code segment provided to Reviewer NZ3B, are required for reproducibility, but these can go into further Appendices as needed. However, they should not be in remarks, which can be easily ignored when read.

This is just an example, but regardless, I believe the manuscript is not yet ready, and will require an additional review process once it is reformatted to better suit practitioners.