rikkathree/vsa-lm

VSA-LM A 100M parameter language model with a fixed-size memory instead of the usual growing attention cache.

I built this to answer one question: can a small model with linear-attention memory layers actually keep up with a regular transformer if you train them identically? Turns out, at this scale, yes — it matches within noise, while using a lot less memory to run. And it can do something transformers can't: you can write a document into its memory and erase it again, exactly.

Details

• 8 layers: 6 VSA (fixed-state, per-head learnable decay) + 2 softmax
• Custom Triton kernels for training, with an adjoint gradient for the decay rates
• Every kernel is checked against a pure-PyTorch reference before training is allowed to start (see vsa_lm/tests.py)
• Trained on a single free Kaggle T4 in about 20 GPU-hours, fp16 + 8-bit AdamW

Limitations

• 500M tokens is roughly 5% of what a model this size would ideally read. Both this and the baseline are data-limited, so this comparison says nothing about what happens at 10B+ tokens.
• Positions only go to 512 (that's the training window).
• Not instruction-tuned, not chat. It completes text, reasonably grammatically, with confident nonsense for facts.
• The memory holds hundreds of tokens cleanly, not whole libraries.

Code and training scripts: github.com/Rikka3/vsa-lm