L1 Augmented Attention as an Improved Vector Similarity Metric

Standard Transformer models rely on scaled dot-product attention to measure similarity between query and key vectors. However, the dot product conflates directional alignment with vector magnitude, potentially masking important geometric information. This paper investigates whether augmenting the attention mechanism with an L1 distance metric can provide a more expressive and accurate similarity measure for language modeling.

The author introduces L1 augmented attention, which modifies the standard attention equation by subtracting a learned, head-specific L1 distance between queries and keys from the dot-product score. To mitigate the computational overhead of calculating L1 distances, the author projects queries and keys into lower-dimensional subspaces. These projections are learned during training, allowing the model to specialize its representation of geometric structure. The method is evaluated on the WikiText-2 dataset using a compact Transformer architecture, comparing performance against the standard Vaswani baseline and an RBF-L2 kernel approach.

L1 augmented attention consistently outperforms the standard dot-product baseline and the RBF-L2 kernel in language modeling tasks. The best-performing model, which used an 8-dimensional projection for the L1 term, achieved a 14.5% reduction in perplexity compared to the baseline. Analysis of the learned weights suggests that different attention heads develop specialized geometric roles, effectively leveraging the complementary nature of dot-product (directional) and L1 (coordinate-wise) metrics to better represent token relationships.

This research demonstrates that the standard dot-product attention used in modern LLMs is not an infallible similarity metric. By enriching the attention mechanism with L1 geometry, researchers can improve model accuracy without abandoning the core Transformer architecture. The use of low-dimensional projections makes this approach computationally feasible, offering a principled way to enhance the expressivity of attention heads in resource-constrained environments.