Titans: Learning to Memorize at Test Time

We thank the reviewer for their time and valuable comments. We respond to your comments below:

Despite promising formal proofs in the ICML rebuttal, the revised paper lacks the theoretical analysis of memory capacity and convergence properties.

Response: We first want to kindly bring to your consideration that providing theoretical results is not a must for publications. While we understand the reviewer's interests and curiosity about the theoretical side of memory modules, we want to clarify that none of the peer studies provide such theoretical results about the capacity of the memory or why one-shot memorization works (e.g., GLA (ICML 2023), Mamba2 (ICML 2024), Mamba (COLM 2024, best paper award), Gated DeltaNet (ICLR 2025), GSA (NeurIPS 2024), Simba (NeurIPS 2024), Longhorn (NeurIPS 2024), TTT (ICML 2025, spotlight), etc.). These are indeed interesting questions that one can ask and explore in future studies. Our submission is a 9-page research paper that aims to provide enough evidence for its usefulness and effect. Even the initial Transformers paper, after 8 years, is still being studied, improved, and evaluated.

In ICML you mentioned you are willing to increase your initial rating of “weak accept”, if we provide theoretical results. Despite providing theoretical results and addressing your concerns, your initial rating remained the same. Therefore, it was not clear for us if our theoretical results are what you expected. In this submission, we decided to make our work more focused, and so leave such exploration for future work, mainly due to the fact that the goal of this work is not on one-shot memory update nor understanding the capacity of recurrent models.

Some important model comparisons are incomplete. For example, the throughput comparison with Samba is still missing from figures despite specific reviewer requests in ICML.

Response: Please note that as we discussed earlier in the ICML review process, Samba is Attention + Mamba. We already have a throughput comparison with Mamba, which indicates that our Titans (LMM) is faster. Therefore, we can conclude that Samba = Attention + Mamba is slower than our MAL = Attention + LMM. However, following your suggestion, here we provide the Samba throughput and compare it with our models:

We will add these results in the camera ready version of the paper.

BABILong evaluation remains flawed, comparing fine-tuned small models against zero-shot large models without acknowledging this unfair advantage.

Response: First, in lines 1374 - 1383 we clearly acknowledged and explained that small models are fine-tuned and large (commercial) models are zero-shot. Therefore, saying ``without acknowledging’’ is a misunderstanding. Second, we have followed the benchmark setups (please see Figure 1 in Peer-reviewed BABILong original paper in NeurIPS 2024). Therefore, we believe if there is a setup in this peer-reviewed paper that the reviewer found flawed, we should not be penalized for that. Please note that it is not expected nor we have the budget to re-run all the benchmark results for all the baselines and change the original setup of the benchmark. If the reviewer believes that it can solve their concern, we can remove the results of large (commercial) models, which definitely are not expected to compare with.

Moreover, in Table 10, we have provided the results for even 128K sequence length in RULER, which we believe should be enough to show the good performance of Titans in longer contexts.

"Momentary surprise" is presented as novel when it's identical to the well-established delta rule. I find this misleading as it stands.

Response: Could you please specify the part that we claim momentary surprise as a novel contribution of our work so we can fix it?

In line 1168, we have clearly mentioned that delta-rule is equivalent to momentary surprise and this is the difference of our update rule, because we consider past surprise as well. Again in line 1328 we clearly mentioned that The Delta Rule is based on momentary surprise.

Different methods are evaluated on different benchmarks without clear justification. Why are some baselines missing from certain benchmarks (e.g. only RULER and not BABILong)? This makes it difficult to form a coherent picture of relative performance.

Response: Please note that the only benchmark with different baselines is the BABILong benchmark. The state-of-the-art recurrent models are considered as the baseline for all language modeling, common-sense reasoning, and RULER tasks. The BABILong benchmark, however, has its own baselines and it has not been considered in any previous similar studies so we can borrow their results. Therefore, it is hard to adapt such models for this downstream task, is out of the scope of our paper, and further requires more compute resources. Once again, please note that we have already shown the good performance of our model until 128K (which has not been evaluated in any similar previous studies on recurrent and linear models). Therefore, if you believe removing BABILong benchmark results can make the evaluation more clear, we would be happy to do it in the camera ready version of the paper.

Despite claims about "larger context windows," rigorous evaluation (RULER) only extends to 16K tokens.

Response: This seems to be a misunderstanding. In Table 10, we have already provided the results for the RULER benchmark until 128K sequence length.

Please note that this is a research paper and more than these sequence lengths are often a goal for commercial LLMs, or research papers that are built on large pre-trained models. Furthermore, our goal is to close the gap between the effective context length of recurrent models and they already perform near zero accuracy in 16K context length of RULER.

"Past surprise" is standard momentum; the main contribution appears to be combining existing techniques

Response: We respectfully disagree with the reviewer that this connection affect the novelty of our method. Momentum term has never been used in this context. With the same logic, one can argue that most of recent architecture studies, including the work of DeltaNet (Schlag et al., 2021) and linear attention (Katharopoulos et al., 2020), are not novel because all use gradient descent, which is a well-known algorithm.

Additionally, the analysis of 2-3 updates improving performance by 0.5% that you promised to include is missing from the current version - could you add it?

Response: Yes, we will make sure to add these results in the final version of the paper.

Fair comparative evaluation: Can you provide (a) BABILong results where all models are either fine-tuned or zero-shot

Response: Please note that all small-models are fine-tuned. If the reviewer believes that the original setup of the peer-reviewed benchmark is not proper, we would be happy to remove the results of zero-shot models. Please once again note that all recent studies in this line has focused on RULER, and we have already provided its results until 128K.

Why does memory-constrained Titans outperform full Transformers on retrieval tasks?

Response: Please note that, the unbounded memory in Transformers is not always the best for all the tasks. The reason is data can be compressible and so the task does not require direct retrieval. Also, it is notable that unbounded memory in Transformers does not mean that they can utilize it in all tasks. As an example, several papers discuss limitations like ''attention sink’’, which can cause Transformers to not scale to a large context. Also, in Table 8, we have provided the results for retrieval tasks that Transformers achieve better results. Please note that this is one the studies in this direction and requires several follow up works to fully address all the aspects.

Moreover, all Transformer and retrieval-based models consider the dependencies of tokens as linear mapping between keys and values. Even given direct access to all past tokens (having perfect memory) the KV cache directly maps keys to values. This might not always the perfect assumption. On the other hand, Titans uses a deep memory and so can learn non-linear mapping between keys and values.

There is no clear guidance on hyperparameter selection, memory size choices, or when to use which variant of Titans.

Response: Please note that memory size is tied to the number of parameters. Given $d$ as the dimension of keys and values, the memory has size $L \times d^2$, where $L$ is the depth of the memory. Therefore, it is not a new hyperparameter, and can be determined based on the computational resources that are available. Also, in our ablation study, we show that for long-context tasks MAC performs better, while for language modeling MAG achieves the best results. We will emphasize more on these in the final version of the paper. Please let us know if there are specific parts that are not clear.

Details "highly optimised kernel" mentioned for MAL?

Response: Please note that as we mentioned in the paper, this is Flash attention kernel and it is not our contribution. The implementation of attention is based on the highly-optimized kernel of Flash attention.

Limitations:

Response: Following your suggestion, we would be happy to discuss them as the limitations and motivated future work to address them.

Details of Implementation:

Response: Following your suggestion, we adhere to provide all the details upon acceptance of the paper.

We hope that our above responses have addressed your concerns, and we would be more than happy to answer any remaining concern.

Dear reviewer,

Re-reviewing is not against policy, provided that the new review reflects the merits of the current version of the paper and does not reveal any names or affiliations. Please adhere to these principles during the discussion and when finalizing your review. Thanks.

Once again, we want to thank the reviewer for their time and comments. Due to the space constraint, we use this comment to address your questions:

1. As we discussed above, the theoretical analysis and understanding of the memory capacity is out of the scope of our paper. We understand that it is an interesting future work and will be discuss it in the limitation and future work section. Regarding the performance multiple update of the memory, while we want to bring to your attention that none of the similar studies mentioned above have done that and again it is out of the scope of our paper (because we do not have any claim about multiple update), we would be happy to provide its results in the final version of the paper. We will add this in the final version, following your suggestion.

2. As we discussed above, all small size models are fine-tuned and only commercial models are zero-shot, which definitely we are not expected to compare with them. However, to fully address your concern, here we provide a comparative evaluation of Titans and baselines we used in other experiments in a fine-tune setup:

Titans show a clear advantage over the baselines.

Also, we will use the 1 additional page in the camera ready version to bring some results from the appendix to the main part, following your suggestion.

3. Could you please clarify, which surprise-based model you are referring to? If you mean EM-LLM, please note that we have already compared Titans with EM-LLM in Table 10, until 128K context length and Titans show a clear improvement +3.5% on average (larger context length show better improvement.). Please note that the hardness of each long-context dataset is different.

Regarding the comparison with Transformers, we have clarified that unconstrained memory does not necessarily mean that the model is able to use all that memory effectively. Moreover, some tasks/dataset might not need unconstrained memory and a more powerful memory system could be more useful in such cases. Furthermore, we already have provided additional results in the initial submission on in-context retrieval tasks that show the superior performance of Transformers.

4. Following your suggestion, we adhere to provide all the details upon acceptance of the paper.

---

Finally please note that as also acknowledged by you, we have improved the paper from different aspects, modified claims, add additional related work discussions (discussing more than 130 studies!), and have provided additional experiments (while such experiments have not been done by any other study on recurrent models). It is indeed discouraging for us to see you keep your initial rating at ICML, before all these changes. Please note that there is no study without limitations and there is always room for further evaluation and discussion. However, in a 9-page conference paper, it is expected to provide enough evidence for the usefulness of the approach, which we believe, as acknowledged by you and other reviewers, we have provided.

We hope that our above responses have addressed your concerns, and we would be more than happy to answer any remaining concern.

I'd like to thank the authors for the detailed responses and the additional results provided in your second comment. I appreciate the constructive engagement, particularly the Samba throughput and additional fine-tuned comparisons, and it is encouraging that you agree to move key results from the appendix to the main paper. However, some core concerns remain:

• First, I note that you've chosen to exclude the theoretical analysis citing scope constraints. However, this analysis would significantly strengthen the paper's contributions. The argument that "no other papers provide such theoretical results" is unconvincing - those papers aren't claiming test-time gradient-based memory updates as their core contribution. Since your approach fundamentally relies on gradient descent dynamics for memory formation, some analysis of convergence properties or memory capacity would differentiate this from purely empirical work. If such analysis proves intractable, the paper should explicitly discuss why (e.g., "We attempted to analyse convergence properties but found that the interaction between momentum and weight decay in the test-time setting makes formal guarantees difficult to establish because..."). This would be more valuable than omitting the discussion entirely.

• Regarding novelty, thank you for clarifying that momentary surprise equals delta rule (L1168, L1328). However, the paper would benefit from being more upfront about building on standard optimisation techniques. The momentum + delta rule combination is valuable, but shouldn't be obscured by terminology that might suggest more novelty than exists.

• The new fine-tuned results you provide (Titans achieving 94% at 1M vs Mamba's 28%) are particularly compelling and should be prominently featured in the main paper. Also, while you acknowledge the fine-tuned vs zero-shot distinction in lines 1374-1383, this critical methodological detail appears only in the supplementary material. Please consider clearly labelling which models are fine-tuned vs zero-shot directly in all figures.

• Additionally, key results demonstrating your long-context claims remain relegated to appendices (128K RULER results in Table 10, competitive baselines). For a paper whose central contribution involves long-context capabilities, these should be featured prominently in the main text.

For score improvement, I suggest:

• Adding a limitations paragraph acknowledging theoretical gaps and test-time computational overhead
• Clarifying in intro (or abstract) that the core innovation is combining existing techniques in a novel architecture
• Ensuring the planned move of appendix content to the main paper includes your key long-context results (e.g., Table 10).
• Adding clear fine-tuned vs zero-shot labels directly on all comparison figures or captions that it makes sense.

The paper has merit and will likely be accepted. These changes would significantly strengthen it and address substantive concerns raised across review cycles.

Finally, regarding ICML, I review all submissions based solely on their technical merits. My goal is helping improve the paper through specific, actionable feedback that would benefit readers. Any consistency in technical standards across venues reflects the importance of these particular issues to the work's contribution.

We thank the reviewer for their time and valuable comments. We respond to your comments below:

However, some architectural choices, such as the specific structure of the memory MLP or hyperparameter tuning strategies, could have been better justified or explored. Why use a deep MLP for memory? Have you compared against shallower versions, recurrent cells, or Hopfield-style mechanisms? Ablations would make this design choice more convincing.

Response: Thank you for mentioning this. However, please note that we already have ablations on shallow models as the memory. In Table 3 (5th row), we replace our deep memory with a linear shallow memory and the results show the effectiveness of the deep memory module. Furthermore, in Figure 6, we ablate the effect of the depth of the memory (including shallow version) on the performance of the model in different model sizes and sequence lengths. While it is interesting, exploring other architectures as the memory is out of the scope of our paper and we leave it for future studies.

Also, we would be happy to clarify any unclear point about the hyperparameters. We do not extensively tune any hyerperparameter and there are only two extra hyperparameters that we use: (1) chunk size: which is 16. The smaller the chunk is the better performance we can get but with the cost of slower model. (2) segment size: in which, we follow previous studies and use 2048 as sliding window (when exists) and 512 as segment size. We however, would be happy to clarify any specific part that has not been clear.

While the surprise-based update is interesting, it lacks deeper theoretical analysis or guarantees on memory stability or convergence.

Response: Please note that while we understand the reviewer's interests and curiosity about the theoretical side of memory modules, we want to clarify that none of the peer studies provide such theoretical results about the capacity of the memory or why one-shot memorization works (e.g., GLA (ICML 2023), Mamba2 (ICML 2024), Mamba (COLM 2024, best paper award), Gated DeltaNet (ICLR 2025), GSA (NeurIPS 2024), Simba (NeurIPS 2024), Longhorn (NeurIPS 2024), TTT (ICML 2025, spotlight), etc.). These are indeed interesting questions that one can ask and explore in future studies. Our submission is a 9-page research paper that aims to provide enough evidence for its usefulness and effect. Even the initial Transformers paper, after 8 years, is still being studied, improved, and evaluated. Therefore, we believe there are indeed many interesting questions that one can study in future work.

The treatment of related work is extensive but might benefit from a clearer contrast with specific prior methods, especially regarding what exactly distinguishes Titans from prior hybrid models.

Response: We kindly refer the reviewer to Appendix F, in which we directly compare our architectural design with most close studies, including test time training models. Also, for hybrid models, we dedicate the entire Appendix B2 to this. Similarly, we extensively discuss all TTT and fast-weight programs study in a separate sub-section (Appendix B3). However, we would be happy to provide more details if there is a specific part that the reviewer believes require more discussion.

The design space for combining memory and attention is still broad, and the paper only explores three specific configurations.

Response: We agree with the reviewer that there are a broad range of ideas that one use to combine memory and attention in different configurations. Please note that this is a 9-page research paper, and we cannot discuss all possible cases. We hope that our work motivates future studies to further study the combination of memory modules and attention.

Can you clarify how the surprise-based memory update behaves over long sequences? It's unclear how the model avoids drift or forgetting. Some visualization or stability analysis would help.

Response: Please note that defiantly forgetting and drift is possible over long sequences as the memory has a fixed size. This has always been a limitations for all recurrent models. (In our RULER benchmark you can see that with increasing the sequence length, the performance drops). However, in several experiments, we show that our memory module is more robust to forgetting and drift (Please see RULER results).

The role of persistent memory is vague. How is it trained, and how sensitive is the model to its size? More detail would clarify whether it's essential or just another form of prefix tuning.

Response: Please note that persistent memory are just learnable tokens (which can also be seen as a form of prefix tuning). Such tokens are proven to be effective from different perspectives, which we have discussed in Appendix C. Looking at our ablation study in Table 3, we can see that while persistent memories are improving the performance of the model, the performance gain is not very significant and so removing them does not affect the performance of the Titans that much. Therefore, it is fair to say that Titans' performance does not rely on them and so it is not sensitive to it at all. We would be happy to discuss this further in the final version of the paper.

We hope that our above responses have addressed your concerns, and we would be more than happy to answer any remaining concern.

We thank the reviewer for their time and valuable comments. We respond to your comments below:

There are some downsides to allowing the memory to be updated during inference-time: for example it might expose further security, safety, and alignment issues by giving an attacker access to a computable attack vector. While these don’t take away from the results presented it would be good to address them as potential limitations.

Response: Thank you very much for mentioning this. We understand the importance of potential safety and alignment risks and it would be a very important direction to understand such risks in the future studies. Following your suggestion, we will make sure to discuss this as the potential limitations and advocate for further study in the future in this direction.

Currently the experiments scale up to 1.3B parameter models, which show promising results. Further experimentation would be needed to validate whether this will continue to scale or if other issues might arise at larger 8B+ scale.

Response: Thank you for your comment. We understand the importance of scaling to larger scales. However, please note that training large models is extremely costly and challenging. Therefore, most academic research papers (such as Gated DeltaNet (ICLR 2025), Longhorn (NeurIPS 2024), TTT (ICML 2025), etc.) in this area focus on similar size models as ours. We followed the literature and reported the results on these scales.

Is there a sense from any of the tasks, either by probing or qualitative testing, what data actually gets stored in practice and what gets forgotten as a result of the surprise mechanism?

Response: Please note that as an illustrative example, we have already provided Figure 10 in Appendix that shows the surprise value of different tokens in a specific prompt from the FineWeb dataset. In this example, we can see that: the surprise metric is effective to capture important tokens that are needed for the context. Words such as “Attention!” are getting higher scores, while words such as “Currently,” “Only”, etc., are mainly ignored and forgotten.

Could the memory module be shared across modalities in combined multi-modal architectures?

Response: Thank you for asking this very interesting question. In this work our focus has been on the language modeling tasks and so there has been only a single modality. However, it would be a very interesting future direction to design multi-modal Titans. In theory, the memory can be shared across different modalities as the memory only memorizes tokens. If different modalities are already tokenized, then the memory module should be able to learn what tokens are more important from data (and so can be shared across different modalities).

How sensitive is Titans’ performance to the size of its long-term memory? Are there additional ablation studies that vary the memory capacity to see whether gains saturate or degrade beyond a certain point or tradeoffs that could be made depending on the tasks?

Response: Please note that in all variants of MAC, MAG and MAL the long-term memory size is tied to the number of parameters of the model (given the key and value features, d, the memory size is d^2). Therefore, increasing or decreasing the memory size will result in #parameter change in the entire architecture. However, in Figure 6, we have ablated the performance of the memory module on different sequence lengths and memory sizes (proportional to memory size). Based on the results across model sizes, we can see that increasing the capacity of the memory enhances its effectiveness for longer context length as the memory has more parameters to store the context in. On the other hand, in each size of memory, we can see that deeper memory modules perform better at longer context length, which further validate the importance of deeper memory. Furthermore, as you might have noticed, for each model size (please see Figure 6 (a) and (b)), increasing the context beyond a point can damage the performance. The reason for this is that memory does not have the capacity to store and compress all necessary tokens. Again, we can see that deeper memories are more robust to this pattern.

We hope that our above responses have addressed your concerns, and we would be more than happy to answer any remaining concern.

We thank the reviewer for their time and valuable comments. We respond to your comments below:

It’s unclear in MAC description: how we control the number of retrieved memories (N_l).

Response: Thank you for bringing this to our attention. As you mentioned, we fix the number of memory tokens in the context and consider the last N_l tokens. When input is shorter, we pad the sequence until it matches the length of the segments/chunks.

Following your suggestion, we will add these details to the final version of the paper and will make sure to properly discuss it.

The majority of used common-sense tasks are short (Table 1), and Titans model is basically full-attention + memory module (because sequence fits within one segment)

Response: Thank you for mentioning this. Please note that we use 512 tokens as our segment size. We will make sure to make this clear in the final version of the paper. Also, please note that even our attention-free model (Titans (LMM)) outperforms all the baselines including hybrid models. Therefore, we believe these results show the strengths of Titans memory module. While we also agree with the reviewer that better understanding of the advantages of each architecture and different methods that one can use to combine attention with memory modules requires further evaluations, we want to kindly bring to your consideration that this has not been the main focus of our submission. For our experimental setups, we aimed to follow the literature (including those with hybrid architectures) and so used the same common-sense reasoning and language modeling tasks that were previously used in similar studies. However, to fully address your concerns, if you believe it is necessary, we would be happy to evaluate the performance of Titans on any downstream tasks.

Thank you very much again for this valuable comment.

Results provided on BABILong on figure 8

Response: Thank you for bringing this to our attention. Please note that our main goal has been to compare sequence modeling backbones as the choice of different architectures can be orthogonal to this direction. That is, ARMT uses delta-rule to update its associative memory and so it can be replaced by our Titans (LMM) module, further enhancing the performance of ARMT. However, to fully address your concern, we provide the results of ARMT and compare it with Titans:

On average, Titans provide a better accuracy than ARMT.

Please clarify your chunk-segmentation procedure.

Response: Please note that we use chunk and segment for two different concepts. Chunks are subsequence that we use to accelerate the training process of the memory and our chunk size is 16. In fact, for each chunk, we take the gradient with respect to the last state of the previous chunk. On the other hand, we use segment to refer to the larger subsequence that we perform attention on. We will make sure to make this clear in the final version of the paper.

Also, regarding the different architectures, yes, you are right. For MAG and MAL, we have the same number of outputs as the input size. While on MAC, we use a fix size memory tokens. We will make sure to discuss this in the final version of the paper.

Why don’t we just use sliding window attention in MAC just like in MAG? It is possible to keep the attention on the last tokens of the previous segment.

Response: Thank you for mentioning this. That is totally possible and is a very interesting extension for the MAC architecture and a promising direction for future work. We will make sure to discuss it in the final paper.

Typos:

Response: Thank you very much for bringing this to our attention. We will make sure to proofread the paper and fix all typos.

We hope that our above responses have addressed your concerns, and we would be more than happy to answer any remaining concern.