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

Zhenke Duan; Xin Li; Jiqun Pan; Dong Xiaofei; Hanwen Ning; Xinyuan Song

cs.LGOct 9, 2026

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

Authors: Zhenke Duan, Xin Li, Jiqun Pan, Dong Xiaofei, Hanwen Ning, Xinyuan Song

Organizations: 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 · Department of Statistics, The Chinese University of Hong Kong, Hong Kong SAR, China

Abstract

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.

Figures & tables

Figure 1: The workflows of DAMCHA and MHA.

Figure 2: A basic DAMCHA-based decoder-only Transformer architecture with stack-wise sharing.

Figure 3: (a) Distribution of BARTScore on WMT data. (b) Distribution of NLL on WMT data.

Image generation		MSE $\downarrow$			FID $\downarrow$			FID@Top-50 $\downarrow$	Text generation		NLL $\downarrow$			BARTScore $\uparrow$			BARTScore@Top-50 $\uparrow$
Dataset	Method	mean	q25	q75	mean	q25	q75		Dataset	Method	mean	q25	q75	mean	q25	q75
CIFAR-10	MHA	0.1260	0.0750	0.1610	358.8570	293.0850	417.5810	91.4320	WMT	MHA	5.7730	5.1760	6.3750	-8.4740	-9.4590	-7.3420	-5.3090
	THA	0.1357	0.0827	0.1726	366.7536	301.3989	421.4758	107.2156		THA	5.2368	4.5104	5.7082	-6.7477	-7.5275	-5.8776	-4.5788
	DCMHA	0.1264	0.0758	0.1622	366.0802	297.8571	418.4777	105.1269		DCMHA	5.2461	4.5396	5.7308	-6.8371	-7.5254	-6.0639	-4.8286
	ColMHA	0.1276	0.0752	0.1630	365.9343	298.6780	416.1373	120.1717		ColMHA	5.2432	4.5387	5.7324	-6.1477	-6.7462	-5.4902	-4.3363
	MMA	0.1292	0.0771	0.1604	359.6663	267.3268	422.7725	84.9717		MMA	5.0301	4.5096	5.6847	-8.3875	-9.3292	-7.3488	-5.7903

Table 1: Comparison with MHA and recent multi-head attention variants on generation tasks. For MSE/FID/NLL, lower is better (

\downarrow

). For BARTScore, higher is better (

\uparrow

Figure 4: (a) Distribution of FID on CIFAR-10 data. (b) Distribution of MSE on data.CIFAR-10.

	CIFAR-10			CIFAR-100			CommonGen			WMT
Method	Mem.	GFLOPs	Lat.(ms)	Mem.	GFLOPs	Lat.(ms)	Mem.	GFLOPs	Lat.(ms)	Mem.	GFLOPs	Lat.(ms)
THA	255.40	0.0797	3.494	255.40	0.0797	3.535	406.80	0.1845	3.611	1416.68	0.4698	7.299
DCMHA	292.71	0.0785	4.443	292.71	0.0785	4.499	463.36	0.1780	4.428	1650.30	0.4399	8.148
ColMHA	207.39	0.0765	3.576	207.39	0.0765	3.552	282.79	0.1699	3.529	903.74	0.4068	3.868
MMA	191.39	0.0755	16.677	191.39	0.0755	17.606	281.05	0.1678	25.863	923.74	0.4027	44.356
MoA	321.07	0.1101	4.078	321.07	0.1101	4.098	484.35	0.2538	4.095	1649.67	0.6417	6.770

Table 2: Efficiency comparison with recent multi-head attention variants.

Setting	Dataset	FID@Top-50 / BART@Top-50					FID / BART
Setting	Dataset	1-layer	3-layer	5-layer	7-layer	9-layer	1-layer	3-layer	5-layer	7-layer	9-layer
DAMCHA	CIFAR-10	77.97	75.79	71.07	74.07	58.64	250.35	231.52	234.37	229.09	229.84
	CIFAR-100	65.20	60.82	57.64	52.14	50.29	275.80	240.98	221.85	237.57	217.89
	WMT	-5.26	-5.14	-5.08	-4.98	-4.51	-7.71	-7.27	-6.00	-6.32	-5.61
	CommonGen	-4.97	-4.90	-4.62	-4.84	-4.60	-7.83	-6.53	-7.65	-6.90	-6.20
Diagonal-only	CIFAR-10	102.51	105.05	108.505	98.703	112.856	248.38	233.05	235.17	233.485	239.609

Table 3: Ablation on the capacity of the standard DAMCHA’s

f(x)

at different depth, and inference using full-, diagonal-only, and off-diagonal-only metric matrices. Within each row (fixed setting and dataset), the best layer is highlighted in red bold and the worst layer in blue for each metric block.

Dataset	Variant	Metric	Metric@Top-50
CIFAR-10 (FID $\downarrow$ )	w/o $f(x)$ +share	233.038 $\downarrow$ -2.37%	76.554 $\downarrow$ -17.59%
	w/o share	228.961 $\downarrow$ -0.58%	77.851 $\downarrow$ -19.59%
	DAMCHA	227.641	65.132
CIFAR-100 (FID $\downarrow$ )	w/o $f(x)$ +share	229.467 $\downarrow$ -8.66%	61.950 $\downarrow$ -15.25%
	w/o share	234.296 $\downarrow$ -10.95%	57.934 $\downarrow$ -7.77%
	DAMCHA	211.290	53.756

Table 4: Ablation results. Here,

f(x)

is implemented as an MLP. Blue

\downarrow

indicates the relative degradation w.r.t. DAMCHA (in %).

Dataset	Cross-entropy loss $\downarrow$			PPL $\downarrow$
	MHA	DCMHA	DAMCHA	MHA	DCMHA	DAMCHA
Enwik8	1.07	1.05	1.03 $\uparrow$ 1.90%	2.90	2.86	2.80 $\uparrow$ 2.10%
Text8	1.06	1.05	1.03 $\uparrow$ 1.90%	2.88	2.86	2.81 $\uparrow$ 1.75%
Mix	1.09	1.08	1.04 $\uparrow$ 3.70%	2.98	2.94	2.82 $\uparrow$ 4.08%

Table 5: Token-level pretraining results. Lower is better for cross-entropy loss and PPL. Red

\uparrow

denotes DAMCHA’s relative improvement over DCMHA (in %).

Model	Top-1		Top-5
Model	epoch 100	epoch 500	epoch 100	epoch 500
MHA	69.02%	79.51%	88.68%	94.68%
DCMHA	67.90%	81.70%	88.71%	95.00%
DAMCHA	70.93%	82.10%	89.71%	96.23%

Table 6: Results on ImageNet-1K classification.

Appendix figures & tables11 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 5: Data flow and tensor shapes of DAMCHA (right, with the value/output projection shapes consistent with standard MHA)

Stage	MHA	DAMCHA
Q, K projections	$\mathcal{O}(Ld^{2})$	— (replaced by $f_{\theta}$ )
Metric generator $f_{\theta}$ $\dagger$	—	$\mathcal{O}(Ld+Pw^{2}+wd^{2})$
RowSum aggregation	—	$\mathcal{O}(d^{2})$
Attention scores	$\mathcal{O}(Ld^{2}+L^{2}d)$	$\mathcal{O}(Ld^{2}/h+L^{2}d)$
Softmax	$\mathcal{O}(hL^{2})$	$\mathcal{O}(hL^{2})$
V projection + output	$\mathcal{O}(Ld^{2}+L^{2}d)$	$\mathcal{O}(Ld^{2}+L^{2}d)$

Appendix

Table 7: Per-layer time complexity: DAMCHA vs. standard MHA.

\dagger

indicates the stage is shared across all

N

layers and amortises to

\mathcal{O}(\cdot)/N

Figure 6: Supplementary distribution plot on CIFAR-100.

Figure 7: Supplementary distribution plot on CommonGen.

Figure 8: Cross-head metric matrix visualization on CIFAR-10. The top row shows the state of the cross-head metric matrix at epoch 1, and the bottom row shows the state at epoch 50.

Figure 9: (a) Gradient contribution heatmap of the cross-head metric matrix when the sharing mechanism is disabled. (b) Gradient contribution heatmap of the cross-head metric matrix under the standard DAMCHA setting.

Figure 10: Contour-map visualizations across training epochs (Epoch 0/12/24/36/49). Top: w/o $f(x)$ ; Bottom: DAMCHA.

Figure 11: 3D surface visualization of the learned metric matrix $M$ . Left: w/o $f(x)$ ; Right: DAMCHA.

Figure 12: Cross-subspace regression: DAMCHA vs. MHA. (a) comparison results (testing MSE) under different entanglement levels. (b) DAMCHA’s relative advantage under different entanglement levels.

Figure 13: Cross-subspace classification: DAMCHA vs. MHA. (a) Comparison results (accuracy %) under different $d_{e}$ . (b) Relative improvement of DAMCHA over MHA across different $d_{e}$ .

Experiment	Dataset	Steps / Epochs	Approx. Training Time
Decoder-only Transformer	CIFAR-10	50 epochs	$\sim$ 3 h
	CIFAR-100	50 epochs	$\sim$ 3 h
	WMT14	50 epochs	$\sim$ 5 h
	CommonGen	50 epochs	$\sim$ 4 h
GPT-2 Pretraining	Enwik8	100K steps	$\sim$ 8 h
	Text8	100K steps	$\sim$ 8 h

Appendix

Table 8: Approximate training time per experimental setting. All runs use a single RTX 3090 GPU.

Multi-Head Attention

Multi-head attention uses parallel attention heads to route and combine information in Transformers, shaping their capacity, interpretability, and efficiency.

See all 74 papers

Adaptive Head Budgeting for Efficient Multi-Head Attention
Apr 24, 2026Bilal Faye, Abdoulaye Mbaye, Hanane Azzag +1Transformer Transformer Attention
Multi-Head Attention Residuals
Jul 22, 2026Cheng Luo, Zefan Cai, Junjie HuTransformer Attention Residual Connections
Independent Researcher · University of Wisconsin–Madison
Interleaved Head Attention
Feb 24, 2026Sai Surya Duvvuri, Chanakya Ekbote, Rachit Bansal +6Multi-Head Attention LLM Reasoning
UT Austin · MIT · Harvard University +2

Apr 24, 2026J/K move · Enter open · S save

Adaptive Head Budgeting for Efficient Multi-Head Attention

Bilal Faye, Abdoulaye Mbaye, Hanane Azzag, Mustapha Lebbah

Transformer Transformer Attention

Open paper AskPDF arXiv

Multi-head attention enables Transformers to capture diverse representations, but all attention heads are typically activated for every input, regardless of task complexity. For coarse-grained tasks such as text classification, where relevant information is often global, this fixed allocation can introduce unnecessary computation. We propose BudgetFormer, a Transformer architecture that dynamically allocates attention heads on a per-input basis. The model learns both a head budget and a relevance distribution to select the most informative heads. To support effective head selection, we introduce a training strategy that balances exploration and exploitation. Experiments on text classification tasks show that BudgetFormer reduces FLOPs and memory usage while matching or surpassing the performance of standard multi-head attention. These results highlight adaptive head allocation as an effective approach to improving Transformer efficiency and performance.

Abstract

Multi-Head Attention

Explore similar work

Adaptive Head Budgeting for Efficient Multi-Head Attention

Multi-Head Attention Residuals

Interleaved Head Attention