titans-pytorch: an unofficial Titans implementation with a swappable neural memory module
Unofficial implementation of Titans, SOTA memory for transformers, in Pytorch
At a glance
- What is it?
- The package ships two usable surfaces, a standalone NeuralMemory module and a MemoryAsContextTransformer, plus a training script. It is research code with a thin API, and the README is honest about how little is settled.
- Who is it for?
- Adopt titans-pytorch if you want a working PyTorch entry point to the Titans memory idea and you are prepared to read the source, because the README documents two constructors and a training script and little else.
- Can I use it commercially?
- Yes. MIT is a permissive licence: you can use, modify and sell software built on it, as long as you keep its copyright and licence notices.
- Is it still maintained?
- Yes. The repository last received commits 64 days ago.
- What is it written in?
- Mainly Python, according to GitHub's language statistics.
Answers come from the project's GitHub data, last synced on September 15, 2026, and from our analysis. They are not legal advice.
DEEP OPEN-SOURCE ANALYSIS
What problem titans-pytorch is trying to solve, and for whom
The repository is an unofficial PyTorch implementation of Titans, the paper Learning to Memorize at Test Time by Behrouz, Zhong and Mirrokni. The README describes it as SOTA memory for transformers, and the topic tags on the repository include long-term-memory and test-time-training. So the target is a specific research problem: giving a sequence model a memory that keeps being written to during inference, rather than a fixed key-value cache that only grows.
The audience is narrow. This is not a library for someone who wants better long-context performance by changing an import. It is for people who have read the paper, or watched the linked paper review by Yannic, and want a PyTorch artifact they can run and modify. The README gives a Colab link for a quick run, which is the intended first contact. The Appreciation section credits Eryk for early experimental results, described as positive for a 2 layer MLP, which tells you the author is still collecting evidence about which memory module shape works.
The honest framing is that this package exists so that the architecture can be tried, not so that it can be deployed. Anyone evaluating it as infrastructure should read the second sentence of the README first: it will also contain explorations into architectures beyond the simple 1-4 layer MLP for the neural memory module, if it works well to any degree.
NeuralMemory as a standalone module: the actual interface
The smallest useful unit is NeuralMemory. The README example constructs it with dim=384 and chunk_size=64, moves it to CUDA, feeds a tensor of shape (2, 1024, 384), and asserts that the returned tensor has the same shape as the input. Two values come back: retrieved and mem_state. The first is the retrieved memory, the second is state you can carry forward.
That shape-preserving contract is the whole public surface of the module as documented. There is no described method for resetting state, merging states across batches, or serialising mem_state to disk, and the README does not say what dtype or structure mem_state has. If you need to checkpoint memory across sessions, you are reading the source.
The one tuning note the README gives is about chunk_size: set to smaller chunk size for better perf on smaller sequence lengths (but more memory usage). That is a direct trade-off statement with no numbers attached. It also means chunk_size is not a neutral parameter. It changes both speed and memory footprint, and the direction of the speed effect depends on sequence length. Anyone benchmarking this should fix chunk_size deliberately and report it, because a comparison across different chunk sizes is not a comparison of the same model.
MemoryAsContextTransformer and the MAC configuration
The second surface is the transformer. The README shows MemoryAsContextTransformer with num_tokens=256, dim=256, depth=2, segment_len=128, num_persist_mem_tokens=4 and num_longterm_mem_tokens=16. The comment on segment_len calls it the local attention window size. So the architecture splits work: local attention covers a window of segment_len tokens, and the memory tokens carry information past that window. The two token counts, persist and longterm, are separate knobs, which suggests the design distinguishes memory that is always present from memory that is learned over a longer horizon. The README does not explain that distinction, and the naming alone is not enough to be sure.
The training call is a single line: loss = transformer(token_ids, return_loss=True), with token_ids drawn as random integers of shape (1, 1023) from a vocabulary of 256. The comment annotates the loss shape as (1, 1023, 256), which is a per-position, per-vocabulary tensor rather than a scalar. That matters for anyone writing a training loop, because loss.backward() on a tensor of that shape is not the same as calling backward on a reduced scalar, and the README does not show a reduction step. Sampling then uses transformer.sample(token_ids[:, :4], 512), taking a four-token prefix and generating 512 tokens.
The architectural claim worth flagging is the one in the README's opening: explorations into architectures beyond the simple 1-4 layer MLP for the neural memory module. That means the memory module is treated as a replaceable component, and the package is a testbed for that substitution. If your interest is the memory module design rather than the surrounding transformer, this is the part of the repository aimed at you.
Getting it running: pip, uv, and the training script
Installation is one command: pip install titans-pytorch. The package name on PyPI is titans-pytorch while the import name is titans_pytorch, which is standard but worth noting when writing setup files.
For experiments the README prescribes a different path. Install uv with pip install uv, then modify train_mac.py and run it with uv run train_mac.py. The phrasing in the README is that you modify the script and run it to query nature. There is no configuration file, no CLI flags, no argument parser documented. Parameters are edited in the script itself. That is a real constraint on reproducibility: if you want to sweep segment_len or num_longterm_mem_tokens, you are editing Python between runs or wrapping the script yourself, and there is no documented way to override values from the command line.
The example code is CUDA-first. NeuralMemory is constructed and then .cuda() is called, and the input sequence is created on CUDA. The transformer example does not call .cuda() at all, so it runs on CPU as written. Neither example pins a torch version, and the README does not state a minimum Python version. Check those against your environment before assuming the snippets run unchanged.
Where the documentation stops and the risk begins
The largest limitation is not architectural, it is evidentiary. The README contains no benchmark table, no comparison against a baseline transformer, and no reported perplexity or accuracy. The only performance statement is the chunk_size note about better perf on smaller sequence lengths with more memory usage, and it is qualitative. The Appreciation section mentions early experimental results that were positive for a 2 layer MLP, which is a single informal data point from one collaborator, not a result you can plan around.
There is also no documented evaluation harness. The repository has a train_mac.py script, but the README describes no eval script, no dataset, and no metric. If you need to know whether the memory actually helps on your task, you are building the measurement yourself, and you should budget for that as part of adoption rather than as an afterthought.
The second limitation is API stability. The package is at 0.5.x with three releases in the space of roughly a month, 0.5.0 in January 2026, 0.5.1 later that month, and 0.5.3 in February. The README does not include a changelog or migration notes, so a version bump can change constructor arguments or return values without a documented path. Pin the version in your requirements and read the diff when you move.
Third, the README itself says the implementation is unofficial. That is not a legal statement, but it does mean there is no upstream maintainer commitment to match the paper's behaviour exactly, and no stated correspondence between a given release and a given section of the paper.
What you would use instead, and how the approach differs
The obvious alternative is to use a long-context transformer with a standard attention mechanism and rely on the existing ecosystem around it. The difference is where the capacity lives. A standard long-context setup extends the attention window or adds positional techniques so that attention itself reaches further, and the model's state at inference is the key-value cache, which grows with context and is not updated by gradient. Titans, as implemented here, adds a memory module that is trained at test time, and the README's own topic list names test-time-training as a defining feature. The mem_state returned by NeuralMemory is a learned artifact, not a cache of past keys and values.
A second reference point sits inside the repository's own citation list. The README cites Learning to (Learn at Test Time): RNNs with Expressive Hidden States by Sun and others, and Test-Time Training Done Right by Zhang and others. Those are the neighbouring lines of work the author considers relevant. If your goal is test-time training in general rather than Titans specifically, reading those alongside this implementation is more useful than treating this package as the only option, because the README treats them as related rather than as competitors.
The practical difference for a team is maintenance. A long-context attention stack has multiple mature implementations and a large body of reported results. titans-pytorch has one author, an MIT licence, and a README that points at a Colab notebook for a quick run. Choosing it means accepting that you are early.
Licence, maintenance and what upgrading actually costs
The licence is MIT. That is permissive and permits commercial use, modification and redistribution provided the copyright notice and permission notice are retained. This is a description of the licence text, not legal advice; if the memory module ends up inside a product, have counsel confirm the notice requirements and check whether the cited papers impose anything the MIT licence does not cover.
The maintenance picture from the supplied material: the repository is not archived, the default branch is main, the last push is dated 2026-07-13, and releases 0.5.0, 0.5.1 and 0.5.3 landed in January and February 2026. So the project is active. Activity is not the same as stability, and the release cadence within a single minor version line suggests the API is still moving.
The upgrade cost is concrete. Because there is no changelog in the README and no documented deprecation policy, moving from 0.5.1 to 0.5.3 means diffing the source for changes to NeuralMemory's constructor arguments, to the tuple returned by its forward pass, and to the MemoryAsContextTransformer keyword arguments you rely on. The parameters most likely to move are the ones the README flags as experimental, namely the memory module architecture beyond the 1-4 layer MLP. Pin titans-pytorch to an exact version in your dependency file, and treat a version bump as a code change rather than a routine update.
Editorial conclusion
Adopt titans-pytorch if you want a working PyTorch entry point to the Titans memory idea and you are prepared to read the source, because the README documents two constructors and a training script and little else. Do not adopt it if you need a supported library with stable APIs, published throughput numbers or a documented evaluation: the repository is an unofficial implementation, the author states it will also contain explorations into architectures beyond the simple 1-4 layer MLP, and it ships no benchmark table to check against. Before you build on it, verify three things in your own environment: that retrieved.shape matches your input shape for your chosen dim and chunk_size, that the loss from MemoryAsContextTransformer with return_loss=True backpropagates through the memory state, and what peak memory your segment_len and num_longterm_mem_tokens values actually cost on your hardware.
Community notes