PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Ido Andrew Atad, Itamar Zimerman, Shahar Katz, Lior Wolf

arXiv:2601.17958v1cs.LG

TL;DR

Existing attention analyses and aggregations do not provide a unified representation of global Transformer computation across blocks and components. TensorLens reformulates the entire model as an input-dependent high-order tensor operator, and experiments show richer, more behavior-reflective representations that support interpretability tools.

  • Problem

    Existing attention analyses and aggregation methods do not fully represent global Transformer behavior across all blocks and components.

  • Method

    TensorLens represents the entire Transformer as a single data-controlled linear operator composed from high-order block tensors that incorporate attention and other model components.

  • Results

    TensorLens better reflects model behavior than prior attention forms, with tensor-based perturbation aggregation reaching AUC above 0.82 when using input and output embeddings on DeiT base.

  • Takeaways & Limitations

    The tensor representation can serve as a drop-in foundation for interpretability, analysis, intervention, and mechanistic model-understanding techniques.

  • Takeaways & Limitations

    Experiments are limited to models up to 1B parameters and moderate input lengths, while the high-order representation is GPU-memory intensive.

Abstract

from arXiv · show

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.

1 Introduction

Transformers are analyzed through attention matrices, but existing formulations do not fully represent global behavior across blocks and components. TensorLens addresses this gap with a theoretically grounded, input-dependent high-order tensor representation of the entire Transformer.

  • Evidence and implications: Figure 1 characterizes the Transformer as an input-dependent high-order attention tensor that captures the architecture’s principal sub-components.The depicted formulation includes attention-related computation together with FFNs, normalization, embeddings, and residual connections.
  • Motivation: Attention matrices support interpretability and model analysis by representing token interactions as input-dependent weighted linear combinations.Their row-wise softmax normalization makes layer updates comparatively interpretable.
  • Motivation: Existing aggregation methods average attention across heads and multiply across layers, with later variants incorporating output projections, FFNs, and other components.These approaches extend attention beyond individual matrices but do not provide the unified formulation proposed here.
  • TensorLens: TensorLens represents the entire Transformer as a single data-controlled linear operator using a high-order attention tensor.The representation includes FFNs, activations, LayerNorm, skip connections, embeddings, and other components.
  • TensorLens: TensorLens introduces generalized attention maps that can replace standard attention matrices and cross-layer aggregations while encompassing all model parameters.The formulation is presented as theoretically principled and more precise than prior attention variants.
  • Evidence and implications: Empirical analyses show that TensorLens better reflects model behavior through perturbation evaluations and supports approximating linear relations from language-model embeddings.The tensor structure is positioned as a foundation for mechanistic interpretability and model-understanding tools.

2 Background & Related Work

Prior work has refined attention analysis by incorporating additional Transformer components and has explored higher-order tensors for expressivity. TensorLens instead reinterprets vanilla self-attention through a comprehensive tensor formulation, while attribution methods provide complementary tools for explaining model outputs.

  • Extended attention formulations: Prior attention formulations progressively incorporated output projections, residual connections, normalization layers, and FFN sub-layers to refine matrix-based analysis.
  • Higher-order tensor methods: Higher-order tensor approaches have primarily modified Transformer architectures to increase expressivity, often with reduced efficiency that motivates computational improvements.
  • Higher-order tensor methods: TensorLens differs by reinterpreting the vanilla self-attention mechanism through a tensor-based formulation rather than changing the Transformer architecture.
  • Interpretability methods: Attribution methods quantify contributions from neurons or input features and are used to make neural-network decisions more understandable.
  • Interpretability methods: Class-specific attribution includes gradient-based methods, whereas class-agnostic approaches include activation maximization, probing, and attention-map extraction.

3 Method: TensorLens

TensorLens tensorizes each Transformer component and composes the resulting block tensors into a unified, input-dependent linear representation of the full model. The fourth-order tensor captures token interactions and cross-channel transformations while retaining the model’s parameters and computation structure.

  • Tensorization: Each Transformer sub-layer is represented as a data-controlled linear operator whose coefficients depend on the input.The formulation covers self-attention, LayerNorm, FFNs, activations, residual connections, and embeddings.
  • Tensorization: The tensor representation is required because attention mixes tokens while FFNs mix hidden dimensions.A single matrix cannot simultaneously capture both interaction types in the combined computation.
  • Block construction: Stacking tensorized self-attention, residual, normalization, and FFN sub-layers produces a tensor T (n) for each Transformer block.The construction supports post-layernorm blocks and provides a corresponding pre-layernorm derivation in Appendix B.
  • Entire Transformer: The full Transformer is obtained by composing the block tensors across layers as a chain of high-order tensor transformations.Starting from X0 = X, the model applies the block operators sequentially through T (N) ◦ T (N−1) ◦ · · · ◦ T (1).
  • Generalized attention: The resulting fourth-order tensor has dimensions L × D × L × D and maps every input token–channel pair to every output token–channel pair.It generalizes attention matrices by modeling both intertoken dependencies and intra-token cross-channel interactions.
  • Analysis and interpretation: TensorLens can be restricted to selected layers or heads to produce generalized attention matrices at different analysis granularities.The representation is also presented as encompassing all model parameters and as a direct local approximation of the full Transformer computation.

4 Experiments

Experiments evaluate TensorLens as a proxy for whole-Transformer behavior through perturbation tests and linear relation decoding. Tensor-based aggregation generally outperforms non-tensor baselines, though decoder-only results are less conclusive.

  • Experimental design: TensorLens is evaluated against attention aggregation baselines using perturbation tests and linear relation decoding.The experiments assess representation power as a proxy for Transformer behavior and its usefulness for mechanistic interpretability.
  • Perturbation tests: Vision perturbation tests show tensor methods consistently outperform baselines across perturbation levels.For DeiT-Base, Tensor,In+Out exceeds 0.82 total AUC, Tensor,Norm reaches 0.66, while non-tensor methods remain below 0.6.
  • Perturbation tests: NLP perturbation tests on BERT likewise show tensor aggregation outperforming all baselines across perturbation levels.Tensor,In+Out exceeds 0.158 total AUC, Tensor,Norm reaches 0.101, and non-tensor methods fall below 0.09.
  • Perturbation tests: Decoder-only evaluations are less conclusive despite TensorLens achieving top or second-best AUC scores across benchmarks.The tested models include Pythia-1B, Pico-570M, and Phi-1.5 on WikiText-103.
  • Relation decoding: Relation decoding forms a per-relation tensor from m examples and compares its outputs with the original model on held-out subject-object pairs.Approximation accuracy is the percentage of examples where the top-predicted object matches the original output.
  • Relation decoding: Tensor-based relation decoding achieves higher accuracy than the LRE baseline on most examined relations, while some tasks remain near random-guess performance.The occupation-age relation is an example where both methods perform close to random, consistent with limits of linear input transformations.

5 Conclusions

TensorLens aggregates attention across Transformer blocks and their sub-components using a high-order, data-controlled linear operator. The formulation captures internal interactions including FFNs, embeddings, and LayerNorm, providing a theoretically grounded basis for interpretation.

  • 5 Conclusions: TensorLens aggregates attention across Transformer blocks and all sub-components within each block as a high-order data-controlled linear operator.The formulation includes contributions from FFNs, embeddings, LayerNorm, and other components.

Limitations

TensorLens has approximation, resource, scope, and maturity limitations. Its experiments cover models up to 1B parameters and moderate input lengths, while broader tensor-based applications remain underexplored.

  • Method limitations: Some linearization choices prioritize simplicity over deriving optimally approximated tensors.The paper specifically identifies the activation decomposition in Eq. 13 as an example.
  • Computational limitations: The high-order tensor representation is GPU-memory intensive despite a memory-optimized computation method.The mitigation is described in Appendix C.
  • Scope and future work: Experiments are limited to models up to 1B parameters and moderate input lengths.The authors also state that the full potential of the tensor-based approach remains underexplored.

Ethics Statement

The work uses publicly available datasets and does not involve human subjects, personally identifiable data, or harmful-content generation. The authors acknowledge that interpretability tools may support transparency or expose model vulnerabilities, and encourage responsible use.

  • Data and study scope: Evaluations use publicly available datasets including ImageNet, IMDB, and WikiText-103 under their respective licenses.The study does not involve human subjects, personally identifiable data, or the generation of potentially harmful content.
  • Responsible use: The authors acknowledge that interpretability tools can enhance trust while also exposing or enabling exploitation of model vulnerabilities.They state that positive implications outweigh potential misuse and encourage alignment with ethical AI principles.

A Additional Perturbation Experiments

Additional perturbation experiments evaluate TensorLens on RoBERTa, DeiT-Small, and modern transformer models, finding higher AUC scores than baselines across perturbation fractions. The tests use final hidden representations, with AUC based on hidden-state mean squared error for the reported NLP settings.

  • Additional perturbation experiments: Across all benchmarks, TensorLens achieves higher AUC scores than baselines for every perturbation fraction.The additional experiments cover RoBERTa, DeiT-Small, ModernBert, and Gemma3.
  • NLP perturbation tests: Figure 6 measures RoBERTa-Base sensitivity using mean squared error between original and perturbed [CLS] final hidden states.Higher MSE indicates a stronger perturbation effect under this evaluation.
  • Vision perturbation tests: Figure 7 evaluates how perturbations affect final hidden representations in DeiT-Small.
  • Next-token prediction: Table 1 reports next-token perturbation AUC using final-token hidden-state MSE and probability-change AOPC.The table notes that GlbEnc is unavailable for Pythia and Phi because their method does not apply to parallel-residual architectures.
  • Modern models: Table 2 reports NLP perturbation AUC for ModernBert-Base and Gemma3-270M trained for sentiment prediction on IMDB.Its metric is HS-MSE between the original and perturbed final-token hidden states.

B Tensor Derivation with Biases

The appendix extends TensorLens to include transformer biases, expressing self-attention, FFNs, LayerNorm, residuals, blocks, and the full model as input-dependent tensor transformations. A memory-efficient Jacobian procedure computes only the output slice needed for the experiments or reconstructs the full tensor through forward-mode passes.

  • Bias-aware formulation: With biases included, each transformer module is represented as an affine transformation whose bias is propagated through subsequent components.The model biases are broadcast across sequence positions before module-wise affine transformations are composed.
  • Tensorized self-attention: Self-attention is vectorized by grouping heads into an attention tensor and combining value-output projections with attention-dependent biases.The resulting flattened operator acts on vec[X] and adds vec[B_attn].
  • Tensorized FFN: The FFN activation is converted into an input-dependent Hadamard product, yielding a tensorized operator characterized by tensor M.The FFN is formed from two linear layers surrounding the element-wise activation.
  • Blocks and entire Transformer: Stacking self-attention, residual, normalization, and FFN tensors produces a tensor for each transformer block and recursively composes the entire model.Adding biases turns the full model representation into an affine transformation.
  • Memory-efficient computation: For efficiency, the experiments compute only the 3-dimensional tensor slice for the single output position being analyzed.The full 4-dimensional tensor can instead be computed with forward-mode differentiation using L · D forward passes, trading memory for computation.

D Relation Decoding Experiment

The relation decoding experiment adapts an existing relations dataset to TensorLens and evaluates the smaller Pythia-1B model under constrained sampling and prediction settings. It compares mean tensor approximations with LRE using input embeddings and averages results across six random train-test seeds.

  • Experimental setup: The experiment uses bias, common-sense, and factual relation categories from the setup of Hernandez et al. (2024).
  • Experimental setup: Samples are filtered to the most common token length, and test examples require the correct object to appear among the model’s top-20 predictions.These changes adapt the dataset to mean tensor estimation and the smaller language model.
  • Experimental setup: The evaluation uses Pythia-1B, a smaller model than the one used by Hernandez et al. (2024), because of limited academic computational resources.
  • Evaluation protocol: For each relation type, three training examples estimate the mean tensor and LRE weights, with results averaged over six random train-test seeds.
  • Evaluation protocol: The mean tensor uses the full affine transformation with biases to obtain an accurate approximation of the original model.The LRE comparison uses input embeddings and its default Pythia GPT-NeoX scaling value β = 2.5.

E Tensor Approximation Error Bound

The approximation analysis bounds the error of TensorLens around an input by combining the tensor operator norm with the transformer’s local change under perturbation. The resulting whole-transformer bound is data-dependent through LayerNorm variance and component-wise spectral-norm bounds.

  • Proposition 1: The tensor approximation error at X + ϵ is bounded by the tensor norm times ∥ϵ∥₂ plus the model’s change from F(X) to F(X + ϵ).This is the central statement of Proposition 1.
  • Whole-transformer bound: The tensor spectral norm is bounded by analyzing each transformer sub-module and combining bounds within blocks and across layers.The analysis covers self-attention, FFNs, LayerNorm, residual structure, and the whole transformer.
  • LayerNorm dependence: LayerNorm contributes a data-dependent term through the minimum variance of hidden states entering the normalization layer.The passage identifies this variance as the only data-dependent quantity in the LayerNorm bound.
  • Whole-transformer bound: The bound combines component-level estimates into a bound for the tensor of the complete post-layernorm transformer.
  • Bound interpretation: The approximation analysis notes that the remaining data-dependent quantity in the bound is typically a small constant.
Loading 2601.17958v1…