Dual Dimensionality for Local and Global Attention¶
会议: NeurIPS 2026
arXiv: 2606.18587
OpenReview: https://openreview.net/forum?id=Zujn6IDCvy
论文: https://arxiv.org/abs/2606.18587
领域: LLM效率
关键词: 注意力机制, 距离自适应表征, KV缓存压缩, 局部与全局注意力, 双重维度
一句话总结¶
针对自回归语言模型中所有历史 token 均采用同等 KV 表征维度的冗余问题，提出距离自适应表征（DAR）架构，在局部上下文窗口内保留全维度表征，而在窗口外采用低维瓶颈表征，在 70M 至 410M 预训练及 1B 微调尺度上均以大幅更低的远端维度比肩乃至超越全维度基线。
研究背景与动机¶
标准 Decoder-only Transformer 的自注意力计算建立在一个隐含的架构假设之上：所有历史 token 无论距离当前预测目标多远，都贡献相同维度的 Key 和 Value 隐状态。这种设计对每个历史位置赋予了一视同仁的表征容量。然而在自然语言生成的内在统计规律中，距离当前位置极近的局部上下文承担着维持短语语法规则、避免局部词汇重复以及保持即时语义连贯等细粒度约束，对预测具有决定性的直接影响；相反，更久远的上下文主要作为粗粒度的长程语义记忆背景存在。
为了压缩日益沉重的 KV 缓存，既有方案主要分为两派：一类基于滑动窗口或驱逐策略（如 StreamingLLM、H2O、SnapKV），通过彻底丢弃非关键远端 token 实现稀疏化，但这些策略无一例外都强制保留近期局部窗口，印证了局部 token 具备更高的关键度；另一类如 DeepSeek 提出的 MLA（Multi-head Latent Attention）以及 CSA，在预训练阶段将 token 投影到统一的低秩隐空间，但它们对所有历史 token 均施加相同压缩率，完全忽视了距离维度的非对称需求。局部与长程 token 在表征容量需求上的天然差异，与当前注意力机制中均一化维度分配之间存在着明显的结构性矛盾。
本文由此提出距离自适应表征（Distance-Adaptive Representation, DAR）：其核心 idea 是打破注意力的维度均一性假设，按 token 距当前预测位置的距离非对称分配表征容量——局部窗口内的 token 保留全维度以保留细粒度辨别力，而超出窗口的远距离 token 则压缩至低维瓶颈空间，从而在无损甚至提升建模困惑度与下游推理表现的同时显著减少远端表征容量需求。
方法详解¶
整体框架¶
DAR 引入了一种双重维度分区机制，将所有被注意力的历史 token 划分为局部窗口与远端全局两个路径。对于给定的序列，局部滑动窗口（大小为 \(w\)）内的 token 保持模型原始的全维度隐藏状态；而窗口之外的 token 则通过轻量线性降维投影压缩至低维瓶颈空间 \(d_{\text{down}}\)。为了保持注意力计算接口的一致性并共享后续的 Key 与 Value 投影矩阵，模型在全局路径上将低维表征线性升维回模型维度，再与局部路径统一计算自注意力分数与加权输出。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["输入序列隐藏状态 h_j"] --> B{"距离判断<br/>i - j < w ?"}
B -->|是: 局部路径| C["1. 局部全维直通路径<br/>保留原维度 h_j 进注意力"]
B -->|否: 全局路径| D["2. 远端瓶颈压缩与升维对齐<br/>W_down 降维至 d_down 再由 W_up 对齐"]
C --> E["3. 混合注意力键值构建<br/>局部与远端分支拼接生成统一 K 与 V"]
D --> E
E --> F["标准注意力输出 Softmax(QK^T / sqrt(d_k)) V"]关键设计¶
1. 远端低维瓶颈投影：剥离长程记忆中的容量冗余
针对远距离 token 仅需提供长程语义背景、过高维度会导致参数冗余这一痛点，DAR 并不对其进行硬性丢弃，而是设计了极轻量的下投影变换： $$ \mathbf{h}{j}^{D} = \mathbf{h} $$ 其中 } \mathbf{W}_{\text{down}\(\mathbf{h}_{j} \in \mathbb{R}^{d}\) 为当前第 \(j\) 个位置的隐状态，\(\mathbf{W}_{\text{down}} \in \mathbb{R}^{d \times d_{\text{down}}}\) 为瓶颈权重矩阵，\(d_{\text{down}} < d\) 为瓶颈维度。在消融探索中，作者尝试在降维后施加 Sigmoid 非线性激活，但实测发现线性映射已具备完全等价的表征能力。因此，DAR 采用纯线性投影，在推理与训练中将远端 token 的内在信息瓶颈精确限制在 \(d_{\text{down}}\) 维度空间内。
2. 混合注意力键值构建：共享投影矩阵下的跨维度对齐
为了使局部全维度 token 与远端低维 token 能够同时被当前位置 \(i\) 的 Query 所注意，且不破坏原本多头注意力机制对 \(\mathbf{W}_K\) 与 \(\mathbf{W}_V\) 参数的通用复用，DAR 引入了上投影矩阵 \(\mathbf{W}_{\text{up}} \in \mathbb{R}^{d_{\text{down}} \times d}\) 对全局路径上的瓶颈向量进行几何空间对齐： $$ \mathbf{h}{j}^{\prime} = \mathbf{h}}^{D} \mathbf{W{\text{up}} $$ 该升维操作并不会恢复瓶颈所丢失的高频信息，其核心机制在于将降维后的长程语义信息对齐到 \(\mathbf{W}_K\) 与 \(\mathbf{W}_V\) 的输入空间。根据距离 \(i - j\) 与窗口大小 \(w\) 的相对关系，历史 token 的 Key 与 Value 按分支条件构建： $$ \mathbf{k}} = \begin{cases} \operatorname{RoPE}(\mathbf{h{j}\mathbf{W}}), & \text{if } i - j < w \ \operatorname{RoPE}(\mathbf{h{j}^{\prime}\mathbf{W}}), & \text{otherwise} \end{cases}, \quad \mathbf{v{j} = \begin{cases} \mathbf{h}}\mathbf{W{V}, & \text{if } i - j < w \ \mathbf{h} $$ 结合旋转位置编码（RoPE），得到拼接后的统一键缓存 }^{\prime}\mathbf{W}_{V}, & \text{otherwise} \end{cases\(\mathbf{K}_i\) 与值缓存 \(\mathbf{V}_i\)，直接代入标准缩放点积注意力公式完成查询。
损失函数 / 训练策略¶
模型训练严格遵循标准自回归语言建模的交叉熵目标，无需任何辅助损失或对比正则： $$ \mathcal{L} = -\sum_{t=1}^{T} \log P(x_{t} \mid x_{<t}; \theta) $$ 针对预训练任务，所有参数从零初始化并在 The Pile 10B 子集上端到端反向传播更新；针对在 OLMo-2-1B-SFT 基础上的持续微调，作者设计了两阶段渐进微调策略： 1. 阶段一（瓶颈预热）：冻结除瓶颈矩阵 \(\{\mathbf{W}_{\text{down}}, \mathbf{W}_{\text{up}}\}\) 之外的所有网络权重，使用学习率 \(3 \times 10^{-4}\) 在 Tülu 3 混合数据上训练 1 个 epoch，使瓶颈层迅速适应并学会提炼远端低维有效表征； 2. 阶段二（全量联合微调）：放开全网参数，以较小的学习率 \(3 \times 10^{-5}\) 联合微调 1 个 epoch，促使全模型各层协同适应这种局部全维、远端低维的双重注意力计算范式。
实验关键数据¶
主实验¶
在 Pythia-70M 架构（隐藏维度 \(d=512\)）下，固定局部窗口大小 \(w=128\)，评估不同瓶颈维度 \(d_{\text{down}}\) 的 DAR 模型，并与基线 Vanilla（\(d=512\)）及全局统一降维的 Uniform Reduction 进行困惑度对比。
| 模型配置 | 瓶颈维度 \(d_{\text{down}}\) | FineWeb-Edu (PPL) | C4 (PPL) | WikiText-103 (PPL) | 相对 Vanilla (%) |
|---|---|---|---|---|---|
| Vanilla | \(d=512\) (基准) | 120.41 | 146.91 | 53.22 | 100.00 |
| DAR | \(d_{\text{down}}=256\) | 118.88 | 144.05 | 52.64 | 98.57 |
| DAR | \(d_{\text{down}}=128\) | 120.02 | 144.65 | 53.58 | 99.61 |
| DAR | \(d_{\text{down}}=64\) | 123.03 | 148.10 | 54.81 | 101.99 |
| Uniform reduction | \(d_{\text{down}}=256\) | 124.37 | 150.49 | 54.36 | 102.63 |
| Uniform reduction | \(d_{\text{down}}=128\) | 127.38 | 155.35 | 55.84 | 105.49 |
| Uniform reduction | \(d_{\text{down}}=64\) | 135.36 | 162.61 | 58.96 | 111.30 |
扩展到 Pythia-160M 与 Pythia-410M 模型时，将远端瓶颈维度固定为 \(d_{\text{down}} = d / 4\)，\(w=128\)。测试表明 DAR 在跨尺度预训练中表现极其稳健：
| 模型尺度 | 配置模式 | FineWeb-Edu (PPL) | C4 (PPL) | WikiText-103 (PPL) | 相对 Vanilla (%) |
|---|---|---|---|---|---|
| Pythia-70M | Vanilla (\(d=512\)) | 120.41 | 146.91 | 53.22 | 100.00 |
| Pythia-70M | DAR (\(d_{\text{down}}=128\)) | 120.02 | 144.65 | 53.58 | 99.61 |
| Pythia-160M | Vanilla (\(d=768\)) | 94.84 | 115.94 | 43.03 | 100.00 |
| Pythia-160M | DAR (\(d_{\text{down}}=192\)) | 95.47 | 116.25 | 43.77 | 100.88 |
| Pythia-410M | Vanilla (\(d=1024\)) | 81.74 | 99.69 | 37.90 | 100.00 |
| Pythia-410M | DAR (\(d_{\text{down}}=256\)) | 79.59 | 97.72 | 37.35 | 97.98 |
消融实验¶
在 Pythia-70M（\(d_{\text{down}}=128\)）下对局部窗口大小 \(w\) 进行系统消融，考察非对称维度的空间敏感性：
| 模型配置 | 窗口大小 \(w\) | FineWeb-Edu (PPL) | C4 (PPL) | WikiText-103 (PPL) | 相对 Vanilla (%) |
|---|---|---|---|---|---|
| Vanilla | – | 120.41 | 146.91 | 53.22 | 100.00 |
| DAR | \(w=256\) | 120.85 | 145.90 | 53.79 | 100.25 |
| DAR | \(w=128\) | 120.02 | 144.65 | 53.58 | 99.61 |
| DAR | \(w=64\) | 119.77 | 144.55 | 52.94 | 99.11 |
| DAR | \(w=16\) | 120.72 | 145.31 | 53.67 | 100.01 |
| DAR | \(w=4\) | 120.90 | 145.68 | 52.02 | 100.36 |
| DAR | \(w=1\) | 122.26 | 145.64 | 54.16 | 100.86 |
| DAR | \(w=0\) (退化为 Uniform) | 127.38 | 155.35 | 55.84 | 105.49 |
在 OLMo-2-1B-SFT 上的下游任务评测（\(w=128\)，基线为无降维瓶颈架构 \(d_{\text{down}}=2048\)）：
| 瓶颈配置 | MMLU | HellaSwag | CSQA | GSM8K | MBPP | Multi-News | 平均分 | 相对表现 (%) |
|---|---|---|---|---|---|---|---|---|
| \(d_{\text{down}}=2048\) | 0.4044 | 0.6502 | 0.5315 | 0.4450 | 0.1320 | 0.1653 | 0.3881 | 100.00 |
| \(d_{\text{down}}=1024\) | 0.3989 | 0.6493 | 0.5214 | 0.4496 | 0.1420 | 0.1675 | 0.3881 | 101.09 |
| \(d_{\text{down}}=512\) | 0.4051 | 0.6462 | 0.5304 | 0.4114 | 0.1440 | 0.1855 | 0.3871 | 102.19 |
| \(d_{\text{down}}=256\) | 0.3978 | 0.6440 | 0.5228 | 0.4182 | 0.1180 | 0.1867 | 0.3812 | 98.68 |
| \(d_{\text{down}}=128\) | 0.4067 | 0.6375 | 0.5209 | 0.4008 | 0.0780 | 0.1388 | 0.3638 | 88.29 |
| \(d_{\text{down}}=64\) | 0.3891 | 0.6338 | 0.5012 | 0.3907 | 0.0300 | 0.1545 | 0.3499 | 82.00 |
关键发现¶
- 非对称容量假设成立：远端 token 的维度可在压缩至原维度的 \(1/4\)（\(d_{\text{down}}=128\)）时达到 99.61% 的优异困惑度，甚至在 \(d_{\text{down}}=256\) 时凭借正则化效果实现 98.57% 的反超；反观在所有 token 上统一应用降维的 Uniform 模型，困惑度在 \(d_{\text{down}}=128\) 时急剧恶化至 105.49%，证明保留近邻全维度对于维系语言建模能力至关重要。
- 局部窗口的极小阈值：窗口消融显示，只要保留 \(w \ge 4\) 个邻近 token 的全维度，困惑度即可稳定在 Vanilla 相当的水平（100.36%）；仅当完全抹除局部窗口（\(w=0\)）时模型发生明显退化（105.49%），证明极少数邻近 token 就足以承担起绝大部分高频细粒度建模职责。
- 下游任务的安全压缩界限：在 1B 模型上，远端表征维度压缩至 \(d/4\) 乃至 \(d/8\) 时，下游综合得分依然能维持在 98% 以上甚至略微提升（\(d/4\) 时达 102.19%），仅当压缩跨过 \(d/16\) 阈值后，代码（MBPP 从 0.132 骤降至 0.078）与长文本摘要（Multi-News）等依赖长程精确记忆的任务才出现明显断崖。
亮点与洞察¶
- 颠覆注意力维度的均一性认知：指出了序列模型长期被忽视的结构性先验——自然语言上下文对表征维度的需求随时间距离呈现阶梯式衰减，为注意力机制重新分配显存与计算开辟了新的设计空间。
- 即插即用的渐进微调策略：展示了即使不在大规模预训练阶段从零引入，仅需对现成开源大模型（如 OLMo-2-1B）施加两阶段微调（冻结骨干训练瓶颈 \(\to\) 全量微调），即可无损无缝注入距离自适应能力。
- 与推理加速硬件的理论契合度：远端 \(d_{\text{down}}\) 向量可直接缓存，仅局部保留全维 KV，显存占用直接从 \(\mathcal{O}(Td)\) 降低至渐进 \(\mathcal{O}(T d_{\text{down}} + wd) \sim \mathcal{O}(T d_{\text{down}})\)，且能与 Decoupled RoPE 和 FlashAttention 机制在工程层面深度融合。
局限与展望¶
- 实验实现的显存与吞吐开销：为了控制变量并验证基础假设，当前实验代码在推演时同时冗余保留了双重 Key/Value 状态并在全局路径上实施重投影，未直接在底层定制融合算子以释放推理显存红利。
- 模型尺度与架构覆盖面受限：由于算力约束，预训练探索主要限制在 70M 至 410M 的 Pythia 骨干上，微调验证止步于 1B 规模，尚未在 7B-70B 级前沿大模型及海量超长上下文（如 32k-128k）上完成全流程实证。
- 未来方向：探索连续渐变衰减的维度调度（而非简单的二段划分），并结合硬件友好的 CUDA 内核实现极致精简的生产级 KV 缓存管理系统。
相关工作与启发¶
- vs Sliding-Window / StreamingLLM / H2O: 传统局部窗口与驱逐方案直接对远端 token 判处“全有或全无”的硬丢弃或稀疏截断，极易破坏长程依赖上下文；DAR 选择“全维近邻 + 低维远端”的软压缩策略，在保全全局语义连贯的同时压缩容量。
- vs DeepSeek MLA / CSA: MLA 将低秩压缩无差别施加在整条序列的所有 token 上，忽略了空间局部性；DAR 证明了距离非对称性维度的优越性，两者结合有望在 MLA 的低秩空间上进一步压榨远端 KV 开销。
- vs Matryoshka Representation Learning (MRL / MatryoshkaKV): MRL 侧重于利用嵌套表征以适配外部动态计算约束，而 DAR 则是立足于自然语言生成本身的内在距离先验，属于模型内部自适应容量配置。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ 首次系统性论证注意力键值表征容量与序列距离的非对称相关性并提出 DAR 范式。
- 实验充分度: ⭐⭐⭐⭐☆ 涵盖从 70M 到 410M 从零预训练、多窗口消融以及 1B 模型下游 6 个基准的全面验证。
- 写作质量: ⭐⭐⭐⭐⭐ 逻辑结构严密，动机阐述自然连贯，图表呈现扎实清晰。
- 价值: ⭐⭐⭐⭐☆ 为下一代长文本大模型的高效注意力架构演进与极限 KV 缓存压缩提供了高度实用的洞见。