用偏序图理论分析注意力掩码,指导高效训练结构设计
Hasse Diagrams for Attention: A Partial Order Framework for Designing Transformer Masks
- 将注意力掩码的信息流建模为偏序图(Hasse图)
- 证明深层Transformer信息流收敛于哈斯图结构
- 可直接从任务需求生成新掩码,适合模型架构设计者
大型Transformer模型训练中,注意力掩码控制序列间信息流动的范围与方向。尽管存在多种掩码变体,且如FlexAttention等工具已支持任意掩码,但对任意掩码所诱导的信息流结构缺乏系统性理论分析。本文建立完整理论框架,证明在足够深度下,多层Transformer的信息流收敛于哈斯图——一种表示偏序关系的有向无环图。基于此,将并行训练任务的设计重构为寻找哈斯图的最小共同超图问题,并确立最小共同超图的判定准则。由此提出构造性方法,可直接从一组任务中推导出注意力掩码。应用该框架,设计出两种新型掩码:确保训练-推理一致性的块生成注意力掩码(Block Two-Stream Attention),以及全监督双向注意力掩码(Butterfly Attention)。结果验证了该框架发现新结构的能力。
原文摘要 · Abstract (English)
During the training of large Transformer models, attention masks regulate the scope and direction of information flow across a sequence. Numerous mask variants exist, and operators such as FlexAttention already support arbitrary attention masks. Nevertheless, a systematic formal analysis of the information-flow structure induced by arbitrary masks has been missing. This paper develops a complete theoretical framework. We prove that, with sufficient depth, the information flow of a multi-layer Transformer converges to a Hasse diagram -- a directed acyclic graph representing a partial order. Building on this, we recast the design of parallel training tasks as the problem of finding a minimal common supergraph of Hasse diagrams, and we establish a criterion for the minimal common supergraph. This yields a constructive method to derive attention masks directly from a family of tasks. Applying the framework, we design two novel masks: a block-generation attention mask that ensures training-inference consistency (Block Two-Stream Attention), and a fully supervised bidirectional attention mask (Butterfly Attention). These results demonstrate the framework's capacity to discover new structures.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。