用马氏距离重构注意力机制,提升长序列状态追踪性能。
Mahalanobis-Based Multi-Head Attention for Complex State Propagation

- 以马氏距离替代点积,实现无参数增长的高维特征注意力计算。
- 仅119K参数下,最终隐藏状态使用教师强制时优于Transformer和GCN基线。
- 支持跨头协同与信息直通,适合符号结构建模与高效推理任务。
本文提出马氏距离驱动的多头注意力(MHA-CSP),将标准点积替换为基于马氏距离的RBF核,在不增加参数量的前提下,有效实现无限维特征空间中的注意力计算。马氏距离的正定性使我们可直接构建树状注意力:注意力分数由累积距离构造,并通过LogSumExp校正,消除边指数和的偏差。此外,多头马氏距离矩阵被重用于构建注意力拼接机制,实现跨头核协作,同时提升精度与训练效率。大量实验表明,仅含119K参数且仅在最终隐藏状态应用教师强制的MHA-CSP,在长序列状态追踪任务中持续优于从零训练的Transformer与GCN基线。相较于依赖密集注意力或图传播的基线,MHA-CSP通过合成距离校正实现鲁棒的结构化推理,其来自CSP主干的高效信息旁路特性,彰显了复杂值状态传播与协同多头校正在捕捉符号结构方面的有效性,为结构化推理建立了新的效率-性能权衡范式。
原文摘要 · Abstract (English)
In this paper, we propose \textbf{Mahalanobis-Based Multi-Head Attention} (MHA-CSP), a novel attention mechanism that replaces the standard dot-product with a \textbf{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 \textbf{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 \textbf{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 \textbf{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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。