基于核引导的互信息学习有向无环图的 Transformer 证明

探寻隐藏在现实数据背后的图结构是横跨多个科学领域的关键挑战。最近，基于注意力机制的 Transformer 模型在捕捉图中复杂依赖关系方面表现出强大的经验成功。然而，关于其训练动力学的理论理解仅限于树状图结构，即每个节点仅依赖于单个父节点。将可证明的保证推广到更一般的有向无环图 (DAG)——即涉及多个父节点 per node——仍然具有挑战性，主要因为难以设计能够使不同注意力头分别学习多个不同父子关系的训练目标。本文通过引入一种新型信息论度量——基于 $f$ 范数（divergence）的核引导互信息 (KG-MI)，来解决这一问题。我们的目标函数将 KG-MI 与多注意力头框架结合，其中每个注意力头与一个独立的边际转换核相关，以有效建模多样化的父子依赖关系。我们证明，给定由 $K$ 个父节点构成的 DAG 生成的序列，通过梯度上升训练单个单层多注意力头的 Transformer 可在多项式时间内收敛到全局最优。此外，我们刻画了收敛时的注意力得分模式。当将 $f$ 范数特化为 KL 散度时，所学到的注意力得分准确反映真实的邻接矩阵，从而可证明地恢复底层图结构。实验结果验证了我们的理论发现。