MesaNet: Sequence Modeling by Locally Optimal Test-Time Training

• URL: http://arxiv.org/abs/2506.05233v1
• Date: Thu, 05 Jun 2025 16:50:23 GMT
• Title: MesaNet: Sequence Modeling by Locally Optimal Test-Time Training
• Authors: Johannes von Oswald, Nino Scherrer, Seijin Kobayashi, Luca Versari, Songlin Yang, Maximilian Schlegel, Kaitlin Maile, Yanick Schimpf, Oliver Sieberling, Alexander Meulemans, Rif A. Saurous, Guillaume Lajoie, Charlotte Frenkel, Razvan Pascanu, Blaise Agüera y Arcas, João Sacramento,
• Abstract summary: We introduce a numerically stable, chunkwise parallelizable version of the recently proposed Mesa layer.<n>We show that optimal test-time training enables reaching lower language modeling perplexity and higher downstream benchmark performance than previous RNNs.
• Score: 67.45211108321203
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: Sequence modeling is currently dominated by causal transformer architectures that use softmax self-attention. Although widely adopted, transformers require scaling memory and compute linearly during inference. A recent stream of work linearized the softmax operation, resulting in powerful recurrent neural network (RNN) models with constant memory and compute costs such as DeltaNet, Mamba or xLSTM. These models can be unified by noting that their recurrent layer dynamics can all be derived from an in-context regression objective, approximately optimized through an online learning rule. Here, we join this line of work and introduce a numerically stable, chunkwise parallelizable version of the recently proposed Mesa layer (von Oswald et al., 2024), and study it in language modeling at the billion-parameter scale. This layer again stems from an in-context loss, but which is now minimized to optimality at every time point using a fast conjugate gradient solver. Through an extensive suite of experiments, we show that optimal test-time training enables reaching lower language modeling perplexity and higher downstream benchmark performance than previous RNNs, especially on tasks requiring long context understanding. This performance gain comes at the cost of additional flops spent during inference time. Our results are therefore intriguingly related to recent trends of increasing test-time compute to improve performance -- here by spending compute to solve sequential optimization problems within the neural network itself.

Related papers

• A Scalable Hybrid Training Approach for Recurrent Spiking Neural Networks [13.220581846415957]
  In this work, we introduce HYbrid PRopagation (HYPR) that combines the efficiency of parallelization with approximate online forward learning.<n>HYPR enables parallelization of parameter update over the sub sequences for RSNNs consisting of almost arbitrary non-linear spiking neuron models.<n>We find that this type of neuron model is particularly well trainable by HYPR, resulting in an unprecedentedly low task performance gap between approximate forward gradient learning and BPTT.
  arXiv  Detail & Related papers  (2025-06-17T12:27:25Z)
• Bidirectional Linear Recurrent Models for Sequence-Level Multisource Fusion [10.867398697751742]
  We introduce BLUR (Bidirectional Linear Unit for Recurrent network), which uses forward and backward linear recurrent units (LRUs) to capture both past and future dependencies with high computational efficiency.<n>Experiments on sequential image and time series datasets reveal that BLUR not only surpasses transformers and traditional RNNs in accuracy but also significantly reduces computational costs.
  arXiv  Detail & Related papers  (2025-04-11T20:42:58Z)
• Fast Training of Recurrent Neural Networks with Stationary State Feedbacks [48.22082789438538]
  Recurrent neural networks (RNNs) have recently demonstrated strong performance and faster inference than Transformers.<n>We propose a novel method that replaces BPTT with a fixed gradient feedback mechanism.
  arXiv  Detail & Related papers  (2025-03-29T14:45:52Z)
• Multi-Layer Transformers Gradient Can be Approximated in Almost Linear Time [17.086679273053853]
  We show that a novel fast approximation method can calculate the gradients in almost linear time. By improving the efficiency of gradient, we hope that this work will facilitate more effective training and deployment of long-context language models.
  arXiv  Detail & Related papers  (2024-08-23T17:16:43Z)
• A Dynamical Model of Neural Scaling Laws [79.59705237659547]
  We analyze a random feature model trained with gradient descent as a solvable model of network training and generalization. Our theory shows how the gap between training and test loss can gradually build up over time due to repeated reuse of data.
  arXiv  Detail & Related papers  (2024-02-02T01:41:38Z)
• Intelligence Processing Units Accelerate Neuromorphic Learning [52.952192990802345]
  Spiking neural networks (SNNs) have achieved orders of magnitude improvement in terms of energy consumption and latency. We present an IPU-optimized release of our custom SNN Python package, snnTorch.
  arXiv  Detail & Related papers  (2022-11-19T15:44:08Z)
• Oscillatory Fourier Neural Network: A Compact and Efficient Architecture for Sequential Processing [16.69710555668727]
  We propose a novel neuron model that has cosine activation with a time varying component for sequential processing. The proposed neuron provides an efficient building block for projecting sequential inputs into spectral domain. Applying the proposed model to sentiment analysis on IMDB dataset reaches 89.4% test accuracy within 5 epochs.
  arXiv  Detail & Related papers  (2021-09-14T19:08:07Z)
• Liquid Time-constant Networks [117.57116214802504]
  We introduce a new class of time-continuous recurrent neural network models. Instead of declaring a learning system's dynamics by implicit nonlinearities, we construct networks of linear first-order dynamical systems. These neural networks exhibit stable and bounded behavior, yield superior expressivity within the family of neural ordinary differential equations.
  arXiv  Detail & Related papers  (2020-06-08T09:53:35Z)
• Communication-Efficient Distributed Stochastic AUC Maximization with Deep Neural Networks [50.42141893913188]
  We study a distributed variable for large-scale AUC for a neural network as with a deep neural network. Our model requires a much less number of communication rounds and still a number of communication rounds in theory. Our experiments on several datasets show the effectiveness of our theory and also confirm our theory.
  arXiv  Detail & Related papers  (2020-05-05T18:08:23Z)
• Convolutional Tensor-Train LSTM for Spatio-temporal Learning [116.24172387469994]
  We propose a higher-order LSTM model that can efficiently learn long-term correlations in the video sequence. This is accomplished through a novel tensor train module that performs prediction by combining convolutional features across time. Our results achieve state-of-the-art performance-art in a wide range of applications and datasets.
  arXiv  Detail & Related papers  (2020-02-21T05:00:01Z)

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.