Decomposing LLM Computation with Jets
Authors: Yihong Chen
Xiangxiang Xu
Pontus Stenetorp
Sebastian Riedel
Luca Franceschi
Venue: ICLR 2026 Poster (2026)
Keywords: decomposition
transformer
neural-symbolic
n-grams
interpretability
controllability
Primary area: interpretability and explainable AI
TL;DR: We introduce jet expansions: operators that "cuts through" LLM entaglement, separating parts of computation of interest and enabling systematic model inspection like n-gram tables

Abstract: Large language models are becoming general knowledge engines for diverse applications. However, their computations are deeply entangled after training, resisting modularization which complicates interpretability, auditing, and long-term maintenance. We introduce Jet Expansions, a framework for expanding computational graphs using jet operators that generalize truncated Taylor series. Our method systematically decomposes language models into explicit input-to-output 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 a (super-)exponential path structure with respect to recursive residual depth, 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.

Published as a conference paper at ICLR 2026
DECOMPOSING LLM COMPUTATION WITH JETS
Yihong Chen', Xiangxiang Xu Pontus Stenetorp", Sebastian Riedel", Luca Franceschi
SOTMM., University of Oxford, UK Y AA Centre, University College London, UK
Indeppedentt researcher, Berlin, Germany University of Rochester, USA
yi hong- chenccs ow.ac.c......oooo.....ooo....... ac.ukk luc. franceschi @proton. me. eenpgiianguuurccccooo edu, (p. (p. stenetorp, riedel (css www.ucc....... acl.laccuu
ABSTRACT
Large language models are becoming general knowledge engines for diverse ap- plications. However, their computations are deeply entangled, resisting modular- ization 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-to- output computational paths and complementary remainders. This functional de- composition provides principled, knife-like operator for cutting through entan- glement 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 sketch- ing a transformer language model with n n-gram statistics extracted from its com- putations and indexing model toxicity levels without curated benchmarks.
INTRODUCTION
Earlier artificial intelligence (AI) featured symbolic systems, representing knowledge in units, such as entities and relations. This mirrors how humans might conceptualize reality into structures for repeatable reasoning about the reality. In contrast, the latest advancement surfaces unstruclured paradigm, particularly since large language models (LL.Ms) trained with massive unorganized web texts (Radford al.) 2019 Brown al. 20200 Touvron al.) 2023a). Both paradigms provide pathways to build knowledge engines that power diverse downstream applications (Chen) 2023).
Unlike symbolic methods, which format knowledge into explicitly addressable units customized for query calculation, LLMs distribute knowledge across billions of highly entangled parameters. This misalignment between knowledge layout and computation layout lies at the core of LLM opacity, raising regulatory concerns about transparency. auditability, and long-term maintainability. Once trained, LLMS cannot be easily inspected updated. Knowledge operations that trivial in sym- bolic AI, are far from straightforward in LLMs (Gehman et al. 2020 Carlini et al. 202 Mitchell al. 202 2022 Meng et 2022). This fuels growing demands for LLM interpretability, partic- ularly in high-stakes domains such as healthcare (Smith) 2021 He et al. 2025 Comeau et al. 2025)
and robotics (Wachter et al. 2017 Fernández- Becerra et 2024 Raptis et al. (025)).
Existing interpretability methods often take data-then- thepellttonnlllaaattoon approach: curate inputs, hy- pothesize which sub-computations matter, and observe activations refine the hypothesis (Wang et al. 2022 Meng et al.) 2022, Goldowsky-Dill et al. (023). We argue that the real challenge is structural: LTMM computations are entangled, preventing us from isolating embedded knowledge gain valuable insights with data-driven interpretability ap- proaches, we posit that the ability to restructure computation into smaller, less entangled, end-to-end components "systematically" rather than "empirically" is central tackle such issues scale.
into meaningful units. While
We present JET EXPANSIONS, principled, general-purpon framework for manipulating LLM computations. Noting that LLMs particular types of residual networks (He et al. 20166) Vaswani ct al. 2017). our key idea is to recursively expand residual computations using jet operators (Ehres- mann 1951 the functional counterpart of truncated Taylor series. This process yields functional

---

Published as a conference paper at ICLR 2026
Functional Space Decomposition
Interpretability
Computation Interest
Application
Readout
taken contribution heatmap component role identification model diff
Jet Expand
LLM
jet/logit lens jet ngrams
training monitoring
finetuning verilication
Remainder
knowledge quantification
Figure 1: JET EXPANSIONS restructure residual computations into explicit input-toutput paths and a com- plementary remainder. From these paths extract logits, and n--rams without retraining additional data. These readouts support downstream applications such token contribution heatmaps, model comparison, training monitoring, and fine-tuning verification.
rewritings of the model into two parts: (i) explicit input-routput polynomial functions, which we call jet paths, and (ii) complementary nonlinear remainders. Crucially, JET EXPANSIONS operates functional level, requiring no additional data, nor training. We show that JET EXPANSIONS encompass existing interpretability tools such the Logit Lens (nostalgebraist, 2021b), and extend them to new instantiations such as extracting n-gram probability tables from LLMs. This enables dataset-free, symbolic sketches of transformer LLMs and allow perform global interpretability studies. Figure [egure[iillsteern Tillustrates the pipeline.
We validate framework through case studies across several autoregressive LLMs (GPT, Llama, and OLMo). JET EXPANSIONS enable several empirical model inspection usages: i) understanding inner mechanisms via jet lens (Section 5.2. Ip: ii) assessing fine-tuning effects, e.g. quantifying toxicity levels with jet n-grams, showing RLHF alignment (Bai et al. (002) reduces but does not eliminate toxic knowledge (Section 5.2.2): iii) analyzing pretraining dynamies, e.g. tracking how bi- grams such as "at least" are promoted then suppressed in OLMo (App. D These results demonstrate that JET EXPANSIONS provide a dataset- free operator for restructuring LLM computations, paving the way for more transparent, interpretable, and maintainable foundation models.
Our contributions.
1. A new angle interpretability: treating it as function decomposition, rather than input- driven attribution or circuit identification on particular datasets.
2. A principled theoretical framework, based on jet operators, formally grounding existing tools such as Logit Lens (nostalgebraist) 2021b a) and path expansion (Elhage al., (021).
3. Preliminary but wide-ranging case studies, revealing insights into LLM internal mecha- nisms, fine-tuning knowledge shifts, and toxicity levels.
BACKGROUND AND PRELIMINARIES
Language models as residual networks. We focus on transformer language models (Vaswani al. 2017). which residual networks (He et al. 2016) consisting of I. stacked residual blocks sandwiched between an encoder Enc and a decoder Dec. Formally, the full computation is
X: x +
Dec (Ol + +7)) Enc,
where ye is the non-linear transformation in block f. Unrolling the recursion, the hidden state after blocks is
hg h ho + E5=13 hy-1- hg = Enc(z).
(2)
This recursive form makes clear that residual links accumulate contributions from all preceding layers. We adopt the notion of residual streams (Elhage et al., (021), where the computation of hy can be viewed as nested terms entangling contributions across blocks; see also (Veit et al.) 2016). Table summarizes the notation.
2

---

Published as a conference paper at ICLR 2026
Table 1: Summary of notation used in the paper.
Symbol
Meaning
Symbol
Meaning
x'
Input space
L
Depth (no. of blocks)
Vocabulary size
klentity map
y RV
Output logits
U
Unemberkling matrix
Hidden dimension
2/
Final normalization
f:x y
network
layer
Enc X Rd
Encoder
Residual block layer
Dec: R
Decoder
Residual transform inside block
Base point (senter)
Variable
D f(za)
j-th differential
zo)"
j-fold tensor product
gk (mo)
k-jjt
3
Jet operator
pk
Degree- & polysomial space
Jet weight for é-th hase point
Set of expanded terms
Remainder of expansion
Taylor expansions and jets. To handle nonlinearities when restructuring residual computations, turn jets [Ehresmann) 1951), which generalize Taylor expansions. For f O (IR RR),, Taylor's 's theorem at base point zo gives
f(z) (zo) E Dz D'
O(I xoll(t1).
(3)
The k-th order jel operator abstracts this expansion as
f:: pk, J fxx))(x)= f(0)) Σi= Jt DJ f(zo) (x x0)
or equivalently, by leaving the polynomial action implicit,
2hh fau) f Σラー Dif(xo)) f(xo))
Intuitively, J fxo) captures the local structure of f up to order k. and we write f(x) Jh (z) to indicate agreement up to order Jets thus provide principled operator for rewriting computations of f into decomposable pieces.
Remark (Base points and variables functions). When tracing back to the input X. base points Zo and variables & may themselves depend on z. In that case, jets define maps X' - y via J (=o(2))(<<=))) For brevity, often omit writing the variable explicitly when clear from context. See Appendix for details.
3 RELATED WORK
Mechanistic interpretability and path rewriting. A large body of work has sought to inter- pret the inner computations of large language models. One prominent category mechanistic lerpretability(M) (Ferrando et al. 2024), which aims to reverse-engineer model computations by identifying, clustering, and labeling behaviors (Shah et al, 2024, Meng et al.) 2022; Bricken et al.) 2023) and attributing them to specific components, such as ML.s (Geva et al. 2021 2022) or cir- cuits Conmy et al., 2023] Ferrando & Voita 2024). However, these approaches often restrict anal- ysis atomic components (neurons, layers, weights), which may not reveal the full mechanism of information processing. For example, Templeton et al. (004) highlight the difficulty of drawing conclusions the neuron level compared with higher-level feature representations, while Bolukbasi et al. (2021): I): Goldowsky et al. (2023) emphasize that many findings depend heavily on the cho- data distribution. A second category of approaches attempts explicit path rewriting. Veit al. (016)) syntactically expand residual networks into exponentially many paths of varying length to study gradient behavior. [Elhage al. (02TT) decompose one- and two-layer transformers into sums of uni uni-gram and bi-gram computation paths. Goodowsky-DDi et al. (2023) extend this line of work by developing path patching methods that aim to preserve functional faithfulness while isolating specific behaviors. Aligning with the second category, our approach manipulates functions directly rather than activations. requires neither probe datasets (Belrose et al., 2023) nor sampling (Conmy) al. 2023 Ferrando & Voita 2024 Voita et al. 2024). By allowing arbitrary portions of computa- tion to be isolated from the monolithic transformer, JET EXPANSIONS abstract and generalize prior bath-based characterizations (Veit et al. 2016 Elhage et al. 2021), where nonlinearities were often ignored or simplified (e.g. omitting layer norms, linearizing components, or implicitly assuming the nonlinear compositionality does not destroy the supposed independence of paths).
3

---

Published as a conference paper at ICLR 2026
N-gram models symbolic counterparts. N-gram models, dating back to Shannon 1948) represent one of the earliest symbolic approaches to language modeling. They store explicit prob- abilities of token sequences, e.g. Pr(w; W;-1...... W;-n+1), in tabular form. This makes their knowledge layout identical to their computation layout: each symbol sequence has a directly ad- dressable probability entry. Such symbolic modularity enabled early successes in language mod- cling (Goodman 2001) and tasks like machine translation (Brants et al. (007). While later work combined n- n-grams with networks (Liu al., 2024), recent studies revisit their role in relation to LLMs: analyzing the ability of transformers to simulate n-gram statistics (Svete & Cotterell 2024) or measuring agreement between LLM predictions and n-gram rulesets (Nguyen (2024). This re- newed attention motivates a direct bridge between 7-grams and LLMs. JET EXPANSIONS provide this bridge, allowing corpus-free extraction of n-gram statistics directly from LLMs and thereby recovering form of symbolic modularity within their entangled computations.
Taylor expansions and jets. Taylor expansions are ubiquitous tools in analyzing learning be- haviours (Jastrzebski et al. 2017), notably with linearization (k = 1). For example, [Belrose et al. (2024) applied Taylor expansion on the loss to demonstrate the learning preference of neural net- works. Xu et al. (2022) used a second-order Taylor expansion over the data distribution to interpret optimal features. The generalized jet notions were introduced in machine learning in the context of automatic differentiation tools by [etteenoout et al. (2019), and is an experimental feature in Jax (Bradbury et al. 2018), but have been studied hefore (see e.g. Griewank & Walther 2008). We lever- age jets not merely as approximation tools, but as operators to restructure residual computations in LLMs into explicit -ntptt paths and complementary remainders.
4
RESTRUCTURING LLM COMPUTATION WITH JET EXPANSIONS
4.1 LINEAR CASE: EASY TO RESTRUCTURE
We begin with the linear case, where residual computations can be reorganized exactly. Assuming Te(x) = AAw for some A E indxd, encoder Enc E, and v = id, eq. ① expands as follows
= o (O'z (id + +w)) Enc = Escat IIee Ar) Ar) E Escon fs.
(5)
and each path fs = U ( E UWSE, with We = T, itself linear map from to y. Thus, the entire network can be written as the sum of 2 explicit input->output paths fs. This exact decomposition makes linear residual networks intrinsically easy to restructure for interpretability because the output is a simple sum of its components:
where 2 is the power set of [Z] (1....,I, } L
can directly analyze the contribution of each path, study their interactions, and understand the global input -ouptut relationships. In the nonlinear case, however, such a clean decomposition no longer holds, motivating the of jets.
4.2 NONLINEAR CASE: JETS TO THE RESCUE
J f (xr) encodes all information about a function f E Ck+1 up order-k derivatives at a base point To, providing vector-free representation of its local behavior. This makes jets principled tool for reorganizing computations LLMs. Lemmal] [ proved in Appendix formalizes their disentan- gling property: jet at sum of inputs can be written as a convex combination of jets at individual inputs, up to higher- order error. This allows us to carve apart nested residual terms into separate, analyzable contributions (Figure 2 Lemma 1 (Disentanglement of Jets). Let f E C (R. Rd),, k N, N N+, {x be a set of jet base points, and w N-1 C N be a set of jet weights (i.e., w; 2 0, Σ, w/ =E 24 and max; we || øl|.
Actual entangled output
"Ideal" disentangled output
f0r + x2)
forg)
f(x2)
Layer
+
Residual stream
Figure 2: Convex combinations of jets
Define the sum F-- Li=1=
disentangle a residual stream he (a sum of terms) into sub-streams in function space, each isolated for further analysis.
Then the k-jet of f at the sum & satisfies
J
EM 1i=1 X3 = -i we Jkt ((x:)) + 0(rkt1).
4

---

Published as a conference paper at ICLR 2026
dependency
jes center
jes variate
Inout encoding
output decoding
Intermediate nonlineartly
(a-1)
(a-ll)
(b)
(c)
Figure 3: Carving two-block net. (a) Nested entanglements. (b) Inner expansion "Y2. (c) Outer expansion Dec, yielding explicit paths. Interpreting a model component entails selection. Jet centers trace paths through modules of interest; jet variates trace paths through those excluded from that focus.
Example (JET EXPANSIONS of ReLU). Consider the ReLU activation function
+ R+
defined as ((x) [. For 0, 7 (a) 1. For < 0, 7'xx) 0. Higher order derivatives zero almost everywhere. If x = Iı + 22, then for almost every æ, there exist w E A such that
(x1 x2) (x=)) + wyJ ",J'''<z)>()) (x2)(x) = wi (1((((11) 7 (x1)=2) + w2 (2(())2)) + (x2))=))
In other words, for almost every base point I1 + z2, there exists convex combinations of jets that can recover the original function y exactly. Indeed, if either T2, T2 > 0 or #1, 22 < 0, then any convex combination is exact. If only one of the two terms is positive, say
and D2 O, then we can set W1 = if 21 + Z2 2 0 and W1 = otherwise (w² = 1 w1). The specular argument applies for the case In < 0 and 22 0. From global perspective,
think of jet weights wy wo (1())1))) as optimizable functions of ι1 and 22 rather then constants and in the ReLU case, we obtain (almost every werrw an exact first-order expansion. Conversely, one can see that the O-th order JET EXPANSIONS of y is not globally exact.
4.2.1 MOTIVATING EXAMPLE: CARVING A TWO-BLOCK RESIDUAL NETWORK.
Now consider how use jets carve typical computation graph. We begin with the simplest nontrivial case: a network with two residual blocks. Using Equation (T) its full computation is
Dec
Enc + 71 o Enc + 72 o (Enc + m Enc
The nested parentheses entangle contributions: the outer (purple) grouping mixes everything, while the inner (orange) ties 72 to both zo and zT. Traditional MI would select paths syntactically, akin to selecting modules in PyTorch computation graph, ignoring these nesting effects. Jets let both levels systematically and isolate the contributions of different uiput>>outpt paths.
Step 1: Inner expansion. At 722 taking fmo, 21 as jet base points and using Lemma the residual stream 22 72 o (xo + x1) be decomposed
m (2 + x1) W (e + wiJ (2())) +0(r++1),
so the original entangled stream 22 separates into two sub-streams, as illustrated in Figure 3(b).
Step 2: Outer expansion. At Dec, the jet base points are updated from (Nn, 31, x2) to
(2o. ,2, F21) after previous expansion. Using Lemma I and jet algebra (Proposition) 1. yields
f my D'Dec + t +.20 + x21) " (DecoEnc) + W.J (Deco)) Enc)
f()
(D (12 (xnne) + w (Deco (w)J J 72(m Enc))) +O(1 +O
122)
f
5

---

Published as a conference paper at ICLR 2026
corresponding four distinct >utut- paths fø. f(2) This is shown in Fig- ure C)). Each term aligns with what one might pick manually as a "path" in the network, but here arises systematically from the JET EXPANSIONS.
The toy example above illustrates two key principles of our approach: recursive expansion of the nesting terms, and the use of disentangling property (Lemma(T) to isolate entangled contributions. In deeper networks with many blocks, however, manual expansion becomes infeasible. This motivates our general-purpose algorithmic framework, bringing such expansion to any depths.
4.3 GENERAL FRAMEWORK
Algorithm T describes the core operation of JET EXPANSIONS. At each block I, the algorithm applies Lemma [ to a set of jet base points For convenience, here I, + indicates the final decoder. Dec. The outputs are: (i) &, the set of polynomial terms (jet paths), where each term the jet expansion centered at each 2; E C: and (ii) 8, nonlinear remainder, collecting ror stemming from the Taylor expansion in Equa- tion (3) and error from Lemma [ A key fea- ture is that jet base points can themselves be the outputs of earlier expansions. This enables recursive application of jet expand through- out the network, unrolling the computation graph into end-to-end input-routput paths. In particu- lar, when we apply jet.expand at the final de- coder layer 1 = L 1, we obtain a functional rewriting the model. Assume (E++1,01+) jet ret_expand(f,.d 1, k) for some choice of jet base points and order k. then
Algorithm 1 jet.expand(f,4.C
Require: Net f as in eq. Q block index E [L + 1]; jet base points
{ai)/= jet
order k E
Ensure: Expanded polynomial terms with weights w, and remainder 8 1: if &<L then
2:
residual block from f " < ffblock(())
3:
4:
I/ residual block computation
5:
he + he-1 + (he-1)
6:
jet expansion at block
7:
& [w,Jh }
8:
jet expansion residual link
9:
& + &U (w:Jid(<;)) W
f(x) = Σ e(x, w) + 6L+1(x, w), CELL+1
(6)
10:
// calculate remainder
11:
+ he Eeee e
else
where the jet weights E N-1 can be man- ually specified or optimized. Hence, recursive ap- plications of jet_expand yield rewriting of the model as a sum of explicit paths, plus a comple- mentary remainder. Each path is atomic unit of computation from input to output, mirroring the original function, but with simpler, additive struc- tures. This decomposition is purely algebraic and requires no extra data collection. 1/N, or optimized
13:
jet expansion decoder
14:
E { f.Dec(r;)) 8 8 f.De(hh)) - Zeee
15:
16: return (5.8)
Remark 2 (Jet weights). The jet weights can be fixed, for example to minimize the remainder at a given v, such as logit space. When the decoder linear, this optimization is efficient, as
||U6z41 w)l (l)(hz())) - EeEsci e(z, w()l
representing the squared distance between the expansion and the residual stream in Rd, with the metric induced by U.
Remark 3 (Remainders). Remainders generally do not vanish with increasing k, as the base points are user-defined. For linear residual networks, however, = 0 for all & > 1, showing that Algo- rithm exactly recovers the rewrite in Equation (5) for any w. In light of Equation (6). jet expansions should be viewed as algebraic rewrites of computational graphs, intended to aid interpretation rather than to minimize approximation emror. In our experiments, 5 is often small, and the cosine similar- ity between expanded and original logits approaches (Figure 4 bottom). See App for detailed discussion.
Lemma 2. Residual nety with only ReL.U nonlinearites admit exact first-order JET EXPANSIONS.

---

Published as a conference paper at ICLR 2026
Runtime. Evaluating & and å at E X' requires computing ktth-order jets at cost (F + kB)), where F and B denote forward and backward passes of f. In practice, higher-order jets can be computed efficiently via recurrence relations and automatic differentiation primitives such as Jacobian-vector products (JVPs) (Grriewank & Walther 2008 [eettencourt et al., 2019). App D
reports empirical runtime scaling of our implementation.
5
APPLICATIONS OF JET EXPANSIONS
5.1 THEORETICAL APPLICATIONS
JET EXPANSIONS offer a principled framework that unifies and grounds existing techniques, such as Logit Lens, while enabling systematic derivations of new methods. Here we present several concrete instantiations of JET EXPANSIONS, and set the stage for the subsequent empirical studies.
(Super-)exponenn expansion. Algorithm eex- tends our two-block example to arbitrary depth, pro- ducing 24 paths with uniform jet weights. This mir- rors Veit et al. (2016)'s exponential view of resid- ual networks, but in an explicit and principled way. We defer the discussion on the connection between the expansion and the ensembling view on LLM computation to C For k 2 1, each polynomial term can be decomposed further by degree, isolating higher-order block interactions, hinting at super- exponential ensemble perspective which we leave as future work.
Algorithm 2 2xxp_jet.expand(f.
Require: Net f as in Equation I jet order l E E.
Ensure: Expanded terms & (with uniform weights, |5| 24) and remainder
// initialize expansion
& + {f.Enc, "1 f.Enc} 3: for l = L + + do
(5.5) jet jet..expand(f,.....
// reweighht terms uniformly 6: & (/K() E 8 + f.Dc(hl) Eeege return (5
Jet lens and logit lens. The logit lens (nostalge- braist 2021b Geva et al. 2021 2022, Merullo et al.
2023
Belrose et al.] (023) is a widely used mecha- nistic interpretability tool that applies the decoder
intermediate hidden states:
Logit Lens, = Uu(he(=)) = Deo(he(z)).
Aimed at highlighting the iterative refinement of the prediction across blocks,it is related to early exiting or carly decoding in the context of conditional computation (see e.g. Panda al., [016]] Elbayad et al. 2020 Geva et al. 2022). We can rewrite the logit lens with jet operator as follows
(ogitLens,(z) Dec(h)(z)) = Ded he(()))))))))))))))( (he))) = =J hh(()) Dec (hz(:)).
Here we retain the argument (z) to emphasize that the zeroth-order jet applied within the full computation, as if the jet operator acts like knife: slicing the network at layer & and replacing the sliced computation with truncated jet expansion. Indeed, Dec(=) R4=0 he(()) Dec(x)
Dec(hs(z)), so the logit lens coincides with the zeroth-order jet of the decoder at the base point he(z), equivalently jet det__exanddff. + 1, [he], 0). This perspective suggests two di- rect generalizations of logit lens. First, iterative jet lens extends logit lens to higher-order jets: Second, jet lens expands with a broader set base points rather than merely {he): jet jet_expand(f. L + 1, (re h [L) k). thereby highlighting contributions cach block instead of the cumulative refinement of the residual stream.
jet_expand([, I + 1, (he),, kh)),k),k
Jet n-ggrmm. 7-gram statistics have gain traction in analyzing LLMs (Elhage et al. 2021 Svete & Cotterell] 2024; Nguyen 2024), but existing methods rely on probing datasets. With JET Ex- extract n-grams directly from the model. Concretely, since model can be rewritten into a sum of polynomial terms or jet paths, we can select paths of interest, particu- larly shorter ones, and evaluate them the entire vocabulary space its cartesian product to record the resulting logits. Formally, given a model's expanded terms (jet paths) ç, each path e E defines function e : X Vn-1 + V. The n-gram score for token i given context æ s(x)[d] = Eese e(x) [+]]gl. By evaluating s(x) for all E Vn-1, we obtain a complete n-gram table ((x)()) x Vn-1, i E V. where (x, i) identifies the n-gram and (x)) gives its score.
PANSIONS,
7

---

Published as a conference paper at ICLR 2026
From theory to applications. Applying jet expansions in practice means choosing which portions of the computation of f one is interested in carving out, and consequently, applying the expansion operator jet expand(f: e, C, k) as many times needed, with appropriate layers l, centers C. and orders k, in order to obtain input-to-output functional, as remarked in eq. (6). After we specify a computation method for the jet weights, we obtain set of functions (the jet paths) and reminder. ! All empirical objects in our experiments correspond to evaluations of such paths either on specific input sequences (e.g.. "new simple neural architecture, the Transformer") or on the entire domain (vocabulary) of the resulting expansions. Following previous work, we use the first input-specific evaluation method for lenses experiments (section (...I). We use the second function-sp specioni-seiiiccco eval- uation method for all experiments concerning jet n-grams (section 5.2.2).
In particualr, we note that specific paths, like bi/tri-gram paths, restrict the effective domain to short sequences, where the domain X' = /n-1 is treatable since E (2,3). This allows us to perform exhaustive function-speeiffe evaluations and produce entire likelihood tables. For instance, evaluat- ing a bi-gram path V [0, 1] + E at weight w yields complete table, (I,, Tj e( w) [z,]) where T; E V are tokens. These symbolic tables, which we dub jet n-grams, enable a holistic, dataset-free characterization of the model's global behavior, supporting global, dataset-free analy- ses. In our experiments, we implements such analysis via: (i) top- K scoring bi-grams, which reveal broad behavioral tendencies of the model, and (ii) keyword-conditioned bi-gram mass, which aggre- gates scores over bi-grams associated with a semantic category (e.g.. toxicity), providing a scalar indicator of how much such knowledge is embedded in the model.
Hence, across both the input-specific (jet lenses) and function-specific (jet n-grams) settings, the connection to theory is direct: each empirical quantity is simply a jet path evaluated on its natural domain, with (e, C, k) translating theoretical choices into concrete experimental readouts.
5.2 EMPIRICAL CASE STUDIES
Setup. We experimented with open-source LLMs including GPT-2 RRdford al. 2019). GPT- Neo Black et al. (2021). Llama (Touvron et al. 2023a b Rozière et al. 2024)) and OLMo (Groen- leeeld et 20024. We run experiments ont n nn-grmms on CPU servers with 1TB memory, using zero order jets and uniform jet weights. We run jet lenses experiments on a single laptop CPU, using various orders and optimizing jet weights with gradient descent. Further details about algorithms and evaluation metrics are given in Epp and FF Code is available here
5.2.1 CASE I: ANALYZING LLM INNER WORKING
Jet lens. We use jer lens to analyze LLMs' mechanisms when processing individual examples. Figure 4 (top) visualize a joint jet lens for GPT-Neo-2.7B (Black et al. 202I). For space reason, other examples are deferred to App G. The first row is the input sentence. The first column indicates the blocks. Here, block, e.g. Block contains one self-attention and one MLP module. All table cells depict top-1 tokens for the corresponding path through the particular block, following conventions from prior work (Belrose et al., (023). We observe that the joint jet lens captures the synergy among different blocks, as the model prediction is decomposed into the contributions of several jet paths. Optimized jet weights are reported the percentages.
In this sense, jet lens with 1: > O may serve as tools to systematically discover such synergic be- haviors. We also find that higher-orders (k > 0) help iterative lenses deliver more meaningful interpretations than the logit lens (k 0) for GPT-Neo-2.70 (see Figures 7 to D). This is potentially due to their capability to trace indirect impacts of early layers on the final logits, which are otherwise missing under logit lens. Our findings are consistent with nostalgebraist (202 Cancedda (2024) where naive implementations of logit lens are shown to fail on GPT-Neo model family. Figure 4 (bottom) present mean cosine similarities of joint and iterative jet lenses with respect model out- puts for various GPT models and orders, averaged over 100 example sentences. The similarities are high and close to 1 for various k. showing however different behavior across model families and sizes. This indicates JET EXPANSIONS highly correlate with model outputs. In particular, the right plot compares the similarities of the logits obtained through iterative jet lenses for (solid, line, the same as LogitLens) and for k (dashed lines), indicating an higher correlation of the latter with model outputs, potentially providing more faithful interpretations.
In our experiments, we use either uniform or optimized weights, as deseribed in remark
8

---

Published as a conference paper at ICLR 2026
W
24
LBack (7 3221
12200
Sen 1208
Uoo 24 19001
Sever k.9%
new 1erE
Taam ELa00
Logits Expan. (3.99))
network network
WINch
nearal
Tormar
which
Joint jet lenses
Iterative jet lenses
Cosine similartyy
1.00
wmopt. (left) (right
Cosine similartt
1.00
GPT2; (leff) == (right)
0.95
GPTZ-4age: wropt. (left) or km 0 tright
0.90
(arg-large; avg. (left) or k ==11 (right)
0.90
GPT-Neo 2.78; w-opp. (left) or - (right)
0.85
GPT-Neo 2.78; w avg, (left) or (right)
0.85
jet order (k)
Block index
Figure 4: (Top) example of joint jet lens on GPT-Neo 2.7B with & = 1. visualizing the seven blocks with highest average jet weights after optimization. Each table cell indicates the most likely token predicted by the jet path of each block. Optimized jet weight are displayed in parenthesis next the most likely token. We used a diverging bluc-to-red color map tracking logit scores, centered at zero. The second table with two rows shows the model logits (Logits) and the expansion logits (Expan.), with cosine similarity (0.993) in parenthesis; in this case, all top-1 tokens perfectly coincide. (Bottom) plots of average cosine similarities between original and jet logits of joint (left) and iterative (right) lenses.
Table 2: MLPs in OLMo-7B and Llama-2-7B performing certain linguistic functions based on jet bi-grams extracted from the corresponding jet paths.
OLMe 78
Liana 1B
MLP Index
9
17
19
18
19
Linguistie Role
-1y
kog
-16
-ity
-Lng
A ogig Intervened
-4.19
0.58
-9.33
-422
-7.42
-14.6
-355
-11.95
-9.09
-0.14
Jet paths of individual components. By examining the representative jet bi-grams captured by jet paths of individual components, we can analyze their roles. our case, we find some MLPs perform special linguistic functions. For example, in OLMo-7B, the jet path which passes through the 3rd MLP promotes the addition of the "-ing" suffixes to the current token. Similar MLPs with certain linguistic functions are listed in Table 2 where the negative Alogit indicates removing the corresponding MLP harms the fulfillment of the particular linguistic functions. Note that the rela- not necessarily one-to-one mappings. Particularly paths through multiple MIPs might work together to complete one linguistic function e.g. MI.P and MLP 18 in Llama-2-77 can add "-ing" suffix. This echos work on circuit discovery (Elhage eet al. 2022 Conmy et al. 2023 Ferrando & Voita 2024). where the role of each component cannot be easily dissected and multiple components collaborate. Similar studies on the roles of attention heads can be be found in App G.2
tionships between functions and components
5.2.2 CASE 2: ASSESSING FINE-TUNING EFFECT
Fine-tuning steers an L.M from pretraining's 's vastness towards focused, task-specific intent. These shifts ripple distributedly across high- nigh-dimensional parameter space, often escaping full capture with- out extensive benchmarking, Jet n-grams, however, render the changes legible directly from the weights, revealing model differences through their n-gram "diffing".
Code fine-tuning. Comparing Llama-2-7 2-7B with its code fine-tuned variants, CodeLlama, reveals that diffing jet bi-grams highlights code-specific patterns such as **kwargs or Assertion (Table (3). confirming the acquisition of programming knowledge. This suggests that jet bi-grams can serve as a practical tool to verify if fine-tuning effectively imparts knowledge in target domains.
RLHF alignment. While ToxiGen scores suggest detoxification of LLAMA-2 2-BB 2-7BB-CHATT -CAA,, jet bi- gram masses remain nearly unchanged (Table 4) indicating toxic associations persist in latent form (additional evidence in App. H). Challenging prompts from RealToxicityPromp (Gehman et al. 2020) confirm that these associations can still be triggered. Thus, RLHF appears mask rather

---

Published as a conference paper at ICLR 2026
Table 3: Bi-grams before and after code fine-tuning. For brevity, show every 50th bi-gram among the top 1000. Bi-grams relevant to coding, such as **wwargs (a Python keyword), are highlighted. This demonstrates that our method can extract representative bi-grams reflecting fine- tne-ttning quality.
Rank
Liama-2-78
Codeliana-78 7B
Babellom--Python--
I_nore, _than)
L_kkk, xiae) Ljust, ification) [ Joasel Lget, terat Iktōber,
[_1kk, wisel LLike, wisel [al1, udes)
50
LNow, here]
100
[_aystem, stically)
[,lll, eng)
(no, isy)
200
L-c,, Lens)
toutput, ted)
300
Lother, world)
Lall, ud)
(0bjact, ive]
350
LJust, ifled)
igebiet, 5)
Las, cii)
400
Litust, ees)
(Brotest, Ldeploy, ment)
(ecan, nab)
Lat, he)
Leransport, ation) (Cag, ging)
500
bbook, mark)
Iclass, zoom)
550
[irrm, ITI
Laccess, ury) I_Tn, varlant)
(pecsonal, ized)
600
LJHEN, ever)
Lexoesa, 1ve)
650
Lahera, about)
LI, fadd, itionally) Linvalid, ate) Idiv, ision] Lprocess, ors) Lset,, up)
(Add, itional)
lag, ged)
L kwarget
750
1 he)
Inama, plates)
[all, anto)
(seeeec,, ive; (Assser, lona) Lean, collaticn)
850
I_Tom, arrow)
950
LBach, elor)
Table 4: Toxicity indexes for Llama-2-4 2- and Llama-24 2-nna----B-aa 2-7--chat using different methods: ToxiGen, jet bi- grams, and RealToxicityProme challenge prompting. Higher numbers indicate higher toxicity scores on the corresponding henchmarks and higher toxic knowledge possession for jet bi-grams.
ToxiGen Score
Jet Bi-grams
RTP Challenging Prompts
Hartvigsen al. 2022
Mass of "toxie" bi-grams
No
Very mild
Modium
Hard
Jama--2-88
21.25
0.102
39%
49%
64%
1lama-2-7B-cluu
0.0
0.093
23%
33%
64%
84%
than erase toxic knowledge, a finding revealed directly by data-free jet bi-gram indices. This show- cases a potential application of jet bi-grams in constructing data-free indices that reveal embedded knowledge, offering complimentary views beyond traditional data-driven benchmark evaluations.
6 CONCLUSION
We introduced jet expansions, principled framework for decomposing the computation of trans- former. Specialized to LLMs, our method systematically disentangles contributions of user-selected input->output paths from the overall computation, yielding interpretable functional components plus a complementary remainder. Operating directly in function space, jet expansions cut through entan- gled computation in recursive residual nets, and ground mechanistic interpretability techniques in approximation theory. This enables modular inspection: one can pull out paths of interest, e.g., logit lens, n--grmm paths, while bracketing the rest as remainder.
Limitations. Jet expansions are not strict function approximations in the Taylor sense; they rewrite the computation into interpretable polynomial terms plus a remainder. Remainder sizes depend on the jet order lt and weight choices (hyperparameters) and expansions are not unique (higher orders contain lower orders). While graph manipulation is lightweight, systematic evaluation of many (and higher-order) paths can be costly; heuristics or subsampling may be needed for large input spaces. Our n-gram studies focused on bi- and tri-grams; longer-context expansions are left to future work.
Implications and future work. Moving beyond polynomial bases, we envision a Fourier-transform style decomposition as critical step toward controllable LLMs. For example, safety control by filtering out toxic "frequencies". We believe interesting future directions include connecting de- compositions to attribution methods (e.g., the Shapley value), formalizing model equivalence via jet spaces to ground model diffing, and studying the implications of exponential ensembling path growth with model depth, hinted in algorithm[2] We see fruitful links to linear algebraic decomposi- tions and to Markov/HMM viewpoints (c.g., structured decoding (Zhang al. (023)). Practically, beyond longer n- n--rams, we aim to develop dataset-free safety tools. Finally, although our experi- ments are mainly observational, jet expansions may help guide interventions, complementing causal tracing (Meng et al. 2022) and path patching (Goldowsky-Di et al. 2023).
10

---

Published as a conference paper at ICLR 2026
ACKNOWLEDGEMENTS
This work was partially funded under the Horizon Europe grant 101213369 DVPS. This work partially used CPUs from the DiRAC@ Durham facility managed by the Institute for Computa- tional Cosmology on behalf of the STFC DiRAC HPC Facility www.dirac.ac.ukk The equipment was funded by BEIS capital funding via STFC capital grants ST/P002293/11 ST/R002371/1 and ST/S002502/1, Durham University and STFC operations grant ST/R000832/1. DiRAC part of the National c-Infrastructure. Yihong thanks Lena, Karen, Javier, Nicola, William, Xinchi, Keenan, Hue, Yao, Jiayi, and Sohee for their valuable feedback on an earlier version of this work in 2024, as well Yong. Yue, and the Tsinghua FIB Lab for hosting her research visit. Yihong is grateful to her family for their steady physical and emotional support throughout this project, to Bocan for her strength and creativity, and to Feier for her clarity and inspiration revive this work when was close to being set aside.
ETHICS STATEMENT
This work focuses on developing a mathematical framework (JET EXPANSIONS) for analyzing large language models. Our study does not involve human subjects, proprietary or sensitive data, or ex- periments that raise privacy, security, or legal concerns. We acknowledge that interpretability tools may potentially be misused to extract or expose harmful content (e.g., toxic or private knowledge) embedded in pretrained models. We use the public datasets for LLM toxicity research. We empha- size that our intent is to promote transparency, safety, and responsible analysis of LLMs, and we recommend future work carefully consider these implications line with the ICLR Code Ethics.
REPRODUCIBILITY STATEMENT
We have taken steps to ensure the reproducibility of our results. All definitions, assumptions, and theoretical proofs are included in the main text and appendix. Detailed algorithms (Algorithm I Algorithm 2 and mathematical derivations are provided for clarity. Experimental procedures, model families used (GPT-2, GPT-Neo, LLaMA, OLMo), and metrics are described Section [5.2 and Appendix F We will open-source the code implementing JET EXPANSIONS, extracting jet n-grams, and reproducing jet lenses, ensuring that all empirical results reported can be replicated.
LLM USAGE ACKNOWLEDGMENT
We used LLMs to assist with grammar and writing polishing. All equations, analysis, and research contributions are our own.
REFERENCES
Andy Arditi, Oscar Obeso, Aaquib Syed, Daniel Paleka, Nina Panickssery, Wes Gurnee, and Neel Nanda. Refusal in language models is mediated by a single direction. Advances in Neural Infor- mation Processing Systems, 37:1 37:136037-136083 2024.
Yuntao Bai, Andy Jones, Kamal Ndousse, Amanda Askell, Anna Chen, Nova DasSarma, Dawn Drain, Stanislav Fort, Deep Ganguli, Tom Henighan, et al. Training helpful and harmless assistant with reinforcement learning from human feedback. arXiv preprint AXXv:2204.0586.. 2022.
Nora Belrose, Zach Furman, Logan Smith, Danny Halawi, Igor Ostrovsky, Lev McKinney, Stella Biderman, and Jacob Steinhardt. Eliciting latent predictions from transformers with the tuned lens. arXiv preprint arXiv:2303.08 2023.
Nora Belrose, Quintin Pope, Lucia Quirke, Alex Mallen, and Xiaoli Fern. Neural networks learn statistics of increasing complexity. arXiv preprint arXiv:2402.04366 04340220046,..... 2024.
Jesse Bettencourt, Matthew J. Johnson, and David Duvenaud. moylor-mode automatic differentiation for higher-order derivatives JAX. In Program Transformations for ML Workshop at NeurlPS 2019, 2019. URL htt.t:ttt/p//emmmt
3FNFF

---

Published as a conference paper at ICLR 2026
Sid Black, Leo Gao, Phil Wang, Connor Leahy, and Stella Biderman. GPT-Neo: Large Scale Autore- gressive Language Modeling with Mesh-Terme Beesh-Tensorflow, March 2021. URL BtLhttps:: /ot.. If you use this software, please cite it using these metadata.
10.528I/zenodo. zenodo. 5297715
Tolga Bolukbasi, Adam Pearce, Ann Yuan, Andy Coenen, Emily Reif, Fermanda Viégas, and Martin Wattenberg. An interpretability illusion for bert. arXiv preprint vrXiv:2104.071433 2021.
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Miine and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL http: //g.ttubb.cc///gogge////aaa gotp://oooooo//aaaaa/
//github. cot.
Thorsten Brants, Ashok Popat, Peng Xu, Franz Josef Och, and Jeffrey Dean. Large language models in machine translation. In Proceedings of the 2007 Joint Conference on Empirical Methods in Natural Language Processing and Computational Natural Language Learning (EMNLP-CoNLL)) pp. 858 867, 2007.
Trenton Bricken, Adly Templeton, Joshua Batson, Brian Chen, Adam Jermyn, Tom Con- erly, Nick Turner, Cem Anil, Carson Denison, Amanda Askell, Robert Lasenby, Yifan Wu, Shauna Kravec, Nicholas Schiefer, Tim Maxwell, Nicholas Joseph, Zac Hatfield-Dodds, Alex Tamkin, Karina Nguyen, Brayden McLean, Josiah E Burke, Tristan Hume, Shan Carter, Tom Henighan, and Christopher Olah. Towards monosemanticity: Decomposing language models with dictionary learning. Transformer Circuits Thread, 2023. https//transformer- circuits.pub/2001 http:///mooooomattttttttt///e///
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language mod- are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020. URL pttpss:/ neurips...eu.ppp......... cc/ppprr files/paper/2020/ File//15coccceeeeo00/111Baaaaaa2 Pap.r..
Paper.
Nicola Cancedda. Spectral filters, dark signals, and attention sinks, 2024.
Nicholas Carlini, Florian Tramer, Eric Wallace, Matthew Jagielski, Ariel Herbert- Voss, Katherine Lee, Adam Roberts, Tom Brown, Dawn Song, Ulfar Erlingsson, al. Extracting training data from large language models. In 30th USENIX security symposium (USENIX Security 21), pp. 2633 225-,2650440 2021.
Yihong Chen. Structure and Destructure: Dual Forces the Making of Knowledge Engines. PhD thesis, University College London, 2025. URL
wcw. ac uk/id/
eprint/102112917
Donnella S Comeau, Danielle S Bitterman, and Leo Anthony Celi. Preventing unrestricted and un- monitored ai experimentation in healthcare through transparency and accountability. npj Digital Medicine, 8(1 8(1):42, 2025.
Arthur Conmy, Augustine Mavor-Parker, Aengus Lynch, Stefan Heimersheim, and Adrià Garriga- Alonso. Towards automated circuit discovery for mechanistic interpretability. In A. Oh, T. Naumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine (eds.). Advances in Neu- ral Information Processing Systems, volume 36, pp. 16318-16352. Curran Associates, Inc., faarrrrfiees/// http://peeee_ileeeeppp////2///
2023. URL httpt:://pptooooo httpps//pproeeeddns...eeee neuridss ffle23eeeccaeeeaaaaaarraaaaaaeee www.pperrr-oocreeeee... .pf
Charles Ehresmann. Les prolongements d'une variété différentiable: l'espace des jets d' d'ordre r de vn dans vm. C. R. Acad. Sci. Paris, 233:777-779, 1951.
Maha Elbayad, Jiatao Gu, Edouard Grave, and Michael Auli. Depth-adaptive transformer. ICLR, 2020.
Nelson Elhage, Neel Nanda, Catherine Olsson, Tom Henighan, Nicholas Joseph, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Nova DasSarma, Dawn Drain, Deep Ganguli, Zac Hatfield-DD Dodds, Danny Hernandez, Andy Jones, Jackson Kernion. Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and
12

---

Published as a conference paper at ICLR 2026
Chris Olah. A mathematical framework for transformer circuits. Transformer Circuits Thread, circuits. wwrcuits.pubb/22221111
2021. https:/ttransformer- wwtprssttoeeerrrrrrsssssesee
html.
Nelson Elhage, Tristan Hume, Catherine Olsson, Nicholas Schiefer, Tom Henighan, Shauna Kravec, Zac HatfieldDDodds, Robert Lasenby, Dawn Drain, Carol Chen, Roger Grosse, Sam McCandlish, Jared Kaplan, Dario Amodei, Martin Wattenberg, and Christopher Olah. Toy models of superposition. Transformer Circuits Thread, 2022. https://transformer- hwwtii/ppub220022/yyymmmme//eee
Laura Fernández-Becerrr, Miguel Ángel Gonzalez-Santamarta, Ángel Manuel Guerrero-Higuerae Francisco Javier Rodríguez-Lera and Vicente Matellán Olivera. Enhancing trust autonomous agents: An architecture for accountability and explainability through blockchain and large lan- guage models. arXiv preprint (rXiv:2403.09567 2024.
Javier Ferrando and Elena Voita. Information flow routes: Automatically interpreting language models at scale. arXiv preprint vrXiv:2403.00822 2024.
Javier Ferrando, Gabriele Sarti, Arianna Bisazza, and Marta Costa-jussà. A primer on the inner workings of transformer-basee language models. arXiv preprint ar arXiv:2405.00200 2024.
Samuel Gehman, Suchin Gururangan, Maarten Sap, Yejin Choi, and Noah A. Smith. RealToxic- ityPrompts: Evaluating neural toxic degeneration in language models. In Trevor Cohn, Yulan He, and Yang Liu (eds.), Findings of the Association for Computational Linguistics: EMNLP 2020, pp. 3356-3369, Online, November 2020. Association for Computational Linguistics. doi: 10. 10.18853//1/0200 http://////22000.mmmmmmmmm...aaa E Eindings emnlp.... 301
URL httph: hnttps://acleee http://laaanttoogg...o.. wowpthhoooo...rg2222...... org/2020.
Mor Geva, Roei Schuster, Jonathan Berant, and Omer Levy. Transformer feed-forward layers are key-value memories. In Marie-Francine Moens, Xuanjing Huang, Lucia Specia, and Scott Wen-tau Yih (eds.), Proceedings of the 2021 Conference on Empirical Methods Natural Lan- guage Processing. pp. 5484-5495, Online and Punta Cana, Dominican Republic, November 2021. Association for Computational Linguistics. doi: 10. 11253//11222..emmlllll 1e..//2.....mmmaaii.. emnip-main.446 URL https://acl antho ogy- www. 2021 emnle ma i in. 446
Mor Geva, Avi Caciularu, Kevin Wang. and Yoav Goldberg. Transformer feed-forward layers build predictions by promoting concepts in the vocabulary space. In Proceedings of the 2022 Confer- ence on Empirical Methods in Natural Language Processing, pp. 30-45, 2022.
Nicholas Goldowsky-Din Chris MacLeod, Lucas Sato, and Aryaman Arora. Localizing model behavior with path patching. arXiv preprint arXiv:2304.05969, 2023.
Joshua T Goodman. bit of progress in language modeling. Computer Speech & Language, 15(4): 403 434.443,222 2001.
Andreas Griewank and Andrea Walther. Evaluating derivatives: principles and techniques of algo- rithmic differentiation. SIAM, 2008
Dirk Groeneveld, Iz Beltagy, Pete Walsh, Akshita Bhagia, Rodney Kinney, Oyvind Tafjord, Ananya Harsh Jha, Hamish Ivison, Ian Magnusson, Yizhong Wang, et al. Olmo: Accelerating the science of language models. arXiv preprint 1xx:::44420008.3 2024.
Thomas Hartvigsen, Saadia Gabriel, Hamid Palangi, Maarten Sap, Dipankar Ray, and Ece Ka- machine-generated dataset for adversarial and implicit hate speech detection. In Smaranda Muresan, Preslav Nakov, and Aline Villavicencio (eds.), Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 3309-3326, Dublin, Ireland, May 2022. Association for Computational Linguis- 234. URL https: / //ac/aala/ canthelooy.orgoe
mar. ToxiGen: A large-scale machine generated
tics. doi: acl-long.oo 234
Kai He, Rui Mao, Qika Lin, Yucheng Ruan, Xiang Lan, Mengling Feng, and Erik Cambria. A survey of large language models for healthcare: from data, technology, and applications to accountability and ethics. Information Fusion, 118: 102963, 2025.
13

---

Published as a conference paper at ICLR 2026
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp- 770-778, 2016.
Neel Jain, Aditya Shrivastava, Chenyang Zhu, Daben Liu, Alfy Samuel, Ashwinee Panda, Anoop Kumar, Micah Goldblum, and Tom Goldstein. Refusal tokens: A simple way to calibrate refusals in large language models. In Second Conference on Language Modeling.
Stanislaw Jastrzebski, Devansh Arpit, Nicolas Ballas, Vikas Verma, Tong Che, and Yoshua Bengio. Residual connections encourage iterative inference. arXiv preprint arXiv: 2017.
Jiacheng Liu, Sewon Min, Luke Zettlemoyer, Yejin Choi, and Hannanch Hajishirzi. Infini- gram: Scaling unbounded n-gram language models to a trillion tokens. arXiv preprint 1rX1v:2401.17377 2024.
Kevin Meng, David Bau, Alex Andonian, and Yonatan Belinkov. Locating and editing factual associations in gpt. Advances in Neural Information Processing Systems, 35:17359-17372, 2022.
Jack Merullo, Carsten Eickhoff, and Ellie Pavlick. Language models implement simple word2vec- style vector arithmetic. arXiv e-prints, pp. arXiv-2305, 2023.
Eric Mitchell, Charles Lin, Antoine Bosselut, Chelsea Finn, and Christopher Manning. Fast model editing at scale. In International Conference on Learning Representations, 2021.
Eric Mitchell, Charles Lin, Antoine Bosselut, Christopher D Manning, and Chelsea Finn. Memory- based model editing at scale. In International Conference on Machine Learning, pp. 15817- 15831. PMLR, 2022.
Aaron Mueller. Missed causes and ambiguous effects: Counterfactuals pose challenges for inter- preting neural networks. arXiv preprint urXiv:2407.04690 2024.
Timothy Nguyen. Understanding transformers via n-gram statistics.
arXiv preprint
arXiv:2407. /2034, 2024.
nostalgebraist. logit lens on non-gpt2 models + extensions, 2021a. URL https://colab. research. cwm. wtt../o//dvvveeeeeeoeeee http:///w.w..../ooooooooooorotaa
nostalgebraist. interpreting gpt: the logit lens, 2021b. URL httpss////w/w
HBf55bbD7hgcAY2690
Priyadarshini Panda, Abhronil Sengupta, and Kaushik Roy. Conditional deep learning for energy- efficient and enhanced pattern recognition. In 2016 Design, Automation & Test in Europe Con- ference & Exhibition (DATE), pp. 475-480, 2016.
Alethea Power, Yuri Burda, Harri Edwards, Igor Babuschkin, and Vedant Misra. Grokking: Gener- alization beyond overfitting small algorithmic datasets, 2022.
Xiangyu Qi, Ashwinee Panda, Kaifeng Lyu. Xiao Ma. Subhrajit Roy, Ahmad Beirami, Prateck Mittal, and Peter Henderson. Safety alignment should be made more than just a few tokens deep. In The Thirteenth International Conference on Learning Representations.
Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners. 2019.
Emmanuel K Raptis, Athanasios Ch Kapoutsis, and Elias B Kosmatopoulos. Agentic llm-based robotic systems for real-world applications: review on their agentieness and ethics. Frontiers in Robotics and AI, 12:1 1604405 2025.
Baptiste Rozière, Jonas Gehring, Fabian Gloeckle, Sten Sootla, Itai Gat, Xiaoqing Ellen Tan, Yossi Adi, Jingyu Liu, Romain Sauvestre, Tal Remez, Jérémy Rapin, Artyom Kozhevnikov, Ivan Ev- timov, Joanna Bitton, Manish Bhatt. Cristian Canton Ferrer, Aaron Grattafiori, Wenhan Xiong, Alexandre Défossez, Jade Copet, Faisal Azhar, Hugo Touvron, Louis Martin, Nicolas Usunier, Thomas Scialom, and Gabriel Synnaeve. Code llaaa: Open foundation models for code, 2024.
14

---

Published as a conference paper at ICLR 2026
Harshay Shah, Andrew Ilyas, and Aleksander Madry. Decomposing and editing predictions by modeling model computation. arXiv preprint arXiv:2404.11534, 2024.
Claude Elwood Shannon. A mathematical theory of communication. The Bell system technical journal, 2773):379-233, 1948.
Helena Smith. Clinical ai: opacity, accountability, responsibility and liability. Al & Society, 36: 535-545, 2021. doi: 10. 00007/00014-00200000001
Luca Soldaini, Rodney Kinney, Akshita Bhagia, Dustin Schwenk, David Atkinson, Russell Authur, Ben Bogin, Khyathi Chandu, Jennifer Dumas, Yanai Elazar, Valentin Hofmann, Ananya Harsh Jha, Sachin Kumar, Li Lucy, Xinxi Lyu, Nathan Lambert, Ian Magnusson, Jacob Morrison, Niklas Muennighoff, Aakanksha Naik, Crystal Nam, Matthew E. Peters, Abhilasha Ravichander, Kyle Richardson, Zejiang Shen, Emma Strubell, Nishant Subramani, Oyvind Tafjord, Pete Walsh, Luke Zeulemoyer, Noah A. Smith, Hannaneh Hajishirzi, Iz Beltagy, Dirk Groeneveld, Jesse Dodge, and Kyle Lo. Dolma: An Open Corpus of Three Trillion Tokens for Language Model Pretraining Research. arXiv preprint, 2024.
abs/ 2002 00159
Anej Svete and Ryan Cotterell. Transformers can represent n-gram language models. arXiv preprint arXiv:2404. 14994, 2024.
Adly Templeton, Tom Conerly, Jonathan Marcus, Jack Lindsey, Trenton Bricken, Brian Chen, Adam Pearce, Craig Citro, Emmanuel Ameisen, Andy Jones, Hoagy Cunningham, Nicholas L Turner, Callum McDougall, Monte MacDiarmid, C. Daniel Freeman, Theodore R. Sumers, Edward Rees, Joshua Batson, Adam Jermyn, Shan Carter, Chris Olah, and Tom Henighan. Scaling monosemanticity: Extracting interpretable features from claude sonnet. Trans- www.mnttrrrrruusseee pwtpiittsppub/2220 pub/2024/
former Circuits Thread, 2024. URL
scaling www.ggmmooemmmatiiicoooooo wiittniciitytyy http:/ideexxxiee
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXh::2302.13977, 2023,
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Niko- lay Bashlykov, Soumya Batra, Prajiwal Bhargava, Shruti Bhosale, et al. Llama 2: Open founda- tion and fine-tuned chat models. arXiv preprint arXiv.230..09288, 2023b.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural informa- tion processing systems, 30, 2017.
Andreas Veit, Michael J Wilber, and Serge Belongie. Residual networks behave like ensembles of relatively shallow networks. Advances in neural information processing systems, 29, 2016.
Elena Voita, Javier Ferrando, and Christoforos Nalmpantis. Neurons in large language models: Dead, n-gram, positional. In Lun-Wei Ku, Andre Martins, and Vivek Srikumar (eds.). Findings of the Association for Computational Linguistics ACL 2024, pp. 1288-1301, Bangkok, Thailand and virtual meeting, August 2024. Association for Computational Linguistics. URL https: //aclanthology.. ort///gg//2202 inddingsaac.. 75
Sandra Wachter, Brent Mittelstadt, and Luciano Floridi. Transparent, explainable, and accountable ai for robotics. Science robotics, 26))eaan6000, 2017.
Kevin Wang, Alexandre Variengien, Arthur Conmy, Buck Shlegeris, and Jacob Steinhardt. Inter- pretability in the wild: a circuit for indirect object identification in gpt-2 small. arXiv preprint arXiv:2211.00593, 2022.
Xiangxiang Xu, Shao-Lun Huang, Lizhong Zheng, and Gregory W Wornell. An information theo- retic interpretation to deep neural networks. Entropy, 24(1):135, 2022.
Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understanding deep learning (still) requires rethinking generalization. Communications of the ACM, 64(3):107- 115, 2021.
15

---

Published as a conference paper at ICLR 2026
Honghua Zhang, Meihua Dang, Nanyun Peng, and Guy Van den Broeck. Tractable control for au- toregressive language generation. In International Conference on Machine Learning, pp. 40932- 40945. PMLR, 2023.
16

---

Published as a conference paper at ICLR 2026
A JETS AND EXPANSIONS
A jet of function represents an equivalence class. We thus can perform algebraic operations among functional equivalence classes using jet algebra stated below.
Proposition 1 (Jet algebra). Let f.g C (Rd Rd)) and k N+. Then,
(i) (af (z)) (f)(()) + +JJ (g)(o)) a, B (linearity): f(x) (ii) Jd e and f g(u) = J (((o)(( (jet after endomorphisms))
(ii) go 2e 6g u _ f(x)) (endomorphism after jet);
(iv) (zo)) = 1 f(g())
2 (zo) (composition of jets);
Properties (i)-(iii follow directly from the definition; (iv) is a consequence of the chain rule and trum- cation. To reorganize residual computations typically used in LLMs, we rely the disentangling property jets, restated below.
Lemma (Disentanglement of Jets). Let C (RR R.), N. N E N {a) be set of jet base points, and N IN be set of jet weights (i.e., we 0. w, w. = Define the
sum I
Dim Ti and r = max; wi - æ||. Then the k-jet of f at the sum & satisfies
J Eو= mz) Ei wi Jh f1x) + 0(1
Proof of Lemma Take y Rd. N 2 1, the set of jet base points æ; E R for i E [N] jet weights E N-1 and an order 2 Since belongs to the simplex A N-1, we have Σ w. 1. Multiplying f(y) on both hands, we obtain
ΣN w.f(y)= f(y).
Applying eq. (3) (Taylor expansion) and the definition of jet with each n; as the center, the left hand side (LHS) becomes
N
N
Σ
m fyy
Σ
wi
f(x,) + Σ D Def(zz)( 00 +Ol1y 2: 2:11(t2)
N
Σ
wh
((x) + Oww:lwwwl 11)
(8)
At the same time, we can expand f(y) at the right hand side (RHS) with Σ 2; as the center
(()) J a.) Σ 1111
Now let take y ΣN 3y and observe that theremainder RHS vanishes Ollly- -x 0 and the remainder at LHS (w,ily =1|l*+1) Ow || E3 11)). Finally we observe that the class of functions in the last O are dominated by the class of function in O(r where r max; [w ||= Σ, 1 is the maximum remainder. This concludes the proof.
As side note, jet weights would not need to form convex combinations, but rather linear combina- tions] E, w;- = 1. However, restricting to convex combinations has two major advantages:
optimizing over a convex set guarantees the existence of maxima and minima (Weierstrass theorem) and uniqueness of minima if we optimizing strictly convex loss in general is the case for expansions that only affect the decoder module.
weights within the probability simplex have a clearer interpretation for interpretability pur- poses.
17

---

Published as a conference paper at ICLR 2026
B REMAINDER SIZE OF JET EXPANSIONS
JET EXPANSIONS does not aim to provide approximation guarantees. Instead, in the following we clarify when small remainders expected (input-specific cases) and when they not (function- level expansions), but the expansions still yield meaningful, interpretable insights.
In input-specific evaluations (where we choose specific input sentences like "new simple neural architecture, the Transformer"), we expect reminders to be small, since we want to draw specific conclusions about how model is behaving lenses experiments of $5.2.1) we do provide empirical studies, and we often find that the remainder ò is small. In order measure the remainder between the expansion logit and the model logit, we compute the cosine similarity between them, since direct difference depends heavily on the input sentences (See App. [F] "Cosine similarity as remainder metrie."). In short, we use cosine similarity as an indicator for checking if the remainder is small. Most of our heatmaps in the submission include this similarity measures the parenthesis e.g. "Expan. (0,993)". We also compute average similarities over 100 examples and summarize the results in Figure 4 (bottom): expansion logits are frequently close to the model outputs, with cosine similarities in the 0.85-1 range, indicating that the extracted components capture most of the behavior for many concrete inputs i.c. smaller remainders. In general, the remainder size will depend on three factors:
particular input. In these cases (i.e. in the Jet
1. the type of jet expansion one performs (e.g. iterative vs joint);
2. the non-linearity on which the jet is applied to (e.g. ReLu vs ELu, vs LayerNorm)
3. how far the variate is from the base point, which in turn fully depends on the chosen input sentence (see Equation
In funtiionn-eee expansions (e.g.. around the embedding function Enc for extracting jet bi- bi-grams), remainders can be large and this is fully expected. Here. the goal is not approximation accuracy but to decompose the computation into interpretable paths. Concretely, let us consider bi-grams ($5.2.2) the model performs far more than the isolated mechanism of predicting a token based solely on the do expect reminders between the model output and the extracted bi-
previous one. Therefore gram paths to be naturally large. Yet the extracted paths remain meaningful precisely because they isolate coherent part of the computation, even if they explain only small fraction of the total behavior. A large remainder therefore does not invalidate the interpretation; it simply reflects that focusing one specific, coherent part of the computation. This is a common situation across interpretability research, or in general when humans try to explain things. Conceptually, find it useful to compare jet expansions to Fourier transform: even when the Fourier transform captures only part of the signal's is spectrum and the resulting approximation error cannot be directly measured, the partial frequency information provides can still be valuable for understanding the signal and diagnosing issues in the generating circuits.
C EXPONENTIAL EXPANSION AND ENSEMBLING PERSPECTIVE ON LLM COMPUTATION
Algorithm recursively expands all residual blocks. We see it the maximal expansion (the upper limit of cxpansion) of JET EXPANSIONS. It aims to demonstrate that JET EXPANSIONS also provides theoretical grounds for Veit al. (2016), whose analysis at the time lacked theoretical foundations and was presented heuristically and empirically, by manually selecting paths (i.c., 24 gradient paths) for analyzing model behavior, where the exponentially many ensemble structures were syntactically noted rather than analytically derived.
With closer look, we were surprised to find that several empirical procedures used in Veit et al. (2016) are exactly recoverable as specific instances of jet expansions. For example, the deletion of a module (Sec. 4.1 /eiit et al. (2016) corresponds precisely to
jet eet.exaand (4-2, 0),
ie., JET EXPANSIONS with the skip-layer upstream stream the base point and with jet order 0.
Another example is their gradient analysis. In the original exposition, the procedure is described operationally as: "To sample path of length k, we first feed batch forward through the whole
18

---

Published as a conference paper at ICLR 2026
netwrrk. During the backward pass, randomly sample lo residual blocks. For those k blocks, we only propagate through the residual module; for the remaining n blocks, we only propa- gate through the skip connection." " The outcome of these operational steps is equal to taking the derivative component of a first-order jet expansion over each block and evaluate it over a batch of inputs, shown below. Note that the following reinterpretation focus the gradients with respect to intermediate representations while empirically collecting gradients with respect to parametes is much easier.
Path selection in backpropagation as first-order JET EXPANSIONS.
For residual block,
3: (1v-11) (id 7n h-11 h111 + (hi-11),
the derivative satisfies D I + D' [1]] Thus, propagating gradients through the full network f = BL B .... o $1 yields the Jacobian using the chain rule
Df(zo) I (I D=())-ll).
On the other hand, the first-order jet of residual block h4-1
(b (3:(hu-1), 1 + Da((u-il).
where we use the pair notation for the jet, the first entry being the zero-th order term and the second entry being the first order term. Since jets compose according to
J' (x) = (g(x)) J (')) (Proposition (imiiionnvv)))) (iv),,
we have the full first-order jet of f as
(z) = J (' / o....
(81)
Extracting the derivative component of this composite jet therefore yields the ordered product of the derivative factors from each block:
Df(=o) = (1 + Di o (1 D71-1(hu-21) o ... o (1 + D' (=o))
Composition of first-order jets multiplies the linear parts in this precise order:
Df( (I Dyu(he-il) (I+ Dy1 (mol).
Expanding this product produces 24 additive terms:
(I [Dnlttull)
Df(xo) = E
SC(I....L) VES
each corresponding to a distinct choice of either I or Dy at every block. These terms match exactly, one-for-one, the "gradient paths" enumerated by Veit al. (2016). Their operational procedure ("selecting" residual or skip gradients per block) amounts to selecting individual terms from this expansion. Consequently, their "gradient paths" are not independent computational input-output paths through the nonlinear network, but the combinatorial derivative components of the first-order jet (i.e., the first order term) of the residual network.
D RUNTIME OF JET EXPANSIONS
We report in fig. plot of the runtime for evaluating expansions originating from the joint jet lenses of section 5.2.1 as a ratio of the input model evaluation (forward pass), for both the uniform and the optimized jet weights w setup, for different jet orders k.
E JET n-GRAMS AND THEIR ALGORITHMS
General concept of n-gram models The general concept of n-gram models linked to counsormer-bassee language models involves defining or constructing mappings that functionally depend only
input tokens (with the n-th token being the output token) to capture and de- scribe the behaviour of the original language models. We are not the first to explore this idea; for instance Nguyen (2024) fits n-grams the same dataset used to train the language models.
19

---

Published as a conference paper at ICLR 2026
Runtime/origina forward runtim
Average runtimes ratios (# std)
GPT2 optimized)
GPT2 (w = average)
15
GPT2-large (w optimized)
GPT2-large (w average)
GPT-Neo 2.7B (w optimized)
5
GPT-Neo 2.7B (w average)
Jet order (k)
Figure 5: Empirical runtime of evaluations of JET EXPANSIONS originating form the joint jet lenses as ratio of the evaluation of the input model.
JET EXPANSIONS for in-model n-grams JET EXPANSIONS allow to define n-grams statistics that are derived solely and directly from the model itself producing in-model n-grams rather than in-data n-grams. This approach offers at least two significant advantages:
No repeated inference runs over prepared datasets: It removes the need to prepare datasets for prompting LLMs, thereby avoiding repeated inference runs to collect activation patterns for interpretability analysis and reducing computational overhead. With small n, direct expansion LLMs into n-grams can be performed on CPUs, which are roughly an order of magnitude less expensive than GPUs.
Avoidance of fitting artifacts: It avoids potential artifacts that could arise from the selee- tion of external n-gram fitting methods.
We describe the detailed relationship between JET EXPANSIONS and bi- bi-grams/tri--ran which we used in studies. We will release code for these procedures and also provide equivalent algorithms that directly use transformer modules.
Jet bi-grams Jet bi-grams are paths that do not pass through self-attention layers. In experiments. we focus on two types of bi-gram paths. a) the embedding-uneni unembeddding path that can be obtained jet_exxanddff... L, {Enc), 0). b) paths that pass through MLP module, assuming MLPs are odd block indices in the residual network architecture, the procedure to extract the path is described from line 3 line 7 in Algorithm 3 This procedure results in series of functions in one for each MLP layer, that depend only on the last input token. Applying softmax normalization to their logit output allows these functions to define (conditional) bi-grams.
Jet tri-grams Jet tri- tri-grams involve paths that pass through least one self-attention layer, with need to isolate the contribution from the first token of the tri-gram. The procedure for extracting 0-th order jet trigram path that passes through the ith self-attention layer (assuming it has one head) is described from line I to line 5 in Algorithm 4 This procedure yields map that depends only on two input tokens, isolating the contribution of the ith self-attention layer on pairs of tokens. Once softmax normalization is applied, this defines tri-gram. The tri- tri-gram& could represent either skip trigram or a contiguous trigram, depending on how positional information is encoded (e.g., absolute positional embeddings versus rotary embeddings).
20

---

Published as a conference paper at ICLR 2026
Algorithm 3 Dataset-free extraction of jet bi-grams (MLP paths).
Require: Model f. total blocks L, vocabulary V
1: // Initialize with encoder path
2: C < [Encc
3: // Iterate over MLP blocks (odd indices)
4: for | 1,3, ...,..... do
5: (5i, ôr) < jet.expand (f, {Enc), 0)
6:
// Collect each expanded MLP term
7:
C C Cu(((()
e * Enc}
8: // Expand over the decoder
9: (E, 8) - jet_expand f, L 1,)
// Evaluate over the input space
11: for all do
12: for all x V do
13: for all i V do Record (I, i, e()) as the bi-gram score
14:
15: return Symbolic bi-gram table {<z, i, e(z)(i)))
Algorithm Dataset- extraction jet tri-grams (Attention paths).
Require: Model f, total blocks L, target self-attention block index I, vocabulary V 1: C (Enc (x))
2: // Expand over the attention
3: 8) + jf
4: C (e(;,) & E # Enc}
// Expand over the decoder
6: (5,8) jet_expand f, L 1, 0,)
7: // Evaluate over the input space
8: for all E E E do
9: for all (I, E 2 do
10: for all E do Record (It-1, Ie, i, e( (.)(() the tri-gram score
11:
return Symbolic tri-gram table (1 2((r,,
Practical computation formal jet n-grams. Algorithms and E express bi-gram and tri- gram extraction the jet-expansion formalism: we expand intermediate paths using ret_expand, collect the resulting functions into and finally obtain family of zeroth-order jet paths
C E,
e: e: RV
where the arity of e (one token for MLP paths, two tokens for ingle-attention paths) determines whether the path represents bi-gram (n.
tri-gram (n 3). Formally, n-gram table obtained by exhaustive evaluation of e on its domain: recording e(x)[i] for x E in the bi- gram case, e(2,-1, Ix] for (x,11,2,) E y2 in the tri-gram case. In practical implementations. these evaluations correspond to batched matrix computations over the entire vocabulary. For bi- grams, the embedding matrix E E RV xd serves as the batch of all input tokens, and each MLP component represented in C acts E via matrix multiplications; after layer normalization and
21

---

Published as a conference paper at ICLR 2026
projection through the unembedding matrix U, the resulting matrix has (æ. entry equal to [i] For tri- tri-grams, an analogous "pairwise embedding" tensor encodes all (zc-1, 24) pairs at once, and the attention and subsequent linear operators act this tensor in batch; projection through U yields a three-dimensional array whose (I4-1,, T1, i) entry is exactly e(x)-1,2;)[1]. Thus, the formal jet description and the practical batched-matrix implementation are two views of the same operation: jet path defines the atomic paths {e, E {} which simplifies the large model, and the bi-gram table obtained by evaluating e simultaneously all vocabulary elements using matrix multiplications.
F EXPERIMENTAL METRICS
In this section, we detail the measured quantities in each empirical case study.
Cosine similarity as remainder metric. In Section 5.2.1 (the lens experiments), we need comparable metric for quantifying the size of the remainder term in the JET EXPANSIONS for a given input sentence. Let m R denote the model output logits over the vocabulary V, produced by a full forward pass, and let E RV denote the logits predicted by the truncated jet expansion. A naive approach would define the remainder as the difference vector r = m and measure mag- nitude using |||]]. However, this direct norm-based measurement is highly sensitive to input-specific variation, such sequence length, and the internal activation scaling of the model. As a result, the magnitude ||mm ell can be dominated by variations in logit norm rather than reflecting the intrinsic approximation error introduced by JET EXPANSIONS. To address this, we adopt cosine similarity ocale-invarian measure of alignment between the model logits and their JET EXPANSIONS approximation. Formally, we compute
m
cos(e, m)
||e||| m
where denotes the standard dot product. A cosine similarity of means that JET EXPANSIONS preserves the model logits structurally. Conversely, lower cosine similarity values correspond to larger remainder term in the expansion. Practically, cosine similarity enables us to disentangle re- mainder size from input-dependent logit scaling, offering a more interpretable and stable measure of the fidelity of JET EXPANSIONS. We therefore report cosine similarity throughout our experiments as our primary metric for assessing remainder size.
A Logit after intervention. In Section 1 to compute A logits, we calculate the logits for the given n-gram both before and after applying the intervention, then determine the change in the logits. For example, consider the trigram (Lemma, let, s). We compute the logit of "s" conditioned on the input "Lemma let". The intervention involves removing the corresponding attention head (e.g., head 2). We then measure and report the change in the logit for "s" as a result of this intervention.
Jet bi-gram comparison for code fine-tuning. In Section 5.2.2, we derive the top 1000 bigrams using Algorithm 3 These bigrams are then saved, for example, CSV files, enabling the inspec- tion and comparison of models via their respective bi-grams. This approach allows us to bypass the challenges of comparing models in the high-dimensional parameter space, where measuring behavioral-level differences can be difficult. We have developed a web UI demonstration where users can perform "model diffs" using the respective jet bi-grams. For example, Figure 6 demon- strates how this UI can be used to compare the base Llama-2-7B model with its coding fine-tuned versions.
we introduce a method to quantify the possession of toxic knowledge. We compute jet bigram probability scores and calculate the cumulative conditional probability mass over a curated set of toxic bigrams, pairs of tokens specifically linked to toxic meanings in predefined word list. The toxic mass (M) is formally defined as the sum of these conditional probabilities across the query set (Q):
Jet bi-gram toxic mass. In Section 5.2.2
M Σ P(22)))) ZEQ
Here, Q represents the query set comprising the toxic bigrams derived from the word list. In this way, we can measure model toxicity using simple query words instead of relying on extensive, curated prompting datasets.
22

---

Published as a conference paper at ICLR 2026
LLM Diff
1odel
Model 2
llamal_78
Ilamal_7B
llama2_78
llama2_7B
Hama2_7B_chie
llama2_70 @amaa_7B_chat
llama2_130
llama2_130
llama2_ 70B
llama2_700 70B
llamacode 7B
llamacode
llamacode 7B_ instruct
llamacode 78 iinstucct
llamacode 7B_ _B_ptthon
llamacode 78 78_python
which level do you want to compare Ilama2_78 and lamaaode_78_ppyteee 78_python?
cond2gram
You selected: 2gram
2gram for 1lama2_78
2gram for Ilamacode, 78_python
ngram
token_1
token_2
tol
ngram
token_1
token 2
toke
925
anarchy
archy
925
editors
edit
926
longstanding
long standing
postgres
post
gres
927
clientsele
clients ele
conditionals
condition
als
928
ascribed
as cribed
929
Millenn
Mill enn
orElse
Else
930
in tum
in turn
930
-//-///
../..
931
Acceptance
Accept ance
931
thematic
them
932
goose
ose
932
customized
custom
933
directors
direct
ors
933
evenings
even
ings
934
Summary
Sum
mary
934
Bindings
Bind
ings
Figure 6:
UI for running LLM Diff with jet n-grams.
23

---

Published as a conference paper at ICLR 2026
One-to-one bi-grams like and many-to-many bi-grams. In Section analyze the pretraining dynamics by checking the learning speed of bi Di-grams from different categories. One-to-one bi bi-grams (approximately) unimodal Di-grams that concentrate all mass (2(22221) R and given 22, (= (2)) R 1 for a specific pair of token and close to O for all others. In the example in the paper, 21 "&", and 22 "amp". PD is the probability distribution induced by the pre-training data. Many-to-many bi-grams we refer to the opposite scenario where both the conditional probabilities are highly multi-modal. In the example _1 "make" and 22 "sure" have that many other tokens can succeed 2 = ""make" or precede = ="sure".
single token: i.c. given 21,
Hit ratios of bi-grams. The Hit Ratio (HR@n) often referred as hit rate, metric commonly used in ranking tasks. In our context of Section we treat each checkpoint of the language model as "ranker" bi-grams. The Hit Ratio measures how effectively the current model checkpoint retrieves high-quality bi-grams from the set of all possible bi-grams. To quantify the model's progress, we define the bi-grams at the final step as the "good" bi-grams and measure how quickly the model approaches these high-quality bi-grams. Specifically, compute the HR @n to evaluate how often the model's output bi-grams match those in the "true" top n ranked bi-grams given by the final step. Formally, the Hit Ratio @ n is given by
HR @n = 1 Σ Ihe i-th bi-gram output by the current model E True_Top.n)
where is the number of top predictions being considered and
is the indicator function that returns the i-th bi-gram output by the model is present in the True Top bi-grams, and 0 otherwise,
_op__op_n n represents the set of "good" bi-grams, which our case is the set the top n scoring bi- i-garmm from the final model step.
Total mass of bi-grams. In Section we use the total mass as a metric to measure the cumulative probabilities of bi-grams from the top 1K bi-grams, weighted by empirical uni-gram distribution derived from real data. Formally, given by: Total Mass Σ (a PP,
where:
is the embedding-unemmedddin path at the (-th pre-training step.
(z1, 22) are the bi- bi-grams being considered,
Pe, (22(21) is the probability assigned by the model e, (the embedding- path) for the token 22 given token Z1.
PP((2) is the probability of 21 under the empirical distribution D, which the uni-gram probability given by the Infini-gram API (Liu et al., 2024) on the Dolma dataset (Soldaini al. 2024) (the dataset used to pretrain the model checkpoints).
This metric is designed to evaluate how much "correc" probability mass the model checkpoints assign bi-grams (21, 22), taking into account the empirical uni-gram probability of 21. It provides insight into how well the model aligns with the empirical distribution real-world data during the pretraining process.
G CASE STUDY 1: ANALYZING LLM INTERNALS WITH JET LENS AND JET PATHS (ADDITIONAL RESULTS)
G.1 ADDITIONAL PLOTS OF JET LENSES
See plots Figures 7 to 17 The details for obtaining the jet lens plots be found in Note that for iterative lenses the last block coincides with the model logits for all by design. We omit the iterative lens for large for k due to low cosine similarity.
24

---

Published as a conference paper at ICLR 2026
Bock
Supportess
Suggortars
Susporters
Supgertars
Susporters
Supsortars
Susportars
Supportars
Moct
Suppeelen lees
Kapeiters
Iupporters
Suppoiters
Suppieters
Sapprteers
Supperies
Vatporieess
BockE
SuEportors
Subcorbers
Supporters
Subcortars
Sunporters
Siboortars
Supporters
Sap0orters
Sisportes
Sugcortars
Supportors
Scocorters
Surporters
Subcortars
Susporters
5.psortars
Bock
Supppr
Suppenters
Suppar
Suggurtera
Supporters
Supportara
Suppae
Supourtera
Roct
S.Epprie's
Vincoleers
Sisporters
Sincorters
Supporters
V.Rcorters
S.Doorfers
BOCK'
sippprters
Mipporters
Sipporters
Sinporters
Supporters
5ipporters
supporters
S-pooreers
Supporters
Supporters
Supporters
Susperters
Supparters
Supporters
Susporters
Supporters
Supportess
Sipporters
Kappnierr
Canonfees
Simpreen
B804 13
Sueporte's
Sasporters
Supcortars
Susporters
Supcorters
Susporters
5.pcortara
BloCK L4
turpportors
Supgortars
Supporters
Supgortars
Sasporters
Supsortars
Susporters
Supportarrs
Supppe lemn
Supponters
Supparten
Vapperfers
Supporten
Supporten
Sapporters
Sugcorters
Supportors
Supcortors
Supporters
Supcorters
Susporters
S-pcortars
Block 18
Suggeiters
Suspurlers
tapportera
Supperera
Suspie lars
Sustortara
Block
Supppe lar'
Suppuiters
Supporter lers
Supcurters
bupperters
Seaperters
Susporters
Savvrrkers
BIG0220
SiEpactes
Sincorers
Sipportess
Sincoifers
Sipporter
5.ncorters
Kippoifers
Cappnifars
Caponfers
kasporiess
Ganerfees
Binalreless
Manelhales
B100K 25
S.6p be
Supgortars
Supportors
Supcortars
Attempgs
Moncchdes
susporteis
Introduce
papportorr
Suggortars
Supportors
Suppertars
Nonethidas
Introduce3
25
letemats
Nonethaless
Intrdducdd
B120 25
Foressen
Suggorters
sasporters
Suncortars
Faeesenn
Nonethars Bes
Charpss
Infodured
Block 30
Tcressen Support
Supperters
Suppor lara
Allampia
fonsseee
Charges
Introdured
Supporters
Suppa
amirr
[omme
Lages
nchlor
which
nOURI
fomer
Figure 7: Iterative jet lens (A: 0), equivalent to logit lens (nostalgebraisi, 202 applied over GPT-Neo-2.7D with the input sentence "new simple neural architecture, the Translormer"
beboork
Bockz
Supporters
Ton
network
70r
fomer
Bock
Suppor tara
ton
network
for
which
1rak
fomme
Bock
Supporters
ton
network
fomme
Supporters
Bock7
turporters
network
which
femmr
Tapppeeeers
1m
setuork
107
whin
1rd
Fomme
Bocks
Suppprtorr
100
behanrk
whre
fomee
Biock
partorsors
WSS 17
Supporr
swtwork
wheh
Forme
S.ppprrters
10
bchwork
167
which
1r2
Former
BIOOK
Sirpprtoors
100
Jehhork
201
112
Fommor
15
Support
network
Supporters
calbork
Black
Su bera
network
10r
which
1ra
Fermar
Block
Suppoorioo bers
Ton
nebwork
061
which
112
fomer
Imm
nehoork
whth
fomme
porttrs
beboork
whic
satwork
whi2
[orme
Bi00l-24
Supppe lers
Im
satwork
whihh
fomme
BI00 25
SiEportors
100
naboork
M0.00
112
Gomer
Supporters
Bicck-28
Supporters
network
whih
fomeer
MI02 25
Foreseen
Im
betwork
whinh
478
fomee
Block 30
foresean
101
network
for
Mhan
162
fomm
Supporters
network
sehior
fommer
Figure 8: Iterative jet lens
1). applied over GPT-Neo-2.77 with the input sentence "new simple neural architecture, the Transformer".
25

---

Published as a conference paper at ICLR 2026
neumi
architecbure
2014
Fommr
Bock2
the
network
502
300
femmr
105
107
C20
Tomee
The
Reboork
Comy
fomm
Actoork
Icural
Wook
1ie
Satsork
forme
Bock
Uie
"elwork
The
fomee
Bock'
1h2
nchlow
100
Tomer
Uhe
network
Torma
the
setwork
Steccure
707
fomee
Block
the
hatwork
for
the
Doura]
fomer
005
T00771
fomm
Block 18
th2
network
307
18
fomer
Bisck:
Uie
setwork
10r
The
Teund
Torme
Block
selwork
Teural
fomer
Block
the
network
meural
femmr
the
retwork
The
beum
Fommer
Blo0k
setwork
neum
Gome
network
noural
forma
1104
B104 25
th2
retwork
007
The
ncural
femmr
B100K 30
1hC
nebrork
707
.300
Jeura
fomer
neural
Loges
natwork
nourol
Tommor
Figure 9: Iterative jet lens (k = 2), applied over GPT-Neo-2.7B with the input sentence "new simple neural architecture, the Transformer"
DOW
Gmmpl
Tormer
Simple
smmpe
Bocks
85
way
3nd
1r2
1ont
183n
Bocks
JW3y
.710
1901
Prime
Mockh
emough
erouTe
Bock
enough
hoari
archiccture
whch
Pi
Block 10
enough
network
archeecuure
whkh
frsk
16m6
Ruvelution
ennurg
seboork
model
2010
1601
ennughh
salbork
model
P70
wnnugh
selwork
model
whih
073
Pra
Block
nchoor
moda
which
076
1ont
Retol.t.on
Systom
Prame
made
J0h10
.070
1e1t
Frmm
Block
systam
moddl
which
only
lem
mm
network
Block
nebwork
model
which
Ar2
Pimms
1m
selwork
model
whith
only
mm
105
befoork
whih
Frmm
Jecwook
Jwh/ch
Block
swtwork
whah
neural
b2hoork
which
neura
lent
B000
__chookk
Jeura
lent
helwwrkk
calborrk
B12k 30
Figure Iterative jet lens 0), equivalent Logit Lens (nostalgebraist with the input sentence "new simple neural architecture, the Transtormer".
2021b). applied over GPT-2-large
26

---

Published as a conference paper at ICLR 2026
REW
-impl
TTans
Tormer
Bock
betwork
6003
Hours
BocEZ
36
Bock'
nebwork
nouro
Mort
neltort
wi125
Bock?
10
tebwork
whish
Teural
betwork
whi05
(eum
Block
_eebork
200m
celwork
Block
retwokk
Teural
Block
Rebwork
.0nd
noura
nelwork
teunal
Block
setwork
Teural
Block
tetwork
ural
teuid
schoor
neum
selbork
Block
nabrork
and
neuro
schlork
3nd
Zame
Block
nelwork
and
teunal
relwork
leural
Loges
setoork
.9h
Figure 11: Iterative jet lens (k 1). applied over GPT-2-large with the input sentence "new simple neural architecture, the Transformer"
J976
wieh A370
MLDPO
(KSP))
www.b 1980
Rase (L0N))
Kow (5545)
Bax 39135-990
(9879)
8.--990 (999)
3900 1172.00
WhoN DROWN
7mh 09
Fax Loghz
12/8
Expan. (1.000)
whrch
Figure 12: Joint jet lens with learnable weightings 0), applied over GPT2 with the input sentence "new simple neural architecture, the Transformer"
J1724
24
(20%
1394
00%
(37)
www TARPE
Jaraa BANN
Jeac DAPA
Mivel (177))
FR sproa Taar
3008 2140
weoring 1i FPN
JRW Rorze
Swo 2117 Wi0
www.
New 117
2300 2112000
STUDEN
JN BILLL
Logit
hcbwork
Biat
widc
hears
IC1K
Expen: (1.000)
cand
Figure 13: Joint lens with learnable weightings (k 1), applied simple neural architecture, the Transformer"
GPT2 with the input sentence "new
27

---

Published as a conference paper at ICLR 2026
130%
Bin 10
Don 121/79
Toeer 2211
SIeE T1ST
Jeb 2014
26 201
Peeilam
24 Drnaam
Iwe 2310
Book 241
CEXS
Jonen XTM
JEIN
.15000
Je PMM
J0133
2011
Jone 11NN
2PM
I0 PM
JPHA D1IN
Logis
16
witcc
Detral
Tomee
Erpan. (0.57))
the
363
157
The
Figure Joint jet lens with learnable weightings (k = 0), applied over GPT-Neo-2.77 input sentence "new simple neural architecture, the Transformer"
Baa
37 Bamn
L8S0 11.09M
Ta USCN
10300
1aa 2042
Teaaees 3E94
1.39%
EBacs 1377%
D10gH MMM
10.00
0.60
-0.0%
Boek (9419)
Irseend
6579%
K.1M%
5,990
Boo 5:221
Nav 9944
196
tr 11197
read 1770
1.99
1804 11.00
11,000
0,900
200 1229
1227%
6.00
1.50%
E,900
107%
(500 20
10.000
50 0.000
PONTT DWN
CILAEN
Iiteek 1.80
K35
are 2.00
New 279
154 1:200
1ife117377 11.9%
J00 7477
Jeae (777)
1ew 7394
KIN
1.100
to 1999
1500000
Nerer 199)
8b0 11.00M
frw 237M
Raak 197744
mow
I- BrW
500 1:20
mhot 7137
7am 117%
Moot 1.0%
7374
Logits
wiich
nerral
fomeer
Eeaan (0.533
f5r
Witch
Figure Joint jet lens with learnable weightings "new simple neural architecture, the Transformer"
1). applied over GPT-Neo-2.7 input sentence
28

---

Published as a conference paper at ICLR 2026
Now
18:00
15.00
ES6. DIM
01.0%
Jes 2214
2007710
190
3001
DIEN
18:00
10.80
2000 410
11.0%
2110
2015
WiN
Boa 121 10000
15.9%)
Boo 1201
11.8%
25 201
Logts
neara
Expon, (03560
367
DAUSS
Figure 16: Joint jet lens with learnable weightings = 0), applied over GPT-2-large with the input sentence "new simple neural architecture, the Transformer"
JWY84
Semer
160280
Nhvt 5.0%
10GN
1079
1999
EBACE 2aWw
LOOKK 1.00
JEPR
14229
10000
LAVA
1208
210
2000 7200
CRAPe
CRa70
LENM
13.99
DATR51
100%
(8779)
1324
www 2110
2390
.0
10.9
zeml
11400
1ite 2.29%
0174
219
Bow 9044
read 1.99)
18 1999
100 1.00
12.29
10370
1278
2em 1370
CR370
1.99
Jaw 11710
(500 210821133 201
16951
N49 TLWN
1a 5.000
Rani
2.399
41.00
Reway AAN)
16429
Nas 1110
13300
100 5:-1
1010
200
4E370
10725
HOOt KIEN
39553
Fam 5376
CEAW
Book 0.004
16395)
5.9%
Logis
wiicch
DeiFml
Expon, (0354))
Figure 17: Joint jet lens with learnable weightings (k = 1), applied over GPT-2-large with the input sentence "new simple neural architecture, the Transformer"
29

---

Published as a conference paper at ICLR 2026
Table 5: Several attention heads in the first residual block of OLMo-7B and their roles identified with jet tri- grams extracted from corresponding jet paths. also include an example tri-gram captured by each head.
Head Index
16
26
MacXLaTeX (Jeerna, _let, .1 (ofor, .056, puurposes)
purposes"
composition (20, 23, -1
accoum/cossieratiinn
Example 3-grom
(_inte, _ccoout, gossible)
Alogit after intervention
0.1570
-00019
0.0093
0.0001
G.2 ADDITIONAL TABLES OF JET PATHS OF INDIVIDUAL COMPONENTS
5 reports role identification study on attention heads in the first self-attention of OLMo-7B using jet tri-grams. Specifically, find heads associated with math and programming, e.g. head on Math/Latex; heads promoting digits and dash composition into dates, e.g. head 25; and heads constituting phrase templates, e.g. head 15 managing "for purposes", where is placeholder. To verify the roles we revealed, we further perform preliminary intervention experiments where we ablate MLPs attention heads and compute variations in model logits. After the interventions, the logits drop consistently in all cases, suggesting our jet n-grams indeed can help identify certain roles for selected components. Varying impact on logit differences is likely due to overdetermination (Mueller (004) and our partial selection of jet paths (e.g. for tri-grams we only selected encoding- ceention-decoding paths, excluding any MLP).
H CASE STUDY ASSESSING FINETUNING EFFECTS (ADDITIONAL RESULTS ON RLHE TOXICITY AND REFUSAL MASS)
30

---

Published as a conference paper at ICLR 2026
Figure 18: Refusal v.s Toxic bi-gram mass distribution across MLP paths.
natural question is whether the toxic bigram mass changes significantly in the more robustly tuned models introduced in (Qi et al.). We thank the reviewers for the reference and the suggestion of this experiment. Our prior findings suggest that existing alignment methods tend mask toxicity rather than eliminate it at the representation level. Here, we further investigate this phenomenon. Figure [8] summarize the results of Experiment and Experiment below. The orange bars indicate the refusal mass, while the blue bars indicate the toxic mass. For each layer, percentage change is computed Mare - Mare where Mang and Mchat denote the bigram mass from the augmented and chat models, Mar
as:
respectively. Positive values indicate an increase in mass.
(the open- sourced weights from Qi al.) with the RLHF model L. Llamm2-73-Chat. Toxic bigram mass is measured in two ways: (1) independently for each Enc-(MLP)- path (P1), and (2) jointly by averaging contributions across paths (P2), which is the metric reported in our main submission.
Experiment 1: Toxic Bigram Mass. We compare -
Under the joint metric (P2), we observe no substantial change in total toxic bigram mass. The augmented model reduces toxie mass by only 0.05% compared to I.1 2
We further analyze the distribution of toxic mass across individual MLP blocks. For most layers, the relative change Mag Mad ), falls between -0.5% and -2%, with only few layers exhibiting Mas
reductions greater than 3%. These results suggest that the augmented model does not dramatically reduce toxic bigram mass, but instead reshapes its block-wise distribution.
Experiment 2: Refusal Mass. If the additional fine-tuning does not substantially reduce toxic mass, what accounts for the marked behavioral change (i.e., safer behavior appearing deeper into the response Llama2-7B- P77B--hht--Aus Augm--tnggt)) We investigate whether behavioral differences arise from changes in refusal-related patterns. Using the same pipeline, compute bigram mass for refusal phrases (textual patterns commonly used to reject unsafe queries) (Jain et al.) Arditi et al.) 2024). We find that L.1 2- 7B-Chat Augatntteeeeeeeeeeeee exhibits increased refusal mass in mid-to- late MLP layers relative to L1 ama2 7----BC-- Chat. Notably, the final layer (MLP 32) shows a +8.43% increase in refusal mass, likely reflecting its proximity the output layer.
Discussion. Together, these results suggest that the additional fine-tuning does not directly subtract toxic knowledge mass. Instead, appears to strengthen representations refusal phrases, thereby shifting texicity-related behavior deeper into the generation process.
This interpretation aligns with the fine-tuning procedure described in QQi et al.), which augments the corpus with examples containing "Safety Recovery" segments embedded within responses. We speculate that the observed increase in refusal mass emerges directly from this augmentation process, while the modest decrease in toxic mass may be indirect consequence of adapting model statistics to the new data. Overall, the augmentation procedure effectively alters model behavior, but likely through reinforcement of refusal patterns rather than explicit removal of toxic knowledge.
31

---

Published as a conference paper at ICLR 2026
1 CASE STUDY 3: TRACING PRETRAINING DYNAMICS WITH JET BI-GRAMS
Pretraining an LLM is usually extremely resource intensive. Therefore it is crucial to monitor the progress pretraining run to prevent wasting of time and compute. In this section, we show how jet bi-grams serve
effective signaling tool to trace the pretraining dynamics, providing insights about the model's maturity. Such signals are especially useful to understand what happens with the model when the pretraining loss shows marginal improvements and fails reflect the changes inside the model.
Identifying the top bi-grams. To assess the model's progression, we extracted jet bi-grams from OLMo-7B model checkpoints across 555K pretraining steps. Table O presents a summary of the top 10 jet bi-gre bi-grams at different stages of training. Due to space reason, we only show the top 10 jet bi-grams every 100K steps. Initially, the network exhibits nonsensical jet bi-grams, such as "ICUirl ing".. As training advances, it gradually learns more meaningful combinations, like "at least". This process of acquiring sensible bi-grams stabilizes around step 200K, indicating that the model is reaching level of maturity where the top 10 bi-grams capture common meaning.
Learning speed. To evaluate the learning speed of jet bi-grams during pretraining, we consider the jet bi-grams the final training step (555K) as the ground-truth bi-grams. We then chart the hit ratios of these ground-truth bi-grams at each pretraining step, as illustrated in Figure 20a Interestingly, even though the pretraining loss (the blue curve) shows only minor improvements after the initial 50K steps, the model's ' acquisition of effective bi-grams continues to progress in steady, consistent manner. Hence bi-grams learning dynamics are active throughout the training procedure, even after the training loss stabilizes. This indicates that there is significant behavior change in the model which is not well captured by the training loss, an observation that is studied also in grokking and double-descent (Zhang et al. 2021 Power et al. 2022). In other words, jet bi-grams may offer another point of view for analyzing the learning dynamics compared to pretraining loss. In addition, fig- 20b characterizes the total pseudo-joint probability
of top 1K bi-grams from empirical data (Liu et al., 2024). We derive pscudo-joint jet bi-gram probability using statistical uni-grams from (Liu et al. 2024). We observe that the model gradually accumulates probability mass that aligns with the real corpus data distribution.
Learning schemes for different bi-grams. To understand there are any differences between the learning schemes of different bi-grams, we can trace the progression of the jet bi-gram scores for selected bi-grams. Figure 19 provides a visual comparison of how different bi-grams are promoted or suppressed during the pretraining process. The different slopes and levels of the lines indicate varying rates of learning for the respective bi-grams. We observe that, the model first acquires ran- dom bi-grams due to random parameter initialization. These random bi-grams, like "ICUirling" and "VENT thanks", are quickly suppressed in the early steps and never regain high scores. In contrast, one-to-many bi-grams like "at. least" are first promoted to very high scores but then get suppressed perhaps due to the model seeing more of the scope of the token "at". One-to-one bi-grams like "&amp" (HTML code) are gradually promoted and stabilize. Many-to-many bi-grams like "make sure" takes the most time to learn and the scores are still increasing even at the end of pretraining. Our findings suggest that the training process effectively promotes certain "good" bi- grams, but at different paces, where they might be suppressed later depending their occurrences and linguistic nature. These insights could inform future training strategies, such as targeted training on more relevant bi-grams adjusting the training data to improve the pretraining speed.
32

---

Published as a conference paper at ICLR 2026
Jet Bigram Score
ICUirlnng
thanks kamp
0.6
Score
at least
0.2
100000
200000
300000
400000
500000
Pretraining Steps
Figure 19: Visualization of OLMo-TB's promotion and suppression dynamics of jet bi-grams scores.
Protraining
Frotrahnigg
Pretraining Loss
0.8
Pretraining Loss
Hitz Ratie
0.01
5-
0.012
4-
-02
0.010
3-
0.000
100000
200000 Pretraining Steps
300000
400000
500000
100000
200000 Pretraining Steps
300000
400000
500000
(a) Top 1K jet bi-gram hit ratios w.r.t the final step.
(b) Top 1K jet bi-gram mass w.r.. empirical data.
Figure 20: Analysis of OLMo-7B's pretraining dynamics via measuring its jet bi-gram progression.
Table Bi-gram evolution across pretraining steps for OLMo 7B, Each column represents distinct step, while each row corresponds to different rank. The table entries are the bi-grams at each step for each rank. The number of tokens in association with the pretraining steps is also annotated. The model gradually picks up meaningful bi-grams after starting from senseless bi-grams (due random initialization).
[#steps]
100K 442B
200K
300K
400K
Rank
[tokenss]
885B
1327B
1769B
2455B
immortal
samp samp
6amp
1
ICUirling
least
at least
iLs
LLs own
2
ords architect
its own
kamp
its own
thei= own
their own
3 yaml Adam
okerly
your own
your own
at least
his awn
231 next
VENT Lhanks
Lheis own
youz own
sure your
clenal8%
iums
iums
Lhan
Charge
you'
can't 2nd
2nd
7 aveir careless
Everything
2nd
his owa
moze than
at least
HOLD worsening
erna already 'my
you guys
2nd
make sure
more than
diemant
auze
can't lums
33


Reviews:
[{"review_id": "cQOKdpnlYM", "reviewer": "Reviewer_F5cx", "summary": "The authors aim to solve the problem of \"entangled\" computation in LLMs, which makes their internal workings difficult to interpret. They use the core insight to treat LLMs as residual networks (a sum of paths) and apply jet operators, the functional equivalent of Taylor series, to analyze them. This approach is well-suited for the problem because jets are mathematically designed to handle compositions of functions and, via Lemma 1, can \"disentangle\" a nonlinear function of a sum into a weighted sum of its parts. This functional decomposition allows the authors to \"carve out\" and analyze specific computational paths. Their main contributions are the \"jet lens,\" which generalizes the popular Logit Lens (showing it's just the $k=0$ case), and \"jet n-grams,\" a novel \"corpus-free\" method to extract symbolic knowledge tables directly from the model's parameters. The authors provide strong empirical results, showing their decompositions are highly faithful (high cosine similarity) and that their $k>0$ jet lens provides a more stable interpretation for models like GPT-Neo where the standard Logit Lens has more chaotic results.", "questions": "The RLHF & toxicity result supports some results from \"Safety Alignment Should be Made More Than Just a Few Tokens Deep\" that alignment usually only masks the problems. It would be interesting to see if the toxic bi-gram mass changed significantly in the \"more robustly\" tuned models from \"Safety Alignment Should be Made More Than Just a Few Tokens Deep\".\n\nCost of $k>0$ Lenses: Could the authors quantify the practical computational overhead of using the $k=1$ iterative jet lens (Fig. 9) versus the $k=0$ lens (Fig. 8)? How much more expensive is it, and does this limit its real-world applicability as a drop-in replacement?\n\nChoice of $k$: The results for GPT-Neo are dramatically better for $k=1$, but the results for GPT-2-large (Figs. 11 vs. 12) seem comparable. Does this suggest $k=0$ is \"good enough\" for some model families? How should a practitioner choose the optimal $k$ for a new model?\n\nWould the authors release the code?\n\nNotation in Example 1: The expression $J^1\\gamma(x_1)(x_2)$ is unclear. Could the authors confirm this is meant to be $J^1\\gamma(x_1)(x_1+x_2)$, as I think it expands as $J^1 \\gamma(x_1)(x_1+x_2) = \\gamma(x_1) + \\gamma'(x_1)((x_1+x_2) - x_1)$?\n\nLemma 1 Proof: In the proof of Lemma 1 (Appendix A), the evaluation point for the jets appears to be mislabeled as 'x' when it should be 'y' (e.g., $\\sum w_i J^k f(x_i)(x)$). Can the authors confirm if this is a typo? Or clarify the difference between x and y?\n\n\n\"Data-Free\" Clarification: The \"data-free\" claim for n-grams seems to be a source of confusion. The method is clearly input-dependent (it evaluates on token embeddings). Would \"corpus-free\" or \"dataset-free\" be a more accurate descriptor?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "Principled Framework: The paper's primary strength is its formalism. It gives interpretability a more formal mathematical framework/tool grounded in approximation theory.\n\nGeneralization of Existing Tools: Proving that the Logit Lens is simply the $k=0$ case of the iterative jet lens is a great benfit. It places a popular tool on solid theoretical footing and simultaneously shows its limitations.\n\nEmpirical Support: The results strongly support the claims.\n\nFaithfulness: The high cosine similarity (e.g., 0.993 in Figure 4) between the expansion logits and the model logits shows the decomposition is not just theoretical but empirically sound.\n\nSuperiority: The $k>0$ iterative lens provides a much more stable and interpretable computational trace for models like GPT-Neo, where the standard $k=0$ Logit Lens is more chaotic (as shown clearly in Appendix I, Figs. 8 vs. 9).\n\nApplications & Insights:\n\nRLHF & Toxicity: The finding in Table 4 provides strong, quantitative evidence that RLHF alignment masks toxic associations rather than erasing them (ToxiGen score drops to 0, but toxic bi-gram mass is unchanged). The method of demonstrating it via a \"corpus-free\" n-gram analysis is novel and powerful.\n\n\"Corpus-Free\" n-grams: The idea of extracting symbolic n-gram tables without a dataset is a powerful one. The application of \"diffing\" models (Llama-2 vs. CodeLlama) by comparing their n-gram tables is a practical method for verifying fine-tuning.\n\nFuture Work: The framework opens up many exciting possibilities for \"functional-level\" interpretability.", "weaknesses": "Tractability and Clarity: The paper's discussion of the $O(2^L)$ exponential expansion (Algorithm 2) versus the practical, $O(L)$ linear-in-depth applications (the lenses) could be clearer. The tractability of these expansions, especially for $k>1$ is a major practical concern that is only briefly touched upon.\n\nNew Hyperparameters: The method introduces new hyperparameters for interpretability, namely the jet order $k$ and the jet weights $w$. The sensitivity of the results to these choices is not fully explored. Also how efficient and practical is it to get arbitrary order expansions?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "CSoJVfGppu", "reviewer": "Reviewer_tf63", "summary": "This paper introduces Jet Expansions, a framework that expands computational graphs using jet operators (generalizing truncated Taylor series) to decompose transformers into input-output computational paths and complementary remainders. The authors demonstrate how this functional decomposition stimulates advances in the area of interpretability applications, such as grounding and subsuming the traditional interpretability techniques such as the Logit Lens technique and defending against jailbreak attacks.", "questions": "Please see Weaknesses.", "limitations": "", "rating": 6, "confidence": 2, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "1. The paper has an impressive theoretical foundation. They introduced the concept of jet operators from differential geometry to define interpretability rigorously, which is both novel and effective.\n\n2. Using the novel framework, the paper elegantly unifies the framework and as a result, introduces a technique to subsume traditional foundations of interpretability.\n\n3. Unlike most interpretability methods, jet expansions operate directly on model structure without requiring probe datasets. This is a significant practical advantage.\n\n4. The topic has broad applications, including model safety and mechanistic analysis", "weaknesses": "My major concern is the remainder interpretation. The remainder $\\delta$ seems to be a critical quantity to ensure performance and credibility, yet it is not analyzed rigorously enough. For instance, in Remark 2 on page 6, they mentioned \"Empirically, remainders are often small and expansion logits nearly collinear with model outputs\". From a theoretical point of view, how small is \"small enough for the task\"?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "EnXcsZD40K", "reviewer": "Reviewer_KiGr", "summary": "- The authors introduce Jet Expansions, a framework for expanding computational graphs using jet operators that give interpretability to complex models (e.g., transformers).\n- A jet expansion uses a transformer's residual structure to expand it to many smaller layer combination terms, plus some remainder.\n- For interpretability, you look at the different orders of jets to understand what the model is trying to predict.", "questions": "How large are remainders in practice, and can we be more explicit about their sizes?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "- The paper's main contribution, turning LLMs into jet-expanded paths, requires only model access and no extra data to probe models or compare two models with a shared vocabulary.\n- The paper does a good job explaining the background and motivation, given the complexity of the work.", "weaknesses": "I question the practicality of jet expansions for interpretability. The paper only demonstrates the method on small to mid-sized open models, and even there, higher-order jets already add significant runtime. It's unclear how feasible this is for larger / frontier models.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "vGP6Pat261", "reviewer": "Reviewer_QKXz", "summary": "The paper proposes a new interpretability method for examining transformer-based language models based on jet expansions. The method works by approximating the full language model by its Taylor expansion and analyzing low-order terms. This extracts low-order behavior from the model and allows one to interpret it functionally, without relying on any additional data or training. The authors propose an implementation of the framework for a transformer-based language model and analyze its efficiency. \nJet expansions subsume existing interpretability techniques such as logit lens and $n$-gram statistics-based interpretability methods. Thus, the authors use their method to analyze open-source language models and find interpretable results. For example, by comparing fine-tuned and non-fine-tuned models, they find that the fine-tuning did not affect certain parts of model behaviour.", "questions": "- Is there any way to understand how good the approximation is, i.e., how large the remainders are?\n- How much of this methodology relies on the transformer architecture? Would it be easy to port to other neural LM architectures?\n- Could the same methodology even be used for things like model compression (where only the jets are retained/used)?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- The proposed methodology seems broadly applicable and useful\n\t- I particularly like that the approach is dataset-free, which removes many possible confounders in the analysis\n\t- Similarly, I like that the method is more holistic than individual neurons, etc.\n- The algorithm is thoroughly described and seems applicable and useful\n\t- It’s interesting that the analysis can even be done on CPUs\n- I liked the thoroughness of the experiments, which cover many open-source models.\n- Experiments clearly show that taking higher-order jets ($k = 1$ instead of $k = 0$ for logit lens) actually matters", "weaknesses": "- Unfortunately, I found it a bit hard to follow the transition between the empirical results and the theoretical setup. More explicitly stating what the measured quantities in the experiments correspond to in terms of the theory and the notation introduced above would be useful here, I think.\n\t- For example, how do $n$-gram statistic map to jets?\n- It is unclear to me to what degree the results should be trusted as a truthful representation of the model. For example, without knowing anything in advance, is it not possible that most of the model behavior is only explained by higher order terms than those in the jet?\n- The work mostly reframes existing techniques as jet expansions but doesn’t, to my understanding, propose novel applications of this framework for analyses that weren’t possible before.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)