论文

为何要事事全神贯注?专注是关键

Why Attend to Everything? Focus is the Key

模型架构Transformer

摘要

标准注意力与序列长度呈二次方关系。有效的注意力方法降低了 O(n^2) 成本,但当改造为预训练模型时,它们通常会降低困惑度和/或下游准确性。我们引入 Focus,这是一种了解哪些词元对重要的方法。 Focus 添加了一小组可学习的质心(每层只有 148K 个参数),充当门:只有属于同一质心组的词元对才能在长范围内相互关注。焦点是可组合的:可以通过仅训练质心同时保持所有原始权重冻结来将其添加到任何预训练模型中。实验表明,将 Focus 组合到预训练模型上,可以在 124M 到 70B 参数的模型大小和五种注意力架构的下游基准上实现零退化。令人惊讶的是,稀疏焦点注意力在 124M 尺度上优于完全注意力(困惑度为 30.3 vs. 31.4),并且在 7B 尺度上从头开始训练时与完​​全注意力相匹配(13.82 vs. 13.89)。焦点也很快:top-k 组成员资格可提供 2 倍的加速,并且质量比原始预训练模型更好。使用我们的 FlashAttention 分解,Focus 在没有自定义内核的情况下以 1M 词元实现了 8.6 倍的加速。