[논문 리뷰] Data-Aware Random Feature Kernel for Transformers

Amirhossein Farzam, Hossein Mobahi외 2명|arXiv (Cornell University)|2026. 03. 04.

Advanced Neural Network Applications인용 수 0

한 줄 요약

DARKFormer은 트랜스포머 어텐션을 위해 데이터에 맞춘 랜덤 피처 커널을 학습하여 중요도 샘플링과 유사한 분산 감소 및 선형 복잡도에서의 파인튜닝 개선을 가능하게 한다.

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.

연구 동기 및 목표

• 2차원 어텐션 비용과 등방성 랜덤 피처 방법의 높은 몬테카를로 분산을 해소한다.
• 비등방성 쿼리-키 분포에 적응하는 데이터에 맞춘 커널 기하학을 도입한다.
• 샘플별 가중치 없이 학습된 공분산을 통한 중요 샘플링의 실용적인 메커니즘을 제공한다.
• 제한된 피처 예산에서 파인튜닝 시 성능 개선과 학습 안정성을 입증한다.
• 리소스 제약 환경에서 실용성을 보이도록 Gemma 기반 모델에서 접근법을 검증한다.

제안 방법

• 기본 점곱을 학습 가능한 공분산 Sigma = M^T M을 사용한 Mahalanobis 내적으로 대체한다.
• 데이터 인식 랜덤 피처를 커널 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 DARKFormer에서 학습된 공분산이 작은 피처 예산으로 정확한 소프트맥스 어텐션과의 차이를 좁힐 수 있는가?
• RQ3 사전 학습 가중치로부터의 파인튜닝 동안 데이터 인지 커널 기하가 학습 안정성과 효율성을 향상시키는가?
• RQ4 학습률과 파인튜닝 체제에 따라 학습된 Sigma가 성능과 강인성에 어떻게 영향을 미치는가?

주요 결과

• DARKFormer는 isotropic PRF(Performer) 베이스라인에 비해 정확한 어텐션과의 성능 격차를 좁힌다.
• 대규모 피처 샘플이나 광범위한 재학습 없이도 이러한 이득을 달성한다.
• DARKFormer는 다양한 학습률 범위에서 파인튜닝 중 학습 안정성을 개선하고 손실 급등을 줄인다.
• 사전 학습 가중치에서의 자원 제약이 있는 파인튜닝에서 특히 유리하다.
• Gemma 실험은 Performer 대비 다음 토큰 예측 정확도를 개선하고 정확한 소프트맥스에 비해 경쟁력 있는 성능을 보여준다.

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