Learning Associative Inference Using Fast Weight Memory
Imanol Schlag, Tsendsuren Munkhdalai, Jürgen Schmidhuber
Abstract
Humans can quickly associate stimuli to solve problems in novel contexts. Our novel neural network model learns state representations of facts that can be composed to perform such associative inference. To this end, we augment the LSTM model with an associative memory, dubbed Fast Weight Memory (FWM). Through differentiable operations at every step of a given input sequence, the LSTM updates and maintains compositional associations stored in the rapidly changing FWM weights. Our model is trained end-to-end by gradient descent and yields excellent performance on compositional language reasoning problems, meta-reinforcement-learning for POMDPs, and small-scale word-level language modelling. 1
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Cited by top-tier papers22
- Linear Transformers Are Secretly Fast Weight ProgrammersImanol Schlag, Kazuki Irie, Jürgen SchmidhuberICML 2021 · 394 citations
- Going Beyond Linear Transformers with Recurrent Fast Weight ProgrammersKazuki Irie, Imanol Schlag, Róbert Csordás, Jürgen SchmidhuberNeurIPS 2021 · 101 citations
- StepGame: A New Benchmark for Robust Multi-Hop Spatial Reasoning in TextsZhengxiang Shi, Qiang Zhang, Aldo LipaniAAAI 2022 · 100 citations
- Gated Slot Attention for Efficient Linear-Time Sequence ModelingYu Zhang, Songlin Yang, Rui-Jie Zhu, Yue Zhang et al.NeurIPS 2024 · 93 citations
- DeltaProduct: Improving State-Tracking in Linear RNNs via Householder ProductsJulien Siems, Timur Carstensen, Arber Zela, Frank Hutter et al.NeurIPS 2025 · 75 citations
Builds on4
- Hopfield Networks is All You NeedHubert Ramsauer, Bernhard Schäfl, Johannes Lehner, Philipp Seidl et al.ICLR 2021 · 620 citations
- Improving Generalization in Meta Reinforcement Learning using Learned ObjectivesLouis Kirsch, Sjoerd van Steenkiste, Jürgen SchmidhuberICLR 2020 · 132 citations
- Self-Attentive Associative MemoryHung Le, Truyen Tran, Svetha VenkateshICML 2020 · 61 citations
- Meta-Learning Deep Energy-Based Memory ModelsSergey Bartunov, Jack W. Rae, Simon Osindero, Timothy P. LillicrapICLR 2020 · 35 citations
Related papers
- Blending Complementary Memory Systems in Hybrid Quadratic-Linear TransformersKazuki Irie, Morris Yau, Samuel J. GershmanNeurIPS 2025 · 13 citations
- Memorizing TransformersYuhuai Wu, Markus Norman Rabe, DeLesley Hutchins, Christian SzegedyICLR 2022 · 231 citations
- Do LLMs dream of elephants (when told not to)? Latent concept association and associative memory in transformersYibo Jiang, Goutham Rajendran, Pradeep Ravikumar, Bryon AragamNeurIPS 2024 · 19 citations
- Compositional Generalization by Learning Analytical ExpressionsQian Liu, Shengnan An, Jian-Guang Lou, Bei Chen et al.NeurIPS 2020 · 79 citations
- Co-occurrence is not Factual Association in Language ModelsXiao Zhang, Miao Li, Ji WuNeurIPS 2024 · 15 citations
