Train a 340M parameter language model using a new content-based attention mechanism.
Reproduce published recall benchmarks (MQAR and overwrite tasks) comparing attention architectures.
Swap the KATA kernel into a Hugging Face model to experiment with long-context recall.
| ayghri/kata | 0xallam/my-recipe | 0xhassaan/nn-from-scratch | |
|---|---|---|---|
| Stars | 0 | — | 0 |
| Language | Python | Python | Python |
| Last pushed | — | 2022-11-22 | — |
| Maintenance | — | Dormant | — |
| Setup difficulty | hard | moderate | moderate |
| Complexity | 5/5 | 2/5 | 4/5 |
| Audience | researcher | general | developer |
Figures from each repo's GitHub metadata at analysis time.
Needs a CUDA GPU, Hugging Face and Weights and Biases accounts, and a multi-step training pipeline.
KATA is a research project that builds a new kind of attention layer for language models, the part of a model that decides which earlier words matter when predicting the next one. Instead of the usual approach, KATA scores how related two tokens are by squaring their combined similarity, which lets it recall information based on content rather than on position in the text. Because it does not rely on position, a model trained on short sequences can still work reasonably well on much longer ones. A variant called DeltaKATA adds the ability to overwrite or erase earlier stored information when new, related information arrives, something that simpler gated models cannot do as directly. The repository contains three main pieces. First, production-ready GPU kernels written in Triton that make the attention computation fast. Second, a model and training setup that plugs into the popular Hugging Face library, so it can be trained as a roughly 340 million parameter language model using standard tools like torchrun and Hydra configuration files. Third, a set of synthetic benchmarks that test how well a model can remember and update key value pairs over long sequences, used to reproduce the tables and figures from the associated research paper. Setting it up requires a Python environment built with the uv tool, a modern GPU (Nvidia Ampere or newer), and accounts with Hugging Face and Weights and Biases for downloading data and logging training runs. This is not a beginner-friendly, install-and-go tool. It is aimed at machine learning researchers who want to train language models from scratch, experiment with new attention designs, or compare KATA against existing methods like Transformers, DeltaNet, and Mamba2 on standardized recall tests. There is no mention of a lighter-weight way to try the ideas without a full training pipeline and GPU hardware.
A research attention layer for language models that recalls information by content instead of position, enabling longer-context reasoning.
Mainly Python. The stack also includes Python, PyTorch, Triton.
Setup difficulty is rated hard, with roughly 1day+ to a first successful run.
Mainly researcher.
This repo across BitVibe Labs
Don't trust strangers blindly. Verify against the repo.