memorizing-transformers-pytorch

收录理由

做记忆增强注意力相关实验的研究者或机器学习工程师，会在这里找到一个干净、自包含的 PyTorch 实现，对应 ICLR 2022 的 Memorizing Transformers 思路。它不局限于当前上下文窗口，而是把过去的 token 表示存入近似最近邻索引，再把相关记忆拉回注意力计算中，这个方向对长上下文和持续学习都值得琢磨。代码基本贴合原论文，但有两处有意改动：用局部和远距离 logits 的混合注意力替代 sigmoid 门控，以及用带可学习温度的余弦相似度注意力。项目用 faiss 做 ANN 索引，也展示了如何与 Transformer-XL 风格记忆结合，还附带了 enwik8 训练脚本。如果你是在复现或扩展检索增强 transformer 的研究，而不是要直接上生产，这份代码很适合拿来研读和改造。