[论文解读] Data-Aware Random Feature Kernel for Transformers

Amirhossein Farzam, Hossein Mobahi+2 more|arXiv (Cornell University)|Mar 4, 2026

Advanced Neural Network Applications被引用 0

一句话总结

简要结论：DARKFormer 学会了与数据对齐的随机特征核用于 Transformer 注意力，实现了类似重要性采样的方差约简，并在线性复杂度下改进微调。

ABSTRACT

Transformers excel across domains, yet their quadratic attention complexity poses a barrier to scaling. Random-feature attention, as in Performers, can reduce this cost to linear in the sequence length by approximating the softmax kernel with positive random features drawn from an isotropic distribution. In pretrained models, however, queries and keys are typically anisotropic. This induces high Monte Carlo variance in isotropic sampling schemes unless one retrains the model or uses a large feature budget. Importance sampling can address this by adapting the sampling distribution to the input geometry, but complex data-dependent proposal distributions are often intractable. We show that by data aligning the softmax kernel, we obtain an attention mechanism which can both admit a tractable minimal-variance proposal distribution for importance sampling, and exhibits better training stability. Motivated by this finding, we introduce DARKFormer, a Data-Aware Random-feature Kernel transformer that features a data-aligned kernel geometry. DARKFormer learns the random-projection covariance, efficiently realizing an importance-sampled positive random-feature estimator for its data-aligned kernel. Empirically, DARKFormer narrows the performance gap with exact softmax attention, particularly in finetuning regimes where pretrained representations are anisotropic. By combining random-feature efficiency with data-aware kernels, DARKFormer advances kernel-based attention in resource-constrained settings.

研究动机与目标

• 解决各向同性随机特征方法的二次注意力成本与高蒙特卡洛方差问题。
• 引入可适应非各向异性查询-键分布的、数据对齐的核几何。
• 提供通过学习协方差实现的可处理的重要性采样机制，而无需逐样本权重。
• 在有限特征预算下的微调中展示性能与训练稳定性的提升。
• 在 Gemma 基模型上验证该方法，以展示在资源受限场景中的实用性。

提出的方法

• 用可学习的协方差 Sigma = M^T M 将标准点积替换为马氏内积。
• 使用数据感知的随机特征，核为 exp(q^T Sigma k) 以及相应的 phi_Sigma 特征映射，omega ~ N(0, Sigma)。
• 证明学习 Sigma 产生隐式的重要性采样效应，降低蒙特卡洛方差而无需显式逐样本权重。
• 提供理论依据：方差最优采样与输入几何对齐；在高斯情形下，当 Λ 为输入协方差时，最优 Sigma* 为 (I+2Λ)(I-2Λ)^{-1}。
• 认为 DARKFormer 提供一个可处理、数据对齐的采样策略，在有限特征预算下提升性能并增强训练稳定性。
• 在 Gemma 模型上进行经验验证，聚焦于具有各向异性的查询-键分布的微调场景。

Figure 1: The random feature attention replaces the softmax kernel with a linear approximation in the feature space, reducing the quadratic complexity in sequence length ( $L$ ) to linear in sequence length times sample size ( $m$ ).

实验结果

研究问题

• RQ1 数据对齐的随机特征注意力是否降低了具有各向异性查询-键分布的蒙特卡洛方差？
• RQ2 在较小的特征预算下，学习到的协方差能否缩小与精确 softmax 注意力之间的差距？
• RQ3 数据感知核几何是否在从预训练权重微调时提升训练稳定性与效率？
• RQ4 学习得到的 Sigma 如何影响在不同学习率和微调模式下的性能和鲁棒性？

主要发现

• DARKFormer 在与精确注意力的性能差距上缩小了相对于各向同性 PRF（Performer）基线的差距。
• 在不需要大量特征样本或大规模再训练的前提下实现了这些提升。
• DARKFormer 在微调过程中在一系列学习率下提升训练稳定性并降低损失尖峰。
• 在从预训练权重进行资源受限的微调时，该方法尤为有利。
• Gemma 的实验显示在下一个词预测准确性方面优于 Performer，并在与精确 softmax 相比的性能上具有竞争力。

Figure 2: Next token prediction accuracy during pretraining (top) and finetuning (bottom) of the Gemma-2B model with a DARKFormer (green), a Performer (orange), learned feature kernel (LFK) (blue), a random baseline (yellow), a constant baseline (lime), and an exact softmax attention. The DARKFormer