# Mahalanobis-Based Multi-Head Attention for Complex State Propagation

arXiv: https://arxiv.org/abs/2608.24462 · HTML: https://mlsanity.com/paper/2608.24462/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 26 of its 48 references that are in the map.
- First posted on arXiv 25 Aug 2026.

## Builds on (most cited first, 26 of 26)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Mamba: Linear-Time Sequence Modeling with Selective State Spaces](https://mlsanity.com/paper/2312.00752/) (2023)
- [Learning Phrase Representations using RNN Encoder-Decoder for Statistical Machine Translation](https://mlsanity.com/paper/1406.1078/) (2014)
- [Efficiently Modeling Long Sequences with Structured State Spaces](https://mlsanity.com/paper/2111.00396/) (2021)
- [Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention](https://mlsanity.com/paper/2006.16236/) (2020)
- [Rethinking Attention with Performers](https://mlsanity.com/paper/2009.14794/) (2020)
- [Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality](https://mlsanity.com/paper/2405.21060/) (2024)
- [DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model](https://mlsanity.com/paper/2405.04434/) (2024)
- [What Does BERT Look At? An Analysis of BERT's Attention](https://mlsanity.com/paper/1906.04341/) (2019)
- [Analyzing Multi-Head Self-Attention: Specialized Heads Do the Heavy Lifting, the Rest Can Be Pruned](https://mlsanity.com/paper/1905.09418/) (2019)
- [Progress measures for grokking via mechanistic interpretability](https://mlsanity.com/paper/2301.05217/) (2023)
- [Are Sixteen Heads Really Better than One?](https://mlsanity.com/paper/1905.10650/) (2019)
- [Simplified State Space Layers for Sequence Modeling](https://mlsanity.com/paper/2208.04933/) (2022)
- [Grokking: Generalization Beyond Overfitting on Small Algorithmic Datasets](https://mlsanity.com/paper/2201.02177/) (2022)
- [Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention](https://mlsanity.com/paper/2502.11089/) (2025)
- [Diagonal State Spaces are as Effective as Structured State Spaces](https://mlsanity.com/paper/2203.14343/) (2022)
- [Revealing the Dark Secrets of BERT](https://mlsanity.com/paper/1908.08593/) (2019)
- [Transformer Dissection: A Unified Understanding of Transformer's Attention via the Lens of Kernel](https://mlsanity.com/paper/1908.11775/) (2019)
- [Tree Transformer: Integrating Tree Structures into Self-Attention](https://mlsanity.com/paper/1909.06639/) (2019)
- [Multi-Head Attention: Collaborate Instead of Concatenate](https://mlsanity.com/paper/2006.16362/) (2020)
- [Insights into DeepSeek-V3: Scaling Challenges and Reflections on Hardware for AI Architectures](https://mlsanity.com/paper/2505.09343/) (2025)
- [Tree-structured Attention with Hierarchical Accumulation](https://mlsanity.com/paper/2002.08046/) (2020)
- [Weight decay induces low-rank attention layers](https://mlsanity.com/paper/2410.23819/) (2024)
- [DeFT: Decoding with Flash Tree-attention for Efficient Tree-structured LLM Inference](https://mlsanity.com/paper/2404.00242/) (2024)
- [Tree Attention: Topology-aware Decoding for Long-Context Attention on GPU clusters](https://mlsanity.com/paper/2408.04093/) (2024)
- [State Propagation Also Satisfies: A Complex-Valued State-Space Model for Deterministic State Tracking](https://mlsanity.com/paper/2608.03425/) (2026)

## Abstract (arXiv)

In this paper, we propose Mahalanobis-Based Multi-Head Attention (MHA-CSP), a novel attention mechanism that replaces the standard dot-product with a Mahalanobis distance-based RBF kernel, which effectively computes attention in an infinite-dimensional feature space without increasing the parameter count. Crucially, the positive definiteness of the Mahalanobis distance enables a direct construction of Tree Attention: attention scores are built directly from accumulated distances, with a LogSumExp correction that rectifies the raw distance by subtracting the log-sum of edge exponentials. Moreover, the multi-head Mahalanobis distance matrices are themselves repurposed to construct an attention meshing mechanism, enabling cross-head kernel collaboration that simultaneously boosts accuracy and training efficiency. Extensive experiments demonstrate that MHA-CSP, with only 119K parameters and teacher forcing applied exclusively at the final hidden state, consistently outperforms Transformer and GCN baselines trained from scratch under identical conditions on long-sequence state tracking tasks. While these baselines rely on dense attention or graph propagation, MHA-CSP achieves robust structured reasoning via synthetic distance rectification---powered by Mahalanobis-based attention---and efficient information bypass inherited from the CSP backbone. This result highlights the effectiveness of complex-valued state propagation with collaborative multi-head rectification in capturing symbolic structures, establishing a new efficiency-performance trade-off for structured reasoning.