TensorLens：通过高阶注意力张量实现端到端 Transformer 分析 TensorLens: End-to-End Transformer Analysis via High-Order Attention Tensors
将整个 Transformer 重写为高阶注意力张量，统一编码所有子组件用于模型理解
前置知识
自注意力机制（Self-Attention）
Transformer 的核心计算模块。给定输入序列 X，通过三个投影矩阵 $W_q, W_k, W_v$ 分别生成查询 Q、键 K 和值 V，然后计算注意力矩阵 $A = \text{softmax}(QK^\top)$，最终输出为 $AVW_o$。多头注意力将这一过程在 H 个子空间中并行执行，再拼接输出。注意力矩阵的每一行本质上是一个概率分布，表示某个 token 对所有其他 token 的关注程度，这使得它具有一定的可解释性。
TensorLens 的核心思路是将标准注意力矩阵从 $L \times L$ 的矩阵扩展为 $L \times D \times L \times D$ 的四阶张量，理解原始注意力机制的数学形式是理解本文张量构造的基础。
Attention Rollout 与注意力聚合方法
由于 Transformer 有多层多头，单个注意力矩阵只反映局部信息。Attention Rollout（Abnar & Zuidema, 2020）提出将同一层内的注意力矩阵跨头取平均，然后跨层做矩阵乘法来追踪信息流。后续工作（Kobayashi et al.）进一步将输出投影矩阵、残差连接、LayerNorm 甚至 FFN 纳入注意力分析，得到更精确的「扩展注意力矩阵」。这些方法的共同目标是用一个单一的矩阵来概括整个模型中 token 之间的信息流动。
TensorLens 是这类注意力聚合工作的最新进展，它通过引入高阶张量而非矩阵，首次将模型的所有参数和子组件统一编码。理解之前的聚合方法才能看出本文的突破在哪里。
Kronecker 积与张量收缩
Kronecker 积 $\otimes$ 是一种将两个矩阵扩展为更大矩阵的运算。在本文中，它被用来将矩阵运算向量化：例如双线性映射 $AXB$ 可以向量化为 $\text{vec}(AXB) = (B^\top \otimes A)\text{vec}(X)$。张量收缩（tensor contraction）则是对高阶张量的特定维度求和，类似于矩阵乘法的高阶推广。这些数学工具是将 Transformer 各子组件「张量化」的关键。
论文的整个技术推导依赖于将矩阵运算转化为向量空间中的线性映射，Kronecker 积是实现这一转化的核心工具。不理解它就无法跟上论文的张量构造过程。
数据控制线性算子（Data-Controlled Linear Operator）
一个线性算子 $T$ 作用于输入 $X$ 上产生输出 $T(X)$，如果 $T$ 本身依赖于输入 $X$（即 $T = T(X)$），则称为数据控制线性算子。在 Transformer 中，注意力矩阵 $A(X)$ 依赖于输入，FFN 中的激活函数 $\phi(X)$ 也依赖于输入。这意味着 Transformer 的每一层在给定输入的条件下都是一个线性变换，但变换本身随输入变化。这一视角是本文将整个 Transformer 重写为张量的理论基础。
论文的核心主张是 Transformer 可以被理解为一个「数据控制的线性算子」，理解这一概念是理解论文理论框架的前提。
研究动机
现有注意力分析方法存在一个根本性的不完备问题。标准的注意力矩阵 $A \in \mathbb{R}^{L \times L}$ 只编码了 token 之间的交互，完全忽略了 FFN、LayerNorm、残差连接和嵌入层的影响。即便是最先进的扩展注意力方法（如 Kobayashi et al. 2023 将 FFN 纳入），也只做到了部分组合，无法同时捕获 token 间交互和通道间交互。更关键的是，这些方法都停留在矩阵层面——一个 $L \times L$ 的矩阵——无法同时建模两种不同类型的信息混合：注意力在 token 维度上混合（跨 token），而 FFN 在特征维度上混合（跨通道）。强行将二者压缩到同一个矩阵中必然丢失信息。Elhage et al. (2021) 曾用四阶张量分析两层纯注意力 Transformer，但未推广到包含 FFN、LayerNorm 等完整架构。
本文的目标是本文的目标是找到一种最完备的数学表示，能够将 Transformer 的所有参数和所有计算步骤（包括注意力、FFN、LayerNorm、激活函数、残差连接、输入/输出嵌入）统一编码为一个单一的、输入依赖的线性算子。具体来说，作者希望这个表示能够：（1）理论上严格，可以提供近似误差的上界；（2）实践中可用，能作为注意力矩阵的即插即用替代品用于可解释性和模型分析；（3）比现有的任何注意力聚合方法都更忠实地反映模型行为。
与已有工作不同的是，本文的独特切入点在于一个关键洞察：由于注意力在 token 维度混合信息，而 FFN 在通道维度混合信息，二者的组合不能用单个矩阵表示，必须使用高阶张量。这一洞察之前被忽视了，因为研究者一直试图将所有子组件压缩到 $L \times L$ 的注意力矩阵中，而没有意识到维度混合的异质性要求更丰富的数学结构。TensorLens 首次明确指出这一点，并给出了系统的张量化推导：将自注意力、LayerNorm、FFN、残差连接逐一转化为四阶张量 $T \in \mathbb{R}^{L \times D \times L \times D}$，然后通过张量乘法组合为整个模型的统一表示。这是第一个显式捕获所有模型参数的张量表示方法。
核心方法
TensorLens 的核心思路可以用一个类比来理解：如果标准注意力矩阵是一张「谁关注谁」的地图，那么 TensorLens 就是一张「谁的哪个特征如何影响了谁的哪个特征」的全景地图。技术路线分为四步：首先，将 Transformer 的每个子组件（自注意力、LayerNorm、FFN、残差连接）分别表示为数据控制的线性算子，并通过向量化和 Kronecker 积将它们转化为四阶张量；其次，将同一层内的子组件张量相乘，得到每层的块张量 $T^{(n)}$；然后，将所有层的块张量按顺序相乘，得到整个 Transformer 的统一张量 $T$；最后，提供三种从四阶张量到二维矩阵的坍缩方法（范数投影、输入-输出投影、类别特定投影），使其可以像标准注意力矩阵一样可视化和使用。
TensorLens 最本质的创新在于揭示了一个被忽视的数学事实：Transformer 中不同子组件混合信息的维度不同——注意力在序列维度 $L$ 上混合 token，FFN 在隐藏维度 $D$ 上混合通道——因此它们的组合必须用四阶张量 $T \in \mathbb{R}^{L \times D \times L \times D}$ 来表示，而不能压缩为 $L \times L$ 的矩阵。这意味着张量的每个元素 $T[i, d_1, j, d_2]$ 捕获了输入 token $j$ 的通道 $d_2$ 对输出 token $i$ 的通道 $d_1$ 的影响。这是之前所有注意力聚合方法都无法表达的粒度。具体而言，自注意力被向量化为 $A = \sum_h (W_{v,h} W_{o,h})^\top \otimes A_h$，FFN 被表示为 $M = (M_2^\top \otimes I_L) \Psi (M_1^\top \otimes I_L)$（其中 $\Psi$ 是激活函数的对角矩阵），LayerNorm 和残差连接也有对应的张量形式。这些张量相乘后得到每层块张量 $T^{(n)} = L^{(n)}_2 (M^{(n)} + I) L^{(n)}_1 (A^{(n)} + I)$，所有层相乘得到完整的 $T$。
方法步骤详情
方法的具体步骤如下。第一步是自注意力张量化：给定输入 $X$，多头注意力的向量化形式为 $\text{vec}[\text{Attn}(X)] = \sum_h (W_{v,h} W_{o,h})^\top \otimes A_h \cdot \text{vec}[X]$，其中 $A_h = \text{softmax}(XW_{q,h}(XW_{k,h})^\top)$，得到张量 $A \in \mathbb{R}^{LD \times LD}$。第二步是 LayerNorm 张量化：LayerNorm 操作 $\text{LN}(X) = \gamma \odot \frac{X - \mu}{\sigma} + \beta$ 被转化为 $\text{vec}[\text{LN}(X)] = (I_D - \frac{1}{D}\mathbf{1}\mathbf{1}^\top) \cdot \text{diag}(\frac{\gamma}{\sigma})^\top \otimes \text{diag}(\frac{1}{\sigma}) \cdot \text{vec}[X]$。第三步是 FFN 张量化：$\text{FFN}(X) = \phi(XM_1)M_2$ 中的逐元素激活函数 $\phi$ 被转化为输入依赖的 Hadamard 积，最终得到 $\text{vec}[\text{FFN}(X)] = (M_2^\top \otimes I_L) \Psi (M_1^\top \otimes I_L) \text{vec}[X]$。第四步是残差连接张量化：$\text{vec}[X + g(X)] = (I + G) \text{vec}[X]$。第五步是块组合：对 post-layernorm 块，$T^{(n)} = L^{(n)}_2 (M^{(n)} + I) L^{(n)}_1 (A^{(n)} + I)$。第六步是全模型组合：$\text{vec}[F(X)] = (\prod_{n=1}^N T^{(n)}) \text{vec}[X]$。最后提供三种坍缩方法将四阶张量压缩为可解释的矩阵。
技术新颖性
TensorLens 的技术新颖性体现在三个层面。首先，在表示层面，它是第一个将整个 Transformer（包括所有参数和所有子组件）统一编码为单一张量的方法。之前的最佳工作（Kobayashi et al. 2023）虽然将 FFN 纳入了注意力分析，但仍停留在 $L \times L$ 矩阵层面，无法同时建模 token 间和通道间的交互。其次，在理论层面，论文给出了张量近似误差的上界 $\|T_X(X+\epsilon) - F(X+\epsilon)\|_2 \leq \|T_X\|_2 \|\epsilon\|_2 + \|F(X+\epsilon) - F(X)\|_2$，这是之前任何注意力聚合方法都无法提供的理论保证。第三，在应用层面，张量结构可以直接用于线性关系解码（Linear Relation Decoding），通过在少量样本上取平均张量来近似模型的行为函数，在大多数关系类别上优于专门设计的 LRE 方法。
实验结果
论文在两个核心实验上验证了 TensorLens 的表示能力。第一个实验是扰动测试（Perturbation Test）：按照注意力分数从高到低逐步遮蔽输入 token，观察模型最后一层隐藏状态的均方误差（MSE）变化。在视觉领域（DeiT-Base，ImageNet-1K），TensorLens 的 AUC 达到 0.82（使用输入+输出嵌入投影），远超所有基线方法（最佳基线 W-Attn+Rollout 仅 0.59），差距高达 0.23。使用范数坍缩的 TensorNorm 也达到 0.66。在 NLP 领域（BERT-Base，IMDB），TensorLens 的 AUC 为 0.158，而所有非张量基线均低于 0.09，RoBERTa-Base 上同样表现一致（Tensor AUC 0.122 vs 基线 0.077）。对于自回归模型（Pythia-1B、Pico-570M、Phi-1.5，WikiText-103），TensorLens 在大多数指标上取得最佳或次佳成绩，但优势不如 encoder-only 模型明显，论文推测这是因为自回归模型本质上具有局部性，降低了扰动评估的信息量。第二个实验是线性关系解码（Relation Decoding）：在 Pythia-1B 上，仅用 3 个训练样本计算平均张量，TensorLens 在大多数关系类别（如 sentiment、superclass、country、language 等）上超越了专门设计的 LRE 基线方法，验证了张量表示能更好地捕获模型的内部功能。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| 视觉扰动测试（DeiT-Base，ImageNet-1K） | AUC（MSE vs 遮蔽比例曲线面积） | 0.820（In+Out），0.660（Norm） | 0.594（W-Attn+Rollout），0.554（GlbEnc+Rollout） | +0.226（In+Out vs 最佳基线），相对提升 38% |
| NLP 扰动测试（BERT-Base，IMDB） | AUC（MSE vs 遮蔽比例曲线面积） | 0.158（In+Out），0.101（Norm） | 0.089（W-AttnResLN+Mean），0.088（GlbEnc+Mean） | +0.069（In+Out vs 最佳基线），相对提升 78% |
| NLP 扰动测试（RoBERTa-Base，IMDB） | AUC（MSE vs 遮蔽比例曲线面积） | 0.122（In+Out），0.096（Norm） | 0.077（GlbEnc+Mean / W-AttnResLN+Mean） | +0.045（In+Out vs 最佳基线），相对提升 58% |
| NLP 扰动测试（ModernBert-Base，IMDB） | HS-MSE AUC | 0.112（In+Out），0.081（Norm） | 0.077（W-AttnResLN+Mean） | +0.035（In+Out vs 最佳基线），相对提升 45% |
| 视觉扰动测试（DeiT-Small，ImageNet-1K） | AUC（MSE vs 遮蔽比例曲线面积） | 0.192（In+Out），0.187（In+Class），0.151（Norm） | 0.138（W-Attn+Rollout） | +0.054（In+Out vs 最佳基线），相对提升 39% |
| 线性关系解码（Pythia-1B，多个关系类别） | Top-1 准确率（%) | 大多数关系类别上高于 LRE 基线 | LRE（Hernandez et al. 2024） | 在 sentiment、superclass、country、language 等类别上一致超越 |
局限与改进
论文明确承认了三个主要局限。第一，某些线性化技术是出于简洁性而非最优性选择的，例如激活函数的分解方式（Eq. 13）是启发式的，可能不是对张量的最优近似。第二，高阶张量表示非常消耗 GPU 显存——完整张量 $T \in \mathbb{R}^{L \times D \times L \times D}$ 包含 $L^2 D^2$ 个元素，对于大模型和长序列来说计算代价极高。论文虽然在附录 C 中提供了基于雅可比矩阵的内存高效计算方法，但仍需要 $L \cdot D$ 次前向传播来计算完整张量。第三，实验仅限于最大 1B 参数的模型和中等输入长度（128 token），在更大规模模型上的表现尚未验证。此外，我的独立观察是：对于自回归模型，扰动测试的结果不够清晰，论文将其归因于自回归模型的局部性特征，但这实际上暗示 TensorLens 的优势可能主要集中在 encoder-only 架构上，其在 decoder-only LLM 上的实用性仍需进一步验证。
独立分析的弱点
TensorLens 存在几个值得深入探讨的弱点。首先，显存瓶颈是一个实际部署的核心障碍：四阶张量 $T \in \mathbb{R}^{L \times D \times L \times D}$ 包含 $L^2 D^2$ 个元素，对于 BERT-Base（$D=768$）和 128 token 输入，这意味着约 $128^2 \times 768^2 \approx 9.4 \times 10^9$ 个浮点数，约 35GB。虽然附录 C 的雅可比方法避免了显式存储完整张量，但 $L \cdot D$ 次前向传播的计算开销仍然巨大。改进方向包括：利用张量的稀疏结构或低秩近似来压缩表示，或者设计只计算特定输出位置（如 [CLS] 或最后一个 token）的局部张量切片的高效算法。其次，激活函数的线性化（将 $\phi(Z)$ 近似为 $\text{diag}(\frac{\phi(Z)}{Z})$）是启发式的，在 ReLU 等非平滑激活函数上可能导致不稳定的近似。可以考虑使用更精确的分段线性近似或自动微分来替代。第三，实验覆盖范围有限，仅在最大 1B 参数的模型上验证，而当前主流 LLM 已达数十亿甚至数千亿参数，TensorLens 在这些模型上的可扩展性存疑。建议在更大的模型上进行小规模的概念验证实验，例如只分析前几层或选定的注意力头。
未来方向
论文在结论中暗示了几个有前景的未来方向。首先是通过张量性质来研究 rank collapse（秩坍缩）现象——Transformer 深层中隐藏表示趋于低秩的已知问题，张量的谱分析可能揭示其根本原因。其次是利用张量的稀疏性结构来理解 Transformer 的训练动态，例如哪些张量元素在训练过程中变化最剧烈，这可能为高效微调提供理论指导。第三是将 TensorLens 与其他可解释性工具结合，例如用张量表示来改进 LRP（Layer-wise Relevance Propagation）中的信息传播规则，或者用张量的低秩近似来指导模型剪枝和知识蒸馏。第四是将方法扩展到其他架构，如 Mamba、RWKV 等非 Transformer 架构，论文已经提到了隐式注意力的概念，张量化分析可能为这些架构提供统一的理解框架。最后，张量表示可能为 Transformer 的理论理解开辟新视角，例如通过张量分解来分析模型的表达能力和泛化特性。
复现评估
论文的代码已作为补充材料提交，并在 GitHub 上开源（https://github.com/idoatad/TensorLens），这大大降低了复现门槛。数据集方面，所有实验均使用公开数据集（ImageNet-1K、IMDB、WikiText-103），无需额外数据准备。算力需求方面，核心实验在 BERT-Base、DeiT-Base/Small、Pythia-1B 等较小模型上进行，单张消费级 GPU 即可完成扰动测试，但完整张量的计算（如关系解码实验）可能需要更大的显存。复现难度中等：张量化推导在数学上是严谨的，但实现时需要正确处理 Kronecker 积和张量收缩的维度映射，这容易出错。附录 C 的内存高效实现方法是复现的关键，建议优先使用基于雅可比矩阵的方法而非显式构造完整张量。总体而言，有张量代数基础的研究者可以在 1-2 周内完成核心实验的复现。
论文图表
