Kaitlin Maile
Research Areas
Authored Publications
Sort By
Preview abstract
Looping, reusing a block of layers across depth, and depth growing, training shallow-to-deep models by duplicating middle layers, have both been linked to stronger reasoning, but their relationship remains unclear. We provide a mechanistic unification: looped and depth-grown models exhibit convergent depth-wise signatures, including increased reliance on late layers and recurring patterns aligned with the looped or grown block. These shared signatures support the view that their gains stem from a common form of iterative computation. In particular, applying inference-time looping to the middle blocks of a depth-grown model yields additional improvements in reasoning primitives accuracy, up to 2x, despite the model never being trained to loop. Building on this connection, we show that the two techniques are both adaptable and composable. Both approaches adapt better than the baseline when given more in-context examples or additional supervised fine-tuning data. Additionally, depth-grown models achieve the largest reasoning gains when using higher-quality, math-heavy cooldown mixtures, which can be further boosted by adapting a looped block in the middle of the network. Overall, our results position depth growth and looping as complementary, practical methods for inducing and scaling iterative computation to improve reasoning.
View details
MesaNet: Sequence Modelling by Locally Optimal test-Time Training
Songlin Yang
Razvan Pascanu
Alexander Meulemans
Seijin Kobayashi
Yanick Schimpf
João Sacramento
Maximilian Schlegel
Luca Versari
Oliver Sieberling
Rif A. Saurous
2025
Preview 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), which could only run sequentially in time and was
therefore not scalable. 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 study up to the billion-parameter
scale, 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.
View details
