LongNet

收录理由

处理超长输入时，标准Transformer的注意力计算量会随序列长度呈平方级增长，训练往往因此卡住。LongNet把论文中提出的膨胀注意力机制用PyTorch实现出来，让注意力覆盖范围随token距离增大而指数扩展，计算开销却大致保持线性，能应对远超常规自注意力承受能力的序列长度。项目提供了一个独立的DilatedAttention模块，可以直接嵌入现有模型，也附带一个由前馈、层归一化和SwiGLU块搭好的LongNetTransformer，以及针对enwiki8的训练脚本。它更像研究用的积木，而不是开箱即用的成品，想用的话需要按自己的架构调整。如果你在探索极长上下文或文档级序列的实验，这份代码是值得研读和改写的参考。