arXivDaily arXiv每日学术速递 周一至周五更新
arXiv 2402.04779cs.CLcs.AI
收藏

StableMask:改进解码器专用Transformer中的因果掩码

StableMask: Refining Causal Masking in Decoder-only Transformer

Zhejiang University(浙江大学);Peking University(北京大学);Tencent AI Lab(腾讯AI实验室);Eastern Institute of Technology(东方理工大学)

Qingyu Yin, Xuzheng He, Xiang Zhuang, Yu Zhao, Jianhua Yao, Xiaoyu Shen, Qiang Zhang

AI总结:

针对解码器专用Transformer中因果掩码的局限,提出无参数方法StableMask,通过伪注意力值和递减掩码比率平衡注意力并编码绝对位置,提升71M至1.4B参数模型性能,支持高效外推与集成。

摘要:

采用因果掩码和相对位置编码(RPE)的解码器专用Transformer架构已成为语言建模的事实标准。尽管其在各种任务中表现出色,我们识别出两个局限性:第一,它要求所有注意力分数非零且总和为1,即使当前嵌入已包含足够自足信息。这迫使模型对特定标记分配不成比例的过多注意力。第二,基于RPE的Transformer并非通用逼近器,因为其在编码绝对位置信息方面能力有限,这限制了它们在位置关键任务中的应用。在本工作中,我们提出StableMask:一种无参数方法,通过改进因果掩码来解决这两个局限性。它引入伪注意力值以平衡注意力分布,并通过逐渐减小的掩码比率编码绝对位置信息。StableMask的有效性在理论上和实证上均得到验证,在参数规模从71M到1.4B的语言模型中,跨多种数据集和编码方法显示出显著提升。我们进一步表明,它自然支持(1)无需特殊技巧(如StreamingLLM)的高效外推,以及(2)与现有注意力优化技术的轻松集成。

英文摘要:

The decoder-only Transformer architecture with causal masking and relative position encoding (RPE) has become the de facto choice in language modeling. Despite its exceptional performance across various tasks, we have identified two limitations: First, it requires all attention scores to be non-zero and sum up to 1, even if the current embedding has sufficient self-contained information. This compels the model to assign disproportional excessive attention to specific tokens. Second, RPE-based Transformers are not universal approximators due to their limited capacity at encoding absolute positional information, which limits their application in position-critical tasks. In this work, we propose StableMask: a parameter-free method to address both limitations by refining the causal mask. It introduces pseudo-attention values to balance attention distributions and encodes absolute positional information via a progressively decreasing mask ratio. StableMask's effectiveness is validated both theoretically and empirically, showing significant enhancements in language models with parameter sizes ranging from 71M to 1.4B across diverse datasets and encoding methods. We further show that it naturally supports (1) efficient extrapolation without special tricks such as StreamingLLM and (2) easy integration with existing attention optimization techniques.

机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。

↑