TensorLens: End-to-End Transformer Analysis via High-Order Attention Tensors

TL;DR

This work addresses the limitation of existing transformer analyses that focus on individual heads or layers by introducing TensorLens, which represents the entire Transformer as an input-dependent high-order attention tensor $\mathcal{T}$. It tensorizes each sub-layer (self-attention, LayerNorm, FFN, residual) and composes them into a full model tensor, enabling a generalized attention view that encompasses all parameters and components. The authors provide theoretical bounds on the forward approximation error and demonstrate empirical advantages over prior aggregation methods through perturbation-based evaluations and relation decoding. The framework offers a principled foundation for mechanistic interpretability and presents a practical, memory-conscious path to richer, end-to-end Transformer analysis beyond traditional attention matrices.

Abstract

Attention matrices are fundamental to transformer research, supporting a broad range of applications including interpretability, visualization, manipulation, and distillation. Yet, most existing analyses focus on individual attention heads or layers, failing to account for the model's global behavior. While prior efforts have extended attention formulations across multiple heads via averaging and matrix multiplications or incorporated components such as normalization and FFNs, a unified and complete representation that encapsulates all transformer blocks is still lacking. We address this gap by introducing TensorLens, a novel formulation that captures the entire transformer as a single, input-dependent linear operator expressed through a high-order attention-interaction tensor. This tensor jointly encodes attention, FFNs, activations, normalizations, and residual connections, offering a theoretically coherent and expressive linear representation of the model's computation. TensorLens is theoretically grounded and our empirical validation shows that it yields richer representations than previous attention-aggregation methods. Our experiments demonstrate that the attention tensor can serve as a powerful foundation for developing tools aimed at interpretability and model understanding. Our code is attached as a supplementary.

Figures (7)

• Figure 1: Transformers are re-formulated as data-controlled linear operators, characterized by an input-dependent high-order attention tensor $\mathcal{T}$. This formulation enables a unified self-attention representation that captures the entire Transformer architecture, including sub-components such as FFN layers, normalization, embedding layers, and residual connections.
• Figure 2: Method: A schematic visualization of our method, where each sub-component of the transformer architecture, including self-attention, LayerNorm, FFNs, input and output embedding layers, and the residual connection (which is omitted here for simplicity), is formulated as a data-control linear operator represented by high-order tensor in $\mathbb{R}^{L \times D \times L \times D}$. These tensors are composed into per-block tensors $\mathcal{T}^{(n)}$ for each layer $n \in [N]$, which are then used to construct the final linear operator representing the entire Transformer.
• Figure 3: Perturbation Tests in Vision: Effect of perturbations on final hidden representations of DeiT-Base. Measured by the mean squared error between the last hidden-state of the [CLS] token in the original and perturbed input (higher is better).
• Figure 4: Perturbation Tests in NLP: Effect of token perturbations on final hidden representations of BERT-Base.
• Figure 5: Relation Decoding: Accuracy relative to original model computation, for different relation categories on Pythia-1B, with $m=3$ training samples per relation. Results are averaged across 6 train-test splits, with standard deviation shown in error bars. Random baselines shown as horizontal dashed lines.

Theorems & Definitions (1)

• Proposition 1