Research map / Architectures and efficiency / Attention and transformers

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

NeurIPS 2026 · Architectures and efficiency › Attention and transformers · approach: Attention head and block redesign

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes DAMCHA, an attention mechanism that replaces MHA's query-key similarity with a learned, input-dependent full cross-head Mahalanobis-like metric matrix shared across layers, reporting gains on generation, GPT-2 pretraining and ViT ImageNet.

Problem

Standard multi-head attention uses block-diagonal, input-independent, low-rank similarity metrics, ignoring cross-head interactions and input-adaptive geometry.

Approach

Learn the full cross-head metric matrix M(X) generated by a shared MLP metric generator from the input; each head's score uses the row-sum of its block row, X m_i(X) X^T, with stack-wise sharing of the generator across layers.

Key results

• On CIFAR-10, DAMCHA reduces FID@Top-50 from 91.43 (MHA) to 63.03, a 31.1% relative reduction.
• On Mix token-level pretraining with GPT-2, DAMCHA reaches 1.04 cross-entropy versus 1.08 for DCMHA (3.70% loss reduction).
• On ImageNet-1K ViT-S/16 at epoch 500, DAMCHA achieves 82.10% Top-1 versus 79.51% for MHA and 81.70% for DCMHA.

Surprising findings

• MMA achieves the lowest NLL on WMT yet much worse BARTScore, so lower token likelihood does not imply better generation quality.
• Diagonal-only inference of the learned metric is much worse than full DAMCHA, while off-diagonal-only inference is worse still.

Limitations

• Gains on several benchmarks are small and the paper reports single numbers with limited variance analysis.
• The method is not positive definite by design and the RowSum-based score is not a true Mahalanobis metric, as the paper acknowledges.
• Theoretical results rest on strong assumptions (unique minimizers, quadratic risk growth, identifiability) that are not verified empirically.

Experimental setup

Datasets: CIFAR-10 CIFAR-100 WMT14 CommonGen Enwik8 Text8 ImageNet-1K

Benchmarks: CIFAR-10 CIFAR-100 WMT CommonGen Enwik8 Text8 Mix ImageNet-1K

Metrics: MSE FID FID@Top-50 NLL BARTScore BARTScore@Top-50 cross-entropy loss PPL Top-1 accuracy Top-5 accuracy

Baselines: MHA THA DCMHA ColMHA MMA MoA

Builds on: multi-head attention DCMHA THA MMA

Closest papers by meaning

• No Head Left Behind – Multi-Head Alignment Distillation for Transformers (AAAI 2024 · similarity 0.86)
• SAS: Simulated Attention Score (NeurIPS 2025 · similarity 0.86)
• Improving Transformers with Dynamically Composable Multi-Head Attention (ICML 2024 · similarity 0.86)
• MoH: Multi-Head Attention as Mixture-of-Head Attention (ICML 2025 · similarity 0.86)
• Memorization Capacity of Multi-Head Attention in Transformers (ICLR 2024 · similarity 0.86)
• Mixture-of-Top-$k$ Attention: Efficient Attention as Scalable Fast Weights (NeurIPS 2026 · similarity 0.85)