memorizing-transformers-pytorch

Why we included this project

Researchers and ML engineers experimenting with memory-augmented attention will find a clean, self-contained PyTorch implementation of the Memorizing Transformers idea from ICLR 2022. Instead of relying only on the current context window, the model indexes past token representations in an approximate-nearest-neighbor store and pulls relevant memories back into attention, a direction worth exploring for long-context and continual-learning work. The code follows the paper closely but with two deliberate changes: hybrid attention over local and distant logits rather than the sigmoid gate, and cosine-similarity attention with a learned temperature. It wraps faiss for the ANN index, shows how to pair the approach with Transformer-XL style memory, and ships an enwik8 training script. If you are reproducing or extending retrieval-augmented transformer research rather than shipping a production system, this is a handy reference to study and adapt.