One-sentence Summary

Researchers from Instituto de Telecomunicações propose LATIM, a token-level decomposition method for Mamba-1 and Mamba-2 that captures precise token-to-token interactions at the layer level, validated on machine translation, copying, and retrieval-based generation.

Key Contributions

• LATIM is introduced as a token-level decomposition method for Mamba-1 and Mamba-2 that extends feature-interaction analysis to state-space models for fine-grained interpretability of long-context sequence modeling.

• The method decomposes individual token contributions within Mamba layers to reveal token-to-token interaction patterns and selective information flow, and it can also be applied to linear recurrent architectures such as DeltaNet and mLSTM.

• Experiments on copying, machine translation, and retrieval-based generation show that LATIM, especially the ell_2 variant, clarifies Mamba's recall behavior, including a sparse and decaying focus on relevant tokens, while an exact counterpart removes approximation error without harming interpretability or task performance.

Introduction

State space models such as Mamba are increasingly used for long context sequence modeling because they can process long inputs more efficiently than Transformers, but their internal decisions are harder to inspect since they lack explicit attention maps. Earlier interpretability methods approximate implicit attention or track layerwise information flow, yet they do not fully capture precise token level contributions across layers. The authors introduce LATIM, a token level decomposition method for Mamba 1 and Mamba 2 that reformulates state space computation so attention style interpretability methods such as ALTI can be adapted to Mamba. The method is evaluated on copying, machine translation, and retrieval based generation, giving a more transparent way to analyze token interactions in state space models.

Method

The authors introduce LATIM, a decomposition framework for Mamba that rearranges the forward pass so that the contribution of an input token $\boldsymbol{x}_{j}$xj​ to an output token $\boldsymbol{y}_{i}$yi​ can be isolated. The method builds directly on the internal recurrence used by Mamba-style state-space models and converts the otherwise entangled token interactions into an explicit token-to-token attribution form.

At a high level, a Mamba block contains a sequence mixing layer, a gating mechanism, and an output projection. In Mamba-1, the input is first projected and passed through a causal one-dimensional convolution. The convolution output is then transformed by a SiLU activation, which produces the signal used by the recurrent state-space module. The SSM parameters are input-dependent, meaning that matrices such as $\boldsymbol{B}_{i}$Bi​ and vectors such as $\boldsymbol{c}_{i}$ci​ are learned projections of the current input. After the SSM produces its recurrent output, a second branch generates a gate through another SiLU activation, the two branches are multiplied element-wise, and a final linear projection maps the result back to the model dimension.

A crucial insight used by LATIM is that the Mamba recurrence can be unrolled into a lower-triangular interaction structure. By expanding the recurrence across positions, the SSM output can be interpreted as the result of multiplying the channel input by a matrix whose entries encode how earlier positions influence later positions. This reveals a hidden attention-like structure. However, unlike standard transformer attention, the contribution of each token is not directly exposed. The main obstacle is the SiLU activation. Because SiLU is nonlinear and non-additive, the pre-SSM signal cannot be split cleanly into separate token-wise contributions without introducing approximation.

For Mamba-1, the authors address this limitation by introducing an additive approximation function $f$f that approximates the SiLU activation applied to the convolution output. Starting from the causal convolution, the preactivation at position $i$i can be written as a sum over local convolution windows:

$$\boldsymbol{\psi}_{i} = \sum_{k=1}^{w} \boldsymbol{W}_{c}^{(k)} \left( \boldsymbol{W}_{x}^{\top} \boldsymbol{x}_{i-w+k} \right) + \boldsymbol{b}_{c}.$$ψi​=k=1∑w​Wc(k)​(Wx⊤​xi−w+k​)+bc​.

If an additive function $f$f approximates the SiLU activation, the SSM input can be decomposed as

$$\boldsymbol{\phi}_{j} = \sum_{k=1}^{w} f\!\left( \boldsymbol{W}_{c}^{(k)} \boldsymbol{W}_{x}^{\top} \boldsymbol{x}_{j-w+k} + \delta_{k,0} \boldsymbol{b}_{c} \right),$$ϕj​=k=1∑w​f(Wc(k)​Wx⊤​xj−w+k​+δk,0​bc​),

where each term captures the local convolutional contribution coming from a shifted input position. Substituting this decomposition into the unrolled SSM recurrence gives an explicit expansion of the recurrent output. After incorporating the skip term associated with the direct feed-forward component, the output at position $i$i can be expressed as

$$\boldsymbol{v}_{i} = \sum_{j=1}^{i} \sum_{k=1}^{w} \left( \boldsymbol{M}_{i,j+k} + \delta_{i,j+k} \boldsymbol{D} \right) \boldsymbol{\varphi}_{j}^{(k)}.$$vi​=j=1∑i​k=1∑w​(Mi,j+k​+δi,j+k​D)φj(k)​.

Here, $\boldsymbol{\varphi}_{j}^{(k)}$φj(k)​ denotes the additive component associated with the $k$k-th convolution kernel position for the $j$j-th input token. The matrix $\boldsymbol{M}$M encodes the token-to-token interaction through the recurrent state-space transition, while the Kronecker delta term accounts for the direct output path.

From this expression, the authors isolate the portion of the SSM output that originates from token $\boldsymbol{x}_{j}$xj​:

$$\boldsymbol{v}_{i \leftarrow j} = \sum_{k=1}^{w} \left( \boldsymbol{M}_{i,j+k} + \delta_{i,j+k} \boldsymbol{D} \right) \boldsymbol{\varphi}_{j}^{(k)}.$$vi←j​=k=1∑w​(Mi,j+k​+δi,j+k​D)φj(k)​.

They then carry this token-specific recurrent contribution through the gating mechanism and the final output projection. If $\boldsymbol{Z}_{i}$Zi​ denotes the gate at position $i$i, the contribution vector from $\boldsymbol{x}_{j}$xj​ to $\boldsymbol{y}_{i}$yi​ is

$$\boldsymbol{T}_{i}(\boldsymbol{x}_{j}) = \boldsymbol{W}_{o}^{\top} \left( \boldsymbol{Z}_{i} \odot \boldsymbol{v}_{i \leftarrow j} \right).$$Ti​(xj​)=Wo⊤​(Zi​⊙vi←j​).

The block output can then be written as the sum of these token-specific contributions:

$$\boldsymbol{y}_{i} = \sum_{j=1}^{i} \boldsymbol{T}_{i}(\boldsymbol{x}_{j}).$$yi​=j=1∑i​Ti​(xj​).

This structure provides the same kind of interpretability interface that attention decomposition gives for transformer blocks: each term $\boldsymbol{T}_{i}(\boldsymbol{x}_{j})$Ti​(xj​) represents the contribution of an earlier token to the current output position.

For Mamba-2, the block design is similar but contains two important differences. First, the recurrence uses a simpler form of the state transition matrix by defining $\boldsymbol{A}_{i}$Ai​ as a scalar multiple of the identity. Second, Mamba-2 introduces a GroupNorm layer after the gating operation, producing

$$\boldsymbol{u}_{i} = \text{GroupNorm} \left( \boldsymbol{Z}_{i} \odot \boldsymbol{v}_{i} \right).$$ui​=GroupNorm(Zi​⊙vi​).

To preserve token-wise decomposition under GroupNorm, the authors observe that, at inference time, GroupNorm can be viewed as an affine map around a fixed value of $\boldsymbol{u}_{i}$ui​. Specifically,

$$\text{GroupNorm}(\boldsymbol{u}_{i}) = \gamma_{i}(\boldsymbol{u}_{i}) \boldsymbol{u}_{i} + \boldsymbol{\beta}_{i},$$GroupNorm(ui​)=γi​(ui​)ui​+βi​,

where $\gamma_{i}(\boldsymbol{u}_{i})$γi​(ui​) is treated as a fixed linear operator once the full value of $\boldsymbol{u}_{i}$ui​ is known. Since the token-specific portion $\boldsymbol{u}_{i \leftarrow j}$ui←j​ passes through the same linear operator, the contribution vector for Mamba-2 becomes

$$\boldsymbol{T}_{i}(\boldsymbol{x}_{j}) = \boldsymbol{W}_{o}^{\top} \left[ \gamma_{i} \left( \boldsymbol{u}_{i} \right) \boldsymbol{u}_{i \leftarrow j} \right].$$Ti​(xj​)=Wo⊤​[γi​(ui​)ui←j​].

The resulting contribution vectors can be reduced to scalar interpretability scores using different aggregation functions, such as norm-based reductions or logit-aware reductions.

The method also includes two decomposition modes. The approximate mode does not require retraining and replaces the non-additive SiLU activation with an additive surrogate $f$f. Taylor expansions are explored as candidates for $f$f, but the authors find that directly setting $f$f to be the same SiLU function yields the lowest approximation error. This approximate mode introduces some mismatch because the original block is not exactly decomposable into additive token contributions.

The exact mode removes this mismatch by modifying the Mamba block itself. The SiLU activation after the convolution is removed, so the SSM input becomes simply the convolution output. In this case, the additive function is effectively the identity, and the decomposition reproduces the block output without approximation error. This exact version requires training the modified model, but it preserves the same type of token contribution structure while giving zero decomposition error.

Experiment

Experiments evaluate LATIM against Mamba interpretability baselines on copying, machine translation, retrieval-based generation, and approximation-error analysis, combining qualitative inspections with faithfulness, alignment, recall, and error assessments. In copying and translation, LATIM variants, especially the l2 variant, recover finer and more reliable token interactions than coarse global attributions, while some aggregated variants retain sparsity but miss token-level alignment. Retrieval experiments show that Mamba can struggle with selective recall in multi-key and frequency-tracking tasks, with LATIM revealing a bias toward the first key and declining attention over repeated tokens. Finally, approximation analysis indicates that replacing nonlinear activations yields an exact and effective method, yet the practical SiLU approximation retains comparable interpretability and translation quality.

The methods compute contribution scores for token interactions by decomposing Mamba computations into token-originated components and scoring those components in different ways. The variants compare direct contribution magnitude, output norm change after removing a contribution, and projection toward output logits. These scores are used to evaluate Mamba interpretability across recall, machine translation, and retrieval settings. Norm-based LATIM scores reflect the magnitude of an individual token contribution. ALTI-based scores estimate importance through the reduction in output norm caused by removing that contribution. Logit-based scoring links token contributions to model output decisions. The variants are tested against Mamba-specific attribution and attention alternatives to assess how well they reveal interaction structure.

In the copying task, LATIM variants are the most faithful interpretability methods for the Mamba models, better recovering the expected source to copied token relationship. They outperform baseline attention, attribution, and LRP approaches across faithfulness metrics. LATIM with ALTI leads on precision and recall oriented measures, while LATIM with l2 is strongest for overall curve area. LATIM variants exceed baseline attention, attribution, and LRP methods on faithfulness metrics. LATIM with ALTI gives the strongest average precision and recall at K in the shown results. LATIM with l2 achieves the highest AUC among the compared methods.

For machine translation token alignment, LATIM l2 produces the lowest alignment error rates, indicating the strongest token-to-token interaction capture. Methods that aggregate relevance across layers are less effective overall, though MambaLRP scores better than LATIM ALTI-Logit. The results also suggest that per-layer alignment estimates can be preferable to collapsed global representations. LATIM l2 gives the best translation alignment accuracy among the compared methods. MambaLRP consistently outperforms LATIM ALTI-Logit when relevances are aggregated across layers. Smaller and larger Mamba variants show broadly similar qualitative differences between interpretability methods.

Mamba-2 shows a strong bias toward the first passkey in multi-key retrieval, with accuracy consistently higher when the target is first rather than second or later. This first-key preference becomes especially severe in the four-key setting, where later keys are much harder to recover. Larger model sizes do not uniformly fix the issue, as 780M can outperform 1.4B in several retrieval conditions. Across model sizes, retrieval accuracy is much higher for the first key than for later keys. The four-passkey setting causes a large drop in accuracy for targets that are not first. Increasing model size does not guarantee better recall for later passkeys. The 780M model performs notably well in several conditions, exceeding the 1.4B model for some second and later key cases.

The comparison shows that the activation choice affects both approximation error and downstream translation behavior. The identity activation yields exact decomposition and the most favorable interpretability metric, while SiLU and continued pretrained SiLU remain close. ReLU produces substantially larger approximation errors, especially in deeper layers, and lower translation quality. Identity activation achieves zero approximation error and the best AER Approximation error grows with layer depth for non-linear activations Continued pretraining slightly improves SiLU error and COMET ReLU performs worst on both approximation error and COMET

The experiments evaluate Mamba interpretability by applying LATIM scoring variants to token interaction tasks across copying, machine translation alignment, multi-passkey retrieval, and activation sensitivity, comparing them with attention, attribution, and LRP baselines. Overall, LATIM methods recover token relationships more faithfully, with ALTI-based scores performing strongest in copying and l2-based scores best for translation alignment, while layer-wise estimates are generally preferable to collapsed cross-layer aggregations. The retrieval and activation analyses further reveal limitations of Mamba models, including a pronounced bias toward early passkeys that scaling does not fully remove and a clear advantage of identity activation for faithful decomposition compared with nonlinear choices such as ReLU.