Test-time regression: a unifying framework for designing sequence models with associative memory

• URL: http://arxiv.org/abs/2501.12352v1
• Date: Tue, 21 Jan 2025 18:32:31 GMT
• Title: Test-time regression: a unifying framework for designing sequence models with associative memory
• Authors: Ke Alexander Wang, Jiaxin Shi, Emily B. Fox,
• Abstract summary: We show that effective sequence models must be able to perform associative recall. Our key insight is that memorizing input tokens through an associative memory is equivalent to performing regression at test-time. We show numerous recent architectures -- including linear attention models, their gated variants, state-space models, online learners, and softmax attention -- emerge naturally as specific approaches to test-time regression.
• Score: 24.915262407519876
• License:
• Abstract: Sequences provide a remarkably general way to represent and process information. This powerful abstraction has placed sequence modeling at the center of modern deep learning applications, inspiring numerous architectures from transformers to recurrent networks. While this fragmented development has yielded powerful models, it has left us without a unified framework to understand their fundamental similarities and explain their effectiveness. We present a unifying framework motivated by an empirical observation: effective sequence models must be able to perform associative recall. Our key insight is that memorizing input tokens through an associative memory is equivalent to performing regression at test-time. This regression-memory correspondence provides a framework for deriving sequence models that can perform associative recall, offering a systematic lens to understand seemingly ad-hoc architectural choices. We show numerous recent architectures -- including linear attention models, their gated variants, state-space models, online learners, and softmax attention -- emerge naturally as specific approaches to test-time regression. Each architecture corresponds to three design choices: the relative importance of each association, the regressor function class, and the optimization algorithm. This connection leads to new understanding: we provide theoretical justification for QKNorm in softmax attention, and we motivate higher-order generalizations of softmax attention. Beyond unification, our work unlocks decades of rich statistical tools that can guide future development of more powerful yet principled sequence models.

Related papers

• Interactive Symbolic Regression through Offline Reinforcement Learning: A Co-Design Framework [11.804368618793273]
  Symbolic Regression holds great potential for uncovering underlying mathematical and physical relationships from observed data. Current state-of-the-art approaches typically do not consider the integration of domain experts' prior knowledge. We propose the Symbolic Q-network (Sym-Q), an advanced interactive framework for large-scale symbolic regression.
  arXiv  Detail & Related papers  (2025-02-05T06:26:49Z)
• Admitting Ignorance Helps the Video Question Answering Models to Answer [82.22149677979189]
  We argue that models often establish shortcuts, resulting in spurious correlations between questions and answers. We propose a novel training framework in which the model is compelled to acknowledge its ignorance when presented with an intervened question. In practice, we integrate a state-of-the-art model into our framework to validate its effectiveness.
  arXiv  Detail & Related papers  (2025-01-15T12:44:52Z)
• ReMatching Dynamic Reconstruction Flow [55.272357926111454]
  We introduce the ReMatching framework, designed to improve generalization quality by incorporating deformation priors into dynamic reconstruction models. The framework is highly adaptable and can be applied to various dynamic representations. Our evaluations on popular benchmarks involving both synthetic and real-world dynamic scenes demonstrate a clear improvement in reconstruction accuracy of current state-of-the-art models.
  arXiv  Detail & Related papers  (2024-11-01T16:09:33Z)
• Enhanced Transformer architecture for in-context learning of dynamical systems [0.3749861135832073]
  In this paper, we enhance the original meta-modeling framework through three key innovations. The efficacy of these modifications is demonstrated through a numerical example focusing on the Wiener-Hammerstein system class.
  arXiv  Detail & Related papers  (2024-10-04T10:05:15Z)
• Understanding the differences in Foundation Models: Attention, State Space Models, and Recurrent Neural Networks [50.29356570858905]
  We introduce the Dynamical Systems Framework (DSF), which allows a principled investigation of all these architectures in a common representation. We provide principled comparisons between softmax attention and other model classes, discussing the theoretical conditions under which softmax attention can be approximated. This shows the DSF's potential to guide the systematic development of future more efficient and scalable foundation models.
  arXiv  Detail & Related papers  (2024-05-24T17:19:57Z)
• Orchid: Flexible and Data-Dependent Convolution for Sequence Modeling [4.190836962132713]
  This paper introduces Orchid, a novel architecture designed to address the quadratic complexity of traditional attention mechanisms. At the core of this architecture lies a new data-dependent global convolution layer, which contextually adapts its conditioned kernel on input sequence. We evaluate the proposed model across multiple domains, including language modeling and image classification, to highlight its performance and generality.
  arXiv  Detail & Related papers  (2024-02-28T17:36:45Z)
• Real-World Compositional Generalization with Disentangled Sequence-to-Sequence Learning [81.24269148865555]
  A recently proposed Disentangled sequence-to-sequence model (Dangle) shows promising generalization capability. We introduce two key modifications to this model which encourage more disentangled representations and improve its compute and memory efficiency. Specifically, instead of adaptively re-encoding source keys and values at each time step, we disentangle their representations and only re-encode keys periodically.
  arXiv  Detail & Related papers  (2022-12-12T15:40:30Z)
• Robust Graph Representation Learning via Predictive Coding [46.22695915912123]
  Predictive coding is a message-passing framework initially developed to model information processing in the brain. In this work, we build models that rely on the message-passing rule of predictive coding. We show that the proposed models are comparable to standard ones in terms of performance in both inductive and transductive tasks.
  arXiv  Detail & Related papers  (2022-12-09T03:58:22Z)
• Towards a Predictive Processing Implementation of the Common Model of Cognition [79.63867412771461]
  We describe an implementation of the common model of cognition grounded in neural generative coding and holographic associative memory. The proposed system creates the groundwork for developing agents that learn continually from diverse tasks as well as model human performance at larger scales.
  arXiv  Detail & Related papers  (2021-05-15T22:55:23Z)
• Self-Reflective Variational Autoencoder [21.054722609128525]
  Variational Autoencoder (VAE) is a powerful framework for learning latent variable generative models. We introduce a solution, which we call self-reflective inference. We empirically demonstrate the clear advantages of matching the variational posterior to the exact posterior.
  arXiv  Detail & Related papers  (2020-07-10T05:05:26Z)
• SEEK: Segmented Embedding of Knowledge Graphs [77.5307592941209]
  We propose a lightweight modeling framework that can achieve highly competitive relational expressiveness without increasing the model complexity. Our framework focuses on the design of scoring functions and highlights two critical characteristics: 1) facilitating sufficient feature interactions; 2) preserving both symmetry and antisymmetry properties of relations.
  arXiv  Detail & Related papers  (2020-05-02T15:15:50Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.