Source-linked AI summary

MesaNet: Sequence Modeling by Locally Optimal Test-Time Training

Johannes von Oswald, Nino Scherrer, Seijin Kobayashi, Luca Versari, Songlin Yang, Sarthak Mittal, 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

arXiv:2506.05233v2cs.LGcs.AIcs.CL

TL;DR

Transformers have inference memory and compute that scale linearly with sequence length, motivating recurrent alternatives with constant costs. This paper introduces a stable, chunkwise-parallel Mesa layer that solves an in-context optimization problem at each time point. MesaNet results are mixed across evaluation settings: they report lower validation perplexity than several baselines, while the discussion notes underperformance in next-token prediction and long-context benchmarks relative to Transformers, alongside higher inference compute.

  • Problem

    Transformers’ per-token inference compute and memory requirements scale linearly with sequence length, motivating recurrent models with constant complexity.

    Source: §1 · p. 1
  • Method

    The Mesa layer computes optimal fast weights by minimizing a cumulative in-context regression objective, using conjugate gradients in a chunkwise-parallelizable, numerically stable implementation.

    Source: §1 · pp. 1–2, §2 · p. 4
  • Results

    MesaNets achieve lower validation perplexity than Mamba2, xLSTM, DeltaNet, Gated DeltaNet, and Transformers with the same base architecture, while the discussion reports underperformance in next-token prediction and benchmarks needing longer contexts versus Transformers.

    Source: §1 · p. 2, §6 · p. 10
  • Takeaways & Limitations

    The Mesa layer dynamically allocates test-time compute to solve its optimization problem, and its reported performance differences across sequence positions and benchmark settings accompany higher inference-time compute.

    Source: §3 · p. 5, §6 · p. 10

    Mesa inference uses roughly k times the flops of GLA, Mamba2, and xLSTM, and roughly k − 1 times those of (Gated) DeltaNet.

    Source: §3 · p. 5

Abstract

from arXiv · show

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.

1 INTRODUCTION

The paper develops an optimal fast-weight layer that minimizes a cumulative in-context regression loss, then makes the Mesa layer chunkwise parallelizable and dynamically allocates test-time computation. MesaNets achieve lower validation perplexity than several same-backbone baselines, with further analyses reporting differing results across sequence positions and downstream task types.

1 INTRODUCTION

  • The Mesa layer minimizes a quadratic loss over context to find optimal fast weights, rather than gradually learning them through gradient descent.The authors describe this as retaining past information while adapting quickly to new evidence.Source: §1 · p. 1
  • The proposed layer supports chunkwise parallel training and dynamically allocates test-time computation according to the conjugate-gradient iterations needed to meet a stopping criterion.The authors connect this allocation to a trade-off between test-time compute and performance.Source: §1 · p. 2
  • MesaNets at 140M, 440M, and 1B parameters achieve lower validation perplexity than Mamba2, xLSTM, DeltaNet, Gated DeltaNet, and Transformers using the same base architecture.Models were trained on SlimPajama.Source: §1 · p. 2
  • MesaNet outperforms modern RNNs on global reasoning, in-context learning, and in-context recall benchmarks, but still trails Transformers in in-context recall.The paper also reports that RNNs focus more on early tokens while Transformers excel at later tokens.Source: §1 · p. 2

2 A PARALLELIZABLE MESA LAYER

The Mesa layer defines its fast-weight map through a cumulative, regularized squared-error objective and solves for optimal weights at each time point. The proposed variant uses conjugate gradients and chunkwise parallel computation to address instability and sequential training constraints, while retaining recurrent inference.

2 A PARALLELIZABLE MESA LAYER

  • The Mesa layer minimizes cumulative regularized squared error over all data seen so far and learns the linear map to optimality at each time point, unlike methods using only the latest input and a single gradient step.Its quadratic objective avoids explicitly storing past tokens; the layer is described as an optimal squared-error linear associative memory that can store an association in one shot.Source: §2 · p. 4
  • The original recurrent Mesa approach is numerically unstable under forgetting and is not parallelizable, limiting its use of matrix-multiplication accelerators during training.The original approach requires a regularizer that decays exponentially with time.Source: §2 · p. 4
  • The layer updates state variables G_t and H_t through linear recurrences using input and forget gates, supporting the query-specific linear solve.The states summarize information used by the Mesa computation.Source: §2 · p. 4
  • The proposed variant solves a linear system for each query using conjugate gradients; this remains numerically stable regardless of forgetting strength, with an additional state matrix below 1% of total model memory.The solver computes q*_t = (H_t + Λ)^-1 q_t.Source: §2 · p. 4
  • Chunkwise parallelization uses conjugate-gradient computations in gated-linear-attention form, enabling parallel training with O(T) complexity and recurrent inference with O(1) complexity.The method exploits matrix-matrix multiplication and established efficient chunkwise linear-attention implementations.Source: §2 · pp. 4–5

3 TRAIN AND INFERENCE TIME OF THE MESA LAYER

The Mesa layer remains competitive in training and autoregressive sampling despite its higher computational cost. Its inference cost depends on the number of conjugate-gradient steps, which can vary with the sequence being processed.

3 TRAIN AND INFERENCE TIME OF THE MESA LAYER

  • MesaNet remains competitive at train time against MHA and RNN alternatives despite solving t · H linear systems per layer and computing gradients through their solutions.Figure 2 compares training time on TPUv5 and H100; its H100 throughput comparison uses 1B-parameter models.Source: §3 · p. 5, Figure 2 · p. 5
  • Mesa's inference cost is roughly k times that of GLA, Mamba2, and xLSTM, and k −1 times that of (Gated) DeltaNet.For a given sequence length, CG cost growing with kn^2 a sets a maximum k for which Mesa uses fewer flops than MHA.Source: §3 · p. 5
  • MesaNet has competitive autoregressive sampling times with other linear attention variants, while linear attention models have improved throughput over transformer baselines.This holds for the 1B-parameter model on H100 and TPUv5, despite Mesa layers consuming significantly more flops.Source: §3 · p. 6
  • The number of CG steps needed to meet a desired error tolerance can vary by head, sequence, and token, making inference cost dynamic.A stopping criterion reallocates compute according to the incoming data, unlike softmax self-attention's sequence-length-dependent compute and memory.Source: §3 · p. 6

4 MESANET IN A LANGUAGE WORLD

On 1B models trained on 50B SlimPajama tokens, this section compares sequence layers on a shared backbone across language modeling and downstream benchmarks. MesaNet and Hawk-Mesa perform strongly against recurrent baselines, with results varying by token position and task.

4.1 LANGUAGE MODELING (WITHIN AND BEYOND TRAIN SEQUENCE LENGTH)

  • MesaNet matches the Transformer’s average per-token PPL, while Hawk-Mesa exceeds it and both outperform recurrent baselines across the reported benchmarks.The comparison evaluates 1B models on sequence length 2048.Source: Table 2 · p. 7, §4.1 · p. 7
  • MesaNet and Hawk-Mesa extend their NLL advantage over the Transformer beyond 512 tokens, despite recurrent layers generally performing better only early in sequences.Most recurrent layers show an advantage up to 64 tokens; Hawk is strongest there but declines sharply afterward.Source: Figure 3 · p. 7, §4.1 · p. 8

4.2 LANGUAGE BENCHMARKS

  • Across most downstream benchmarks, MesaNet matches or exceeds recurrent baselines; on global reasoning it leads recurrent models but remains below the Transformer.The grouped results cover five benchmark subsets.Source: Table 3 · p. 8, §4.2 · pp. 8–9

4.1 LANGUAGE MODELING (WITHIN AND BEYOND TRAIN SEQUENCE LENGTH)

  • At sequence lengths up to 32k, MesaNet extrapolates better than Hawk, GLA, and xLSTM and on par with Gated DeltaNet, but SWA-1024 remains competitive.SWA-1024 is even superior at 32k on some benchmarks.Source: §4.1 · p. 8, Figure 4 · p. 8

4.2 LANGUAGE BENCHMARKS

  • Local-reasoning scores are similar across models, and Hawk performs well despite weaker results on global reasoning and in-context recall.The authors suggest the subsets measure different capabilities.Source: §4.2 · p. 9
  • On in-context recall, MesaNet outperforms recurrent models and SWA-1024, but trails the Transformer with a 2048-token attention window.Source: §4.2 · p. 9
  • MesaNet exceeds the Transformer on word-scrambling few-shot tasks but not on translation tasks.The evaluations include two word-scrambling and three translation tasks.Source: §4.2 · p. 9

5 TEST-TIME COMPUTE ANALYSIS

The analysis tests whether dynamic stopping can reduce MesaNet’s inference-time conjugate-gradient work. A stopping criterion of ϵ = 10^-4 matches the fixed-step model’s performance with fewer average steps.

5 TEST-TIME COMPUTE ANALYSIS

  • Condition numbers vary across heads: roughly 1–2 heads per layer grow with sequence length, motivating dynamic CG-step allocation.Source: §5 · p. 9
  • Uniformly reducing CG steps increases NLL, especially for later tokens; increasing the dynamic stopping threshold degrades performance uniformly over sequence length.Source: §5 · p. 10, Figure 5 · p. 10
  • ϵ = 10^-4 performs on par with the 30-step base model while reducing average CG steps to ≈9.Source: Figure 5 · p. 10, §5 · p. 10

6 DISCUSSION

The paper presents a scalable, numerically stable Mesa layer that computes locally optimal fast weights, then examines its language-modeling behavior and inference costs. Its discussion finds a mixed performance profile: recurrent models excel early in sequences but can lag on longer-context tasks.

6 DISCUSSION

  • MesaNet scales to 1B parameters with a chunkwise-parallel, numerically stable Mesa layer that dynamically allocates test-time computation.The layer solves an optimization problem to produce a linear model fitting the sequence.Source: §6 · p. 10
  • The Mesa layer replaces slower gradient-descent updates to fast weights with locally optimal learning.The paper relates this to fast weights and describes the change as replacing Hebbian with locally optimal learning.Source: §6 · p. 10
  • RNNs, especially MesaNets, outperform transformers early in sequences but underperform in next-token prediction and benchmarks requiring longer contexts.The authors note that RNNs have inference-time advantages over transformers in the long-context regime and call for further investigation.Source: §6 · p. 10
  • Increased test-time compute remains MesaNet’s biggest shortcoming, despite its dynamic allocation.The authors suggest warm-starting from neighboring solutions or using chunkwise-parallel CG during training and Sherman–Morrison recursion at inference as directions for future work.Source: §6 · p. 10

REPRODUCIBILITY STATEMENT

The appendix documents implementation and comparison practices, then situates MesaNet among efficient sequence models. It connects recurrent updates to online regression, forgetting mechanisms, and fixed-point methods.

REPRODUCIBILITY STATEMENT

  • The authors provide conjugate-gradient pseudocode, numerical-precision details, and plan to release a Triton implementation and educational notebooks.They report extensive hyperparameter scans of related work and an apples-to-apples comparison using the same backbone while varying only the sequence layer.Source: REPRODUCIBILITY STATEMENT · p. 11

A RELATED WORK

  • Linear attention replaces the softmax kernel with a feature-map inner product, yielding recurrent key-value accumulation; this general form connects to Fast Weight Programmers.The passage traces linear-attention designs to work on feature maps and key-value normalization.Source: A RELATED WORK · p. 20
  • DeltaNet interpolates new values with stored associations and performs a gradient-descent step on squared prediction error; related models adapt or batch this update.The passage also describes Titans adding momentum and identifies Atlas as a sliding-window variant of Mesa.Source: A RELATED WORK · p. 20
  • Mesa’s linear fast-weight structure permits more efficient conjugate-gradient optimization than DEQ-style fixed-point methods.The comparison also places Mesa alongside recurrent-depth models that use iterative blocks or solve fixed points.Source: A RELATED WORK · p. 20
  • Modern linear-attention models use increasingly input-dependent forgetting, including per-row gates, while SSMs range from convolution-compatible LTI transitions to input-dependent Mamba transitions.Mamba-2 constrains its transition matrix to a scalar times identity and is equivalent to gated linear attention.Source: A RELATED WORK · pp. 20–21

B DERIVATION OF PREVIOUS TEST-TIME TRAINING RULES

This section derives several recurrent sequence-layer updates from online learning objectives. It contrasts their losses and update rules, including Mesa’s use of the full history of forget-weighted errors.

B DERIVATION OF PREVIOUS TEST-TIME TRAINING RULES

  • Table 4 organizes test-time-training recurrent layers by their update rules and corresponding online objectives.The section introduces these models as optimizing sequences of objectives.Source: Table 4 · p. 21, B DERIVATION OF PREVIOUS TEST-TIME… · p. 21
  • GLA follows from a quadratic continuous-state Hopfield energy with a forget- and input-gate-dependent regularizer, while DeltaNet follows from squared error with the regularizer disabled.DeltaNet’s squared-error update is identified with the original model’s rule.Source: B DERIVATION OF PREVIOUS TEST-TIME… · pp. 21–22
  • LongHorn uses a vector-valued error precision and a diagonal input-dependent step size, but its implementation approximates the derived update with elementwise operations.Its objective uses the current key-value pair’s squared error and a memory quadratic potential pulling toward the previous solution.Source: B DERIVATION OF PREVIOUS TEST-TIME… · p. 22
  • Unlike LongHorn’s instantaneous loss, Mesa optimizes the full forget-weighted sum of squared errors from the sequence start through the current time.Concurrent Atlas work uses a sliding-window version of the Mesa objective and optimizes losses online with a second-order Muon method.Source: B DERIVATION OF PREVIOUS TEST-TIME… · pp. 22–23

C RANK-ONE UPDATE CONJUGATE GRADIENT METHOD

The appendix develops a conjugate-gradient solution for Mesa and a chunkwise-parallel implementation that avoids storing every recurrent matrix. It then details the layer and model design used for comparisons.

C RANK-ONE UPDATE CONJUGATE GRADIENT METHOD

  • The CG solver approximates q*_t without materializing H_t at every time step, storing H only at chunk boundaries to reduce memory overhead.The method uses chunkwise-parallel computation and applies CG to (H_t + Λ_t)^-1 q_t; its inference cost includes O(kn^2a) work.Source: C RANK-ONE UPDATE CONJUGATE GRADIENT… · pp. 24–25

D CHUNKWISE PARALLEL FORM OF GATED LINEAR ATTENTION AND THE MESA LAYER

  • Mesa’s chunkwise form replaces the standard query q_t with q*_t = (H_t + Λ)^-1q_t, which CG computes using operations compatible with chunkwise parallelism.The resulting computation can leverage matrix-matrix accelerators on GPUs or TPUs without extensive memory overhead.Source: D CHUNKWISE PARALLEL FORM OF GATED… · p. 26

E A FULL DESCRIPTION OF THE MESA LAYER, RELATED WORK AND THE MESANET

  • The layer uses recurrent state variables with forget and input gates, normalized keys and queries, and a fixed positive regularization strength over time.Keys and queries are RMS-normalized, convolved with a window of four, passed through SiLU, then normalized to unit L2 norm; the regularization parameter is made positive with softplus.Source: E A FULL DESCRIPTION OF THE MESA… · p. 28

E.1 MODEL DESIGN

  • For close sequence-layer comparisons, the models share key size and query, key, and value parameter counts, with common gate and convolution choices across RNNs except mLSTM.The overview also describes a shared embedding and blocks that pair a sequence layer with an MLP.Source: §E.1 · p. 30

F.1 MAD BENCHMARK SUITE

The appendices specify the synthetic-benchmark procedures and the data and model choices used for language experiments. They describe shared evaluation and architecture controls, along with a possible bias toward shorter sequences in the data preparation.

F.1 MAD BENCHMARK SUITE

  • MAD evaluation averages accuracy across subtasks of varying difficulty and reports each model’s maximum average task accuracy after learning-rate and weight-decay sweeps.The subtasks vary sequence length, training examples, vocabulary size, and task-specific parameters.Source: §F.1 · pp. 30–31

F.2 REGBENCH IN-CONTEXT LANGUAGE LEARNING BENCHMARK

  • RegBench reports test accuracy for configurations selected by grid search over shared hyperparameters across models.Source: §F.2 · p. 31

G EXPERIMENTAL DETAILS: MESANET IN A LANGUAGE WORLD

G.1 DATA

  • Language models use SlimPajama and a GPT-2 tokenizer; training and validation sequences are filled to length 2048 with randomly sampled sequences, which may bias evaluation toward shorter sequences.The experiments use datasets of 15B and 50B tokens.Source: §G.1 · pp. 31–32

G.2 MODEL DESIGN

  • Models share a transformer-style residual-block backbone with a sequence layer and gated MLP, while sequence-layer comparisons use matched key sizes and parameter counts.The RNN comparisons also standardize one-dimensional write and forget gates and use key/query convolutions with window size 4.Source: §G.2 · p. 32

G.3 TRAINING DETAILS

These appendices document training and solver choices, compute and throughput measurements, the original recursive Mesa formulation, and a preliminary parity state-tracking experiment. They note solver-precision limitations and stability considerations.

G.3 TRAINING DETAILS

  • Training uses AdamW, batch size 256, and mixed precision; BF16 activations with FP32 weights slightly improve next-token perplexity over using FP32 throughout.The training procedure uses cross-entropy on next tokens and gradient-norm clipping.Source: §G.3 · pp. 32–33

G.4 HYPERPARAMETER SCANS

  • The language experiments use 140M, 440M, and 940M parameter models, with learning rates selected through scans for the 15B- and 50B-token runs.Source: §G.4 · p. 33

G.5 NOTES ON PRECISION USED IN THE CG-SOLVER, MESA LAYER DESIGN CONSIDERATIONS OR Why you shouldn’t scream at your Mesa layer

  • Key/query normalization and a learned regularizer Λ lower-bounded at 0.25 are used to stabilize Mesa; the CG solver uses FP32 matrix multiplication, while lower-precision convergence remains uninvestigated.The authors caution that lower precision on GPUs may require additional work for stable solver convergence.Source: §G.5 · pp. 33–34

G.6 EXPERIMENTS COMPUTE RESOURCES

  • A 1B run on 50B tokens took about 36 hours on TPUv5; Mesa-layer investigations at 400M parameters took 3–12 hours depending on CG steps and data parallelization.Source: §G.6 · p. 34

G.7 TOKEN THROUGHPUT COMPARISONS OF RECURRENT MODELS FOR 1B MODELS

  • MesaNet throughput is competitive, especially with 10 CG steps, which suffice for the reported superior results; Gated Linear Attention is faster, while Transformer throughput degrades as sequence length grows.The throughput comparison measures 1B models on H100 and TPUv5 systems.Source: §G.7 · p. 34, Figure 6 · p. 35

H THE ORIGINAL RECURSIVE LEAST-SQUARES MESA LAYER

  • The original recursive Mesa update is an online second-order Newton method, but its inversion recursion lacks parallel-in-time computation and can be unstable with small forget gates.The original formulation also forgets its regularization term exponentially fast.Source: H THE ORIGINAL RECURSIVE LEAST-SQUARES… · pp. 34, 36

I A PRELIMINARY INVESTIGATION INTO STATE TRACKING WITH THE MESA LAYER

  • With γ_t ∈ (−1, 1) and high Λ, MesaNet reaches perfect parity accuracy, but the authors characterize this setting as functionally close to GLA rather than the desired use of Mesa’s extra computation.The experiment trains on sequences of length 40 and evaluates up to length 256.Source: Figure 8 · p. 36, I A PRELIMINARY INVESTIGATION INTO… · p. 37

J FURTHER DISCUSSION POINTS

The discussion and appendices report synthetic-task results, examine language-modeling metrics and benchmark behavior, and identify limitations in solver training and architecture comparisons. MesaNet performs strongly on several synthetic and recall tasks, while results vary with task scope and evaluation choices.

J FURTHER DISCUSSION POINTS

  • The authors identify two open directions: training through unrolled CG steps to support fewer-step inference, and comparing sequence layers under alternative architecture backbones.They also describe learned fast inversion algorithms and continuous-data sequence modeling as possible research directions.Source: J FURTHER DISCUSSION POINTS · pp. 37–38

K MESANET TRAINED IN SYNTHETIC ENVIRONMENTS

  • MesaNet achieves the highest average MAD performance among linear recurrent architectures and matches Transformers; on RegBench it surpasses other RNNs and matches Transformers.MAD covers recall, memorization, compression, and copying; RegBench tests inference of held-out probabilistic finite automata from context.Source: K MESANET TRAINED IN SYNTHETIC… · p. 39, Table 9 · p. 39, Figure 9 · p. 39

L.1 LANGUAGE MODELLING / PERPLEXITY ANALYSES

  • Position-conditioned NLL is used because aggregate perplexity weights tokens equally and can mask qualitative differences between early and late sequence positions.Mean-so-far metrics can also obscure transition points because they aggregate cumulatively up to each position.Source: §L.1 · pp. 39–40

L.1.1 WITHIN TRAIN CONTEXT-LENGTH

  • Within the training context, most recurrent models outperform the Transformer early in sequences, but the Transformer surpasses them beyond a token position; Hawk-MesaNet maintains its advantage to the greatest depth.Hawk has the strongest early performance up to a point, then declines sharply.Source: §L.1.1 · p. 40, Figure 10 · p. 41, Figure 11 · p. 41

L.1.2 BEYOND TRAIN CONTEXT-LENGTH

L.2 DOWNSTREAM BENCHMARKS

  • The downstream evaluation covers zero-shot common-sense reasoning, in-context recall, and few-shot learning, with raw scores and sliding-window reference models reported by benchmark.Source: §L.2 · p. 44

L.2.1 ZERO-SHOT COMMON-SENSE REASONING BENCHMARKS

  • MesaNet and Hawk-MesaNet perform strongly on the global reasoning subset, while local-set scores vary little across models; the split may change with larger models or different datasets.The split protocol uses short-window attention comparisons to classify benchmarks as local or global reasoning.Source: §L.2.1 · p. 45, Table 14 · p. 46

L.2.2 IN-CONTEXT RECALL BENCHMARKS

  • MesaNet attains the best or second-best average in-context recall scores across evaluated settings, while sliding-window scores generally rise with window size.SQUAD is an exception where SWA-1024 exceeds the 2048-window Transformer baseline.Source: §L.2.2 · p. 47, Table 15 · p. 47, Table 16 · p. 48

L.2.3 FEW-SHOT LEARNING BENCHMARKS

Few-shot evaluations find strong MesaNet results on most word-scrambling settings and among recurrent models on translation, but recurrent models trail transformers on translation. Needle-in-a-haystack results are strong and stable on noise haystacks, while essay-haystack variability limits conclusions.

L.2.3 FEW-SHOT LEARNING BENCHMARKS

  • MesaNet attains the strongest scores in most word-scrambling settings and significantly outperforms the transformer baseline.The reported results average over 10 randomly drawn few-shot prefixes; three of the five tasks had scores close to 0% and are omitted.Source: Table 17 · p. 49, §L.2.3 · p. 49
  • MesaNet scores best among recurrent models in most translation settings, but recurrent models trail transformers by a relatively large margin, especially at 1B parameters.Scores differ across French, German, and Romanian tasks; the authors suggest training-corpus language prevalence as a likely contributor.Source: Table 18 · p. 50, §L.2.3 · p. 50

L.3 NEEDLE IN THE HAYSTACK (NIAH) RESULTS

  • For 1B models trained on 50B tokens, the NIAH sweep covered noise and essay haystacks, nine key/value combinations, and context lengths 2048 and 4096.Results report mean ± standard deviation over 1,000 samples per setting.Source: §L.3 · p. 50, Table 19 · p. 51
  • MesaNet shows strong NIAH scores with low fluctuations across key/value combinations on noise haystacks, while essay-haystack fluctuations make conclusions difficult.The authors nevertheless highlight Hawk-Mesa’s strong essay-haystack performance.Source: §L.3 · p. 50

M VARYING THE NUMBER OF CONJUGATE GRADIENT STEPS WHEN TRAINING MESANETS

Training with fewer conjugate-gradient steps modestly increases loss, with a larger increase when steps fall below 10. The section also specifies evaluation procedures and relates high forget-gate values to greater solver effort.

M VARYING THE NUMBER OF CONJUGATE GRADIENT STEPS WHEN TRAINING MESANETS

  • Dropping training CG steps below 10 produces a more drastic loss increase; reducing them otherwise causes a small, uniform increase relative to 30 steps.The 400M-parameter, 15B-token comparison shows a substantial increase when steps fall to 5; the authors chose 30 steps to prioritize a fully converged Mesa layer over training flops.Source: M VARYING THE NUMBER OF CONJUGATE… · p. 51, Figure 13 · p. 51

N EVALUATION METHODOLOGY

  • Multiple-choice tasks rank options by token-normalized conditional loss, greedy matching checks whether every target token is an argmax, and in-context recall checks whether a greedily generated completion contains the target.Recall generation stops at 48 tokens or a newline and uses a non-case-sensitive target match.Source: N EVALUATION METHODOLOGY · pp. 51–52

O AN INTERNAL ANALYSIS OF THE MESANET

  • Heads with γ_t near 1 usually have increasing condition numbers over time, requiring more CG steps; their no-CG outputs also tend to resemble Mesa outputs more closely by cosine similarity.CG step counts are measured to an error of 0.001, and heads are ordered by average γ for plotting.Source: Figure 14 · p. 52, O AN INTERNAL ANALYSIS OF THE MESANET · p. 52
Loading 2506.05233v2…