A Coulomb Particle Model for Learning Kernel Attention in Transformers

2026-07-26

AI Summary

Existing random feature methods suffer from limited kernel approximation performance due to their reliance on fixed distributions, which hinders task adaptivity. This work proposes a novel particle optimization framework that, for the first time, integrates Riesz/Coulomb repulsive potentials with kernel target alignment. By introducing learnable repulsive interactions, the approach enhances feature diversity and is underpinned by a rigorous McKean–Vlasov mean-field theory. The resulting task-adaptive random features are seamlessly incorporated into a linearized Transformer attention mechanism, preserving linear inference complexity while significantly improving accuracy, calibration, and robustness on both synthetic classification and sentence-level tasks.

Abstract

Randomized features provide a scalable approximation to kernel machines, but their performance depends strongly on the choice of feature distribution. We propose a particle-based method that learns this distribution by optimizing kernel-target alignment while regularizing particles with a Riesz/Coulomb repulsive potential. The resulting Hamiltonian yields diverse, task-adaptive random features and admits a mean-field description through a McKean--Vlasov equation. We instantiate the method in linearized Transformer attention by learning positive random-feature maps in a first alignment phase, then freezing the kernel and training the remaining network parameters with cross-entropy. Experiments on synthetic classification and sentence-level benchmarks show that learned kernelized attention can improve accuracy, calibration, and robustness for several feature maps while preserving linear-attention inference complexity.

Problem

• random features
• kernel alignment
• feature distribution
• Transformer attention
• kernel machines

Innovation

• Coulomb particle model
• kernel attention
• random features
• kernel-target alignment
• linearized Transformer