高达1000倍的长度外推能力——仅通过替换softmax即可实现。密集注意力会在长输入中分散。ASEntmax为α-entmax引入了可学习的温度参数，能将不相关的token注意力权重精确归零，在稀疏与密集模式间灵活切换。在训练长度8倍的条件下实现更优检索性能。https://t.co/7U7qOjeUX0

基于稀疏注意力的长上下文泛化

来源：https://arxiv.org/html/2506.16640

作者：Pavlo Vasylenko（里斯本高等技术学院、电信研究所）、André F. T. Martins（里斯本高等技术学院、电信研究所、TransPerfect、ELLIS里斯本单元）、Marcos Treviso（里斯本高等技术学院、电信研究所、ELLIS里斯本单元）

摘要

传统Transformer架构采用softmax计算注意力权重，产生基于序列中所有token的稠密分布。虽然在许多场景下有效，但这种稠密性对于需要精准聚焦固定模式的任务存在弊端：随着序列长度增加，非信息性token不断累积注意力概率质量，导致注意力分散与表征坍缩。本文证明，使用α-entmax的动态稀疏注意力机制能规避这些问题，因其可为无关token分配精确的零概率值。进一步，我们提出自适应可扩展Entmax（ASEntmax），该方法为α-entmax赋予可学习的温度参数，使注意力分布能在稀疏（模式聚焦）与密集（类softmax）模式间插值。在合成任务与语言建模的实证评估中，ASEntmax显著优于softmax、可扩展softmax及固定温度α-entmax基线方法，在合成基准测试中实现高达1000倍长度外推，并在保持短上下文性能的同时获得优越的长上下文泛化能力（包括在训练长度8倍条件下更优的困惑度趋势与更高检索准确率）。

源代码：https://github.com/deep-spin/asentmax

1 引言

Transformer架构（Vaswani et al., 2017）已成为现代大语言模型（LLM）的基石，在多个领域建立新基准。然而，当研究者将这些模型推向越来越长的上下文——从数千到数百万token——时，若干基本局限性逐渐显现，这些局限可追溯至注意力机制中使用的softmax变换。三个关键缺陷尤为突出：

表征坍缩源于softmax在序列长度增长时无法维持清晰的注意力模式，导致token间的有意义区分被消除（Barbero et al., 2024）；过度平滑因softmax的稠密概率分布而加剧，导致梯度指数级衰减（Alon & Yahav, 2021; Barbero et al., 2024）；注意力分散则源自softmax的基本属性——迫使概率质量必须分布于所有token，随着上下文增长，注意力权重必然趋向均匀分布（Veličković et al., 2025; Nakanishi, 2025）。

现有解决方案包括ALiBi（Press et al., 2022）和RoPE（Su et al., 2024）等位置编码创新，有助于缓解位置偏差问题。近期研究则直接针对核心症结——softmax函数本身。Nakanishi（2025）提出可扩展Softmax，基于上下文长度缩放logit值；Veličković et al.（2025）则指出softmax在分布外泛化中的根本局限，提出通过学习自适应温度来控制softmax的锐度。尽管这些方案有效，但通常需要精细调参或仅能解决部分挑战。

本文通过用α-entmax（Peters et al., 2019）替换softmax来解决问题根源——这是一种可微分的稀疏变换，能为无关token分配精确的零注意力。尽管α-entmax已在Transformer中成功应用（Correia et al., 2019; Gonçalves et al., 2025），但据我们所知，其长度泛化特性从未被系统研究。我们从理论和实证两个层面证明，α-entmax能有效应对长上下文建模中的挑战。

主要贡献包括：

• 抗分散性：我们证实α-entmax注意力分布能保持稳定的焦点集中度，其熵由O(log s)约束（其中s是具有非零概率的token数量），而非像softmax那样趋向最大熵O(log n)（s≪n）。

• 表征保持：α-entmax的稀疏支持可避免表征坍缩，将梯度路径数从O(n^L)降低至O(s^L)，通过增强长程依赖的梯度流缓解过度平滑。

• 自适应可扩展α-entmax：我们提出ASEntmax，能根据序列长度自适应调整稀疏性，即使在极长上下文中也能保持最优的token选择。

• 实证结果：我们证明ASEntmax在合成与真实任务中均取得卓越性能。如图1所示，ASEntmax仅通过64个token训练，在65K token的关联记忆任务中达到95.3%准确率——实现1000倍长度外推。

图1：多查询多token关联记忆（左）与最大检索（右）的长上下文泛化对比。SSMax代表Nakanishi（2025）的可扩展Softmax方法，自适应温度（Adapt. Temp.）代表Veličković et al.（2025）的方法。尽管所有方法都受益于使用NAPE（无位置编码+ALiBi），但我们的自适应缩放α-entmax版本展现出最佳外推效果，能有效处理极长序列。

2 背景

2.1 Transformer

本文研究（因果）Transformer中通过替换softmax为α-entmax产生的稀疏注意力分布。以下数学表述严格遵循Barbero et al.（2024）的符号规范。

给定token嵌入序列X∈ℝ^(n×d)（n为序列长度，d为隐藏维度），Transformer计算查询、键和值投影：Q=X****W_Q、K=X****W_K、V=X****W_V。令q_i、k_i、v_i∈ℝ^d表示第i个token的d维查询、键、值向量。对于每个查询位置i，第l层的第i个token表示计算如下：

ui^(ℓ) = Σ{j≤i} p_{ij}^(ℓ) norm₁^(ℓ)(v_j^(ℓ-1)) + v_i^(ℓ-1) v_i^(ℓ) = FFN^(ℓ)(norm₂^(ℓ)(u_i^(ℓ))) + u_i^(ℓ)

其中p_{ij}^(ℓ)为注意力权重，FFN^(ℓ)为前馈网络，norm(·)表示LayerNorm模块（Xiong et al., 2020）。输出计算为y_i = norm₃(v_i^(L))。

注意力权重p_{ij}^(ℓ) = π(z_i^(ℓ))j通过将注意力logit z{ij}^(ℓ) = ⟨q_i^(ℓ), k_j^(ℓ)⟩/√d代入变换π: ℝ^n → Δ_n得到，其中Δ_n = {p∈ℝ^n : p≥0, 1ᵀp=1}表示概率单纯形。标准Transformer采用softmax函数作为π。本文将π设为α-entmax变换。

2.2 α-entmax

图2：不同α值的α-entmax(z/θ)可视化。每图显示当温度参数减小（θ⁻¹增大）时，概率质量在z=[2.0,1.8,1.6,1.4,1.2]五个元素上的分配情况。竖线标注了产生零概率的临界温度。

α-entmax（Peters et al., 2019）是softmax的泛化形式，能生成稀疏概率分布的可微分变换。对于输入向量z∈ℝ^n和α>1，定义为：

α-entmax(z)i = [(α-1)z_i - τ(z)]+^{1/(α-1)}

其中[·]_+ = max(0,·)，τ: ℝ^n → ℝ为阈值函数，确保结果分布总和为1。α-entmax的关键特性是：得分低于阈值的token将获得精确的零概率，形成稀疏注意力模式。当α→1⁺时退化为标准softmax函数。稀疏度随α增大而增强，α=2对应sparsemax函数（Martins & Astudillo, 2016）。

图2展示了不同α值的α-entmax行为。更多细节见附录A。尽管α-entmax是稀疏注意力的理想选择，但其对长输入的理论与实证影响尚不明确。下节将阐述它如何从根本上改变长上下文的注意力行为。

3 α-entmax在长上下文中的理论特性

本节分析α-entmax使长上下文建模更具优势的理论特性，重点阐述其如何解决softmax的根本局限。

3.1 注意力概率非衰减性

Transformer中softmax的关键局限在于注意力权重必然随序列长度增加而衰减。首个结论展示α-entmax如何避免此问题。

引理1 [注意力非衰减特性] 考虑标量a₁,…,a_{n-1}, c∈ℝ。令x=[a₁,…,a_{n-1}, c]ᵀ∈ℝ^n和x^*=[a₁,…,a_{n-1}, b, c]ᵀ∈ℝ^{n+1}（所有元素有界）。以下性质成立：

• 对所有α≥1，有α-entmax(x)n ≥ α-entmax(x^*){n+1}
• 对所有α>1，存在b_max∈ℝ，使得对任意b≤b_max，有α-entmax(x)n = α-entmax(x^*){n+1}
• 对α>1，差值α-entmax(x)n - α-entmax(x^*){n+1}可通过选择合适的b取[0, α-entmax(x)_n]中的任意值

（证明见附录D.1）

此结果揭示了softmax与α-entmax的根本区别。softmax中添加新token必然严格降低现有token的注意力概率，而α-entmax允许不同行为：注意力概率可保持不变。这是因为α-entmax的阈值效应使低于特定阈值的token获得精确零注意力，让模型仅聚焦相关token。

在确定α-entmax能防止单个注意力权重衰减后，我们进一步形式化“注意力分散“的广义概念，以理解注意力分布整体随序列长度增长的行为。

3.2 注意力分散与集中

Nakanishi（2025）和Veličković et al.（2025）近期研究强调注意力分散是softmax在长上下文泛化中的根本局限。基于这些洞见，我们给出注意力分散的形式化定义，并展示α-entmax如何自然具备解决这些局限的集中特性。

定义1 [注意力分散] 令f: ℝ^n → Δ_n为将logit映射到概率单纯形Δ_n = {p∈ℝ^n : p≥0, 1ᵀp=1}的变换（如softmax）。

1. 若对任意有界logit序列(zn){n∈ℕ}，归一化熵随序列长度增长趋近1：lim_{n→∞} H(f(z_{1:n}))/log n = 1，则称f具有完全分散性。
2. 若存在有界logit序列使归一化熵远离1保持有界：lim_{n→∞} H(f(z_{1:n}))/log n < 1，则称f具有集中韧性。

这些定义使我们能分析softmax和α-entmax在序列增长时的行为：

命题1 [注意力机制的分散特性] 比较softmax与α-entmax（α>1）注意力机制：

1. α-entmax…