论文阅读:Log-Linear Attention — 基于 Fenwick 树与层次矩阵的对数线性复杂度注意力机制

论文:Log-Linear Attention
作者:Han Guo*, Songlin Yang*, Tarushii Goel, Eric P. Xing, Tri Dao, Yoon Kim
机构:MIT, Princeton University / Together AI, CMU / MBZUAI / GenBio AI
链接:arXiv:2506.04761v3, ICLR 2026

核心贡献

本文提出 Log-Linear Attention,将线性注意力中的单一固定大小隐藏状态替换为 个随序列长度对数增长的层次化隐藏状态。通过 Fenwick 树分区和层次矩阵(HODLR)理论,实现了 训练复杂度和 解码内存的新复杂度折中。

研究动机

高效注意力的结构化矩阵统一框架

论文首先建立了一个关键观察:广泛的高效注意力机制可以统一表示为

其中 是类注意力矩阵(如 ), 是下三角因果掩码矩阵。不同模型通过对 施加不同结构约束来实现不同的效率-表达能力折中:

的结构 模型 训练 解码时间 解码内存
全 1 下三角(半可分) Linear Attention / Mamba-2
1-半可分 RetNet / Mamba-2(含门控)
Toeplitz 长卷积模型(Multi-Hyena)
准层次矩阵(Quasi-H) Log-Linear Attention
无结构 Softmax Attention

核心洞察:决定计算效率的不是 的形式(是否包含 softmax),而是 的结构。 使用无结构的 (如随机下三角矩阵)即使去掉 softmax 也无法降低复杂度。

线性注意力的根本瓶颈

线性注意力的循环形式 , 将整个前缀 压缩为单一 隐藏状态 。即使加入数据依赖门控(Mamba-2)或 delta 规则(Gated DeltaNet),固定大小的隐藏状态仍是关联回忆等需要精确检索上下文信息的任务的根本瓶颈。

本文的核心问题:能否在保持次二次训练复杂度和亚线性解码内存的前提下,突破固定大小隐藏状态的表达瓶颈?

Fenwick 树分区与多尺度隐藏状态

前缀的层次化分解

从解码视角看,注意力机制可被视为对前缀 分区-聚合过程。Softmax attention 将每个 token 独立存储( 个大小为 1 的桶),线性注意力将所有 token 合并为一个桶。Log-Linear Attention 采用 Fenwick 树分解(Ryabko, 1992; Fenwick, 1994),将前缀划分为至多 个不相交的桶 ,桶大小 ),加上一个大小为 1 的哨兵桶

分区方法:从 开始,反复减去 ,每步切出一个桶。

(二进制 111)为例:

  • :桶 ,大小
  • :桶 ,大小
  • :桶 ,大小
时间轴:  0  1  2  3  4  5  6  [7]
        |--- B⁽³⁾ ----|--B⁽²⁾--|B⁽¹⁾| B⁽⁰⁾
        |  大小 4      | 大小 2 |  1 |  1

(二进制 1000)为例,,所有 8 个历史 token 合并为单一桶 ,退化为线性注意力。 的二进制中有多少个 1,就有多少个非空桶。

层级判定算法

给定 ,可通过 Fenwick 树遍历确定 所属的桶层级

def level(t, s):
    if s == t: return 0
    cursor = t
    while cursor > 0:
        lb = cursor & (-cursor)       # lowbit
        if cursor - lb <= s < cursor:
            return lb.bit_length()    # ℓ = log₂(lb) + 1
        cursor -= lb

时各 的层级:

0 1 2 3 4 5 6 7
3 3 3 3 2 2 1 0

循环形式与输出计算

每个桶维护独立的 隐藏状态 。输出由数据依赖的非负层级权重 加权求和:

通过对当前 token 隐藏表示 施加线性投影得到,为 per-head 参数。额外参数量:Mamba-2 增加 < 3%,Gated DeltaNet 增加 < 0.4 %。

退化条件:当 在各层级间相同(或线性相关)时,Log-Linear Attention 退化为标准线性注意力。

层次矩阵 与并行训练形式

并行形式

将循环形式重写为矩阵形式以支持并行训练:

其中 由 Fenwick 树分区确定。 的结构( 简记为 ):

         s=0   s=1   s=2   s=3   s=4   s=5   s=6   s=7
  t=0: [ λ⁰                                              ]
  t=1: [ λ¹    λ⁰                                        ]
  t=2: [ λ²    λ²    λ⁰                                  ]
  t=3: [ λ²    λ²    λ¹    λ⁰                            ]
  t=4: [ λ³    λ³    λ³    λ³    λ⁰                      ]
  t=5: [ λ³    λ³    λ³    λ³    λ¹    λ⁰                ]
  t=6: [ λ³    λ³    λ³    λ³    λ²    λ²    λ⁰          ]
  t=7: [ λ³    λ³    λ³    λ³    λ²    λ²    λ¹    λ⁰    ]

HODLR 结构与低秩性

属于 HODLR(Hierarchically Off-Diagonal Low-Rank) 矩阵。其核心性质可通过递归二分验证:

矩阵沿中点二分为四个 子块。左下非对角块为:

这是一个秩 1 矩阵。两个对角子块的结构与原矩阵自相似——递归二分后左下角仍为秩 1。

这种“递归划分后每层非对角块均为低秩”的性质正是 H-matrix 的定义特征。

复杂度推导

矩阵 与向量的乘法复杂度为 而非 层低秩块(每层秩 1),因此与 的结构化矩阵乘法复杂度为

解码:基于 lssb 的状态更新

定义 的二进制中最低有效设置位的索引(即 )。隐藏状态 按如下规则更新:

该规则等价于 Fenwick 树的更新操作: 层接收当前 token 的外积; 以下的层级合并后提升一级;更高层保持不变。

为例,各时间步的状态演化:

二进制 lssb
0 0
1 01 0
2 10 1
3 11 0
4 100 2

时, 层清零,原有内容合并至 层:

每步更新涉及至多 个状态,解码时间和空间均为

训练:分块并行扫描算法

矩阵分解

给定块大小 可分解为块对角矩阵 (块内交互)与各层级块间矩阵 之和:

其中 为与块大小对齐的最低层级。 的层级折叠进

两阶段计算

块内阶段): 个大小为 的因果对角块,各块独立做稠密注意力并行处理。每块 ,总计

块间阶段):每个 为缩放的半可分结构,可复用 Mamba-2 / Gated DeltaNet 的分块并行原语。以 为例:

层级 2(相邻块间依赖):块 0 的状态 传递至块 1;块 2 的状态传递至块 3。两对独立,可并行执行。对应一次 Mamba-2 分块扫描原语,

层级 3(更远块间依赖):块 0+块 1 的合并状态 传递至块 2 和块 3。同样一次扫描原语,

输出聚合

复杂度分析

Mamba-2 Log-Linear Mamba-2
块内
块间 1 次扫描 次扫描
总计

额外开销仅为对数因子。 时块间扫描次数为 ,且每次调用的均为已有的高效 Triton 内核。

与已有模型的组合

Log-Linear Attention 保留原模型中 的参数化形式,仅将掩码组合为

其中 为原模型的半可分掩码()。半可分矩阵与 H-矩阵的 Hadamard 积仍为 H-矩阵。

任何具有结构化记忆和高效分块原语的线性注意力均可通过此方式提升为对数线性变体。

实验结果

合成任务:MQAR(多查询关联回忆)

序列长度 256,4-64 个键值对,5 个种子的平均准确率(标准差):

模型 维度 16 维度 32 维度 64
Transformer
Mamba-2 46.9 (2.3) 75.1 (4.9) 89.6 (6.1)
  1. Log-Linear
55.9 (9.1) 76.5 (4.8) 92.9 (2.7)
Gated DeltaNet 38.4 (1.0) 79.0 (2.1)
  1. Log-Linear
40.0 (1.4) 84.4 (1.2)

对数线性变体在所有设置下均优于对应的线性基线。

语言建模

500 亿 token 预训练,序列长度 16K,21 层:

模型 Wiki. ppl ↓ LMB. ppl ↓ LMEval avg ↑
Transformer (693M) 21.56 22.14 44.0
Transformer (778M, 24L) 21.13 21.17 45.6
Mamba-2 (802M) 22.44 24.14 44.8
LL Mamba-2 (825M) 22.11 21.86 44.9
Gated DeltaNet (793M) 21.73 19.71 45.0
LL Gated DeltaNet (796M) 21.45 18.09 45.5

Log-Linear Gated DeltaNet 在所有指标上优于层数匹配的 Transformer(21 层),在半数指标上优于参数匹配的 Transformer(24 层)。

长上下文评估

NIAH(Needle-In-A-Haystack):多针检索任务中 Log-Linear Gated DeltaNet 全部 9 个指标改进。

逐位置损失分析:对数线性变体在各位置上持续降低损失,表明长程上下文利用能力得到改善。

训练效率:自定义 Triton 内核在序列长度 >8K 时前后向速度超过 FlashAttention-2;完整模型在 32K 时吞吐量超过 Transformer。

总结与讨论

Softmax Attention Linear Attention Log-Linear Attention
隐藏状态数 1
解码内存
训练复杂度
的结构 无结构 半可分 准层次矩阵(Quasi-H)

Log-Linear Attention 的理论贡献在于建立了对数线性注意力与层次矩阵理论(HODLR)之间的精确联系:注意力算子等价于与准 H-矩阵的结构化矩阵乘法。Fenwick 树分区作为 H-矩阵的具体构造方案,同时满足训练端的 并行复杂度和推理端的 在线更新复杂度。

局限性方面,Fenwick 树分区引入了“近处高分辨率、远处低分辨率”的固定归纳偏置,这在 为 2 的幂时退化为单桶(等价于线性注意力)。此外,与 Transformer 相比仍存在性能差距,最优的 参数化策略有待进一步探索。

上一篇

论文阅读:InfLLM-V2 — 把可训练稀疏注意力做成全注意力的“受限版本”

下一篇

ASC26 参赛记录