Mahalanobis-Based Multi-Head Attention for Complex State Propagation

Xiaohe Li

cs.AIAug 25, 2026

Mahalanobis-Based Multi-Head Attention for Complex State Propagation

Authors: Xiaohe Li

Organizations: GuangDong Police College

Abstract

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.

Figures & tables

Model	Parenthesis	Arithmetic+Repeat	Parity
LSTM	12.4 $\pm$ 1.8	8.7 $\pm$ 2.1	62.3 $\pm$ 1.2
GRU	15.6 $\pm$ 1.2	11.2 $\pm$ 1.8	65.1 $\pm$ 0.8
GDN	22.1 $\pm$ 1.4	18.3 $\pm$ 1.5	78.4 $\pm$ 0.5
ARFormer	38.7 $\pm$ 1.0	32.1 $\pm$ 1.2	91.2 $\pm$ 0.3
Vanilla CSP [ 26 ]	99.8 $\pm$ 0.06	30.8 $\pm$ 2.3	100.0 $\pm$ 0.0
MHA-CSP (Ours)	50.3 $\pm$ 1.6	50.3 $\pm$ 1.6	100.0 $\pm$ 0.0

Table 1: Accuracy comparison on structured state tracking tasks. All models use approximately 119K parameters. Values are mean accuracy over 3 random seeds.

Training Format	Airthetic Accuracy
Direct answer only (e.g., 8 )	34.7%
Answer with input copy (e.g., (3+5) mod 9 = 8 )	41.2%
Repeat format (e.g., 8 repeat (3+5) mod 9 = 8 )	50.3%

Table 2: Ablation on the “repeat” target format for parenthesis matching.

Model Variant	Parenthesis Accuracy
MHA-CSP (full)	50.3%
w/o tree-structured accumulation	41.8%
w/o LogSumExp correction	44.2%
w/o confusion matrix $\mathbf{B}$ (mean fusion instead)	42.6%
w/o complex normalization	38.1%

Table 3: Ablation study on parenthesis matching accuracy.

Figure 1: t-SNE visualization of digit embeddings in the complex state space. The circular arrangement indicates the model has learned the modular structure of mod-9 arithmetic.

Figure 2: PCA cumulative explained variance of hidden states from MHA-CSP trained on parenthesis matching. The first 8 principal components cumulatively explain approximately 90% of the total variance, indicating a highly structured and low-rank representation.

Attention Mechanisms

Attention mechanisms let neural networks weight information from other tokens or regions according to relevance, enabling context-dependent processing.

See all 511 papers

Aug 4, 2026

State Propagation Also Satisfies: A Complex-Valued State-Space Model for Deterministic State Tracking

Transformer-based architectures have dominated sequence modeling, largely due to the expressive power of attention mechanisms. However, for a class of deterministic state tracking tasks---such as parity checking, modular counting, and parenthesis matching---attention may be overkill. In this paper, we show that state propagation alone is sufficient. We propose the Complex State Propagator (CSP), a minimalistic recurrent architecture that only propagates hidden states across layers without output projections at intermediate steps. The state is represented as a complex-valued vector, updated via input-dependent rotations in the complex domain. To enable deep propagation without gradient vanishing or degradation, we introduce a block-level skip connection alongside element-wise complex normalization and SiLU activation at sequence boundaries. Applied with Focal Loss, CSP achieves 100% accuracy with perfect F1 scores across canonical tasks.

Xiaohe Li, Yang Lu

GuangDong Police College · Department of Computer Science and Technology, School of Informatics, Xiamen University

Recurrent Neural Networks State Tracking State Space Models

Ask PDF

Oct 9, 2026

Data-Adaptive Mahalanobis Metric Learning for Cross-Head Attention in Transformers

In this paper, we propose Data-Adaptive Mahalanobis Cross-Head Attention (DAMCHA), a simple and effective attention mechanism that generalizes Multi-Head Attention (MHA) through the lens of metric learning. Standard MHA implicitly learns a static, head-wise Euclidean similarity corresponding to the diagonal blocks of a unified cross-head metric matrix, discarding all off-diagonal interactions across heads. DAMCHA addresses this limitation via two key innovations. First, it learns the full cross-head metric matrix, inducing a Mahalanobis-like similarity that explicitly captures inter-head interactions beyond the restricted block-diagonal structure of MHA. Second, it parameterizes the metric matrix as an input-dependent function, enabling the attention geometry to dynamically adapt to the intrinsic manifold structure of the data rather than remaining fixed after training. We further introduce principled regularization strategies, most notably stack-wise parameter sharing, to ensure computational efficiency and stable optimization. DAMCHA serves as a drop-in replacement for MHA and its variants in various Transformer architectures, delivering superior expressive power and learning performance with comparable or reduced computational overhead. Extensive experiments demonstrate that DAMCHA-based Transformers consistently outperform competing baselines across a range of benchmark tasks. Code is available at https://github.com/erwinmsmith/DAMCHA.

Zhenke Duan, Xin Li, Jiqun Pan +3

School of Statistics and Mathematics, Zhongnan University of Economics and Law, Wuhan, China · CodeSoul.Co, Hangzhou, China · Innovation and Talent Base for Digital Technology and Finance, ZUEL +1

Multi-Head Attention Transformer Metric Learning

Ask PDF

May 21, 2026

Structured-Sparse Attention for Entity Tracking with Subquadratic Sequence Complexity

Entity tracking requires maintaining and updating latent states for entities and attributes over long sequences. Recent task-specific attention operators can compress deep Transformer stacks into a few layers by performing multi-hop state propagation within a single layer, but their dense evaluation remains expensive. We show that in this setting, learned attention is strongly structured: most mass concentrates in local block-diagonal neighborhoods with a light cross-block residue. Exploiting this, we derive a blockwise evaluation of a resolvent-style operator that keeps within-block interactions exact and routes cross-block interactions through a reduced system. The resulting evaluation is subquadratic in sequence length

O(n^{4/3}d)

(and

O(n^{7/3})

when

d\approx n

). On controlled tracking benchmarks, our method matches the dense operator's accuracy while reducing wall-clock time by

12-29\%

under a standardized measurement protocol, and is up to

2.4 \times

faster than a compact dense Transformer at comparable exact-match accuracy. We further provide ablations over block size and model capacity, and identify a limitation: performance collapses when the number of simultaneously evolving properties exceeds the number of attention heads.

Hangyue Zhao, Paul Caillon, Erwan Fagnou +1

ESPCI PSL, Paris, France · LAMSADE, Université Paris Dauphine - PSL, Paris, France

Self-Attention Structured Sparsity Transformer Attention Efficient Attention Sparse Attention State Tracking

Ask PDF

Abstract

Attention Mechanisms

Explore similar work

State Propagation Also Satisfies: A Complex-Valued State-Space Model for Deterministic State Tracking

Data-Adaptive Mahalanobis Metric Learning for Cross-Head Attention in Transformers

Structured-Sparse Attention for Entity Tracking with Subquadratic Sequence Complexity