From CUDA to MLX: How K-Search Brings Decades of Kernel Expertise to Apple Silicon
Figure 1: CUDA-to-MLX optimization translation map. CUDA optimization knowledge can be translated into architecture-native MLX strategies rather than copied instruction-for-instruction.
We face a new epoch in computing. Hardware is changing rapidly — not just faster GPUs, but a growing range of chips from different vendors, each with its own architecture and often tailored to specific AI workloads. Software is changing just as fast, and AI coding tools now generate in minutes what took months of effort a few years ago.
With so much of computing now centered on AI, GPU kernels are a crucial component of its success. These are the low-level programs that run inside the GPU, and writing efficient ones is far from obvious — it takes years of expertise to get right. Transferring a kernel from one vendor’s hardware to another is harder still, and often means rediscovering the same optimizations from scratch. The CUDA ecosystem, for example, has accumulated decades of hard-won kernel expertise: hand-tuned implementations of attention, state space models, and other critical operations representing thousands of engineering hours. Newer hardware ecosystems (Apple Silicon, custom AI accelerators, and others) are growing fast but lack this depth.
In this work we ask whether that expertise can be transferred automatically. We built on K-Search, an evolutionary kernel search framework introduced by Cao et al. at Berkeley Sky Lab that uses AI to optimize GPU kernels, and extended it with a backend for MLX — Apple’s machine-learning framework for its own Apple Silicon chips. We developed a novel structured CUDA-to-MLX translation layer that lets K-Search take existing CUDA kernels as a knowledge base and adapt them into high-quality GPU kernels for Apple Silicon, rather than rebuilding from scratch.
We show that our approach reaches near-expert level performance on Apple Silicon with 0.97x speedup compared to the native MLX Attention kernel, and up to a 20x prefill speedup over the community mlx-lm implementation on the Mamba SSM kernel; we report the numbers, and how much of the gain comes from the translation layer, in the sections below. Although we focus on MLX kernels for Apple Silicon, the method is not specific to MLX and applies to any ecosystem where CUDA expertise is transferable.
Why MLX?
Apple’s MLX framework has seen remarkable adoption since late 2023. With Apple Silicon in hundreds of millions of MacBooks and Mac Studios, MLX enables local AI inference without cloud costs. The unified memory architecture makes it especially attractive for mid-sized models (7B–70B parameters on M series chips).
Yet beneath this momentum lies a significant gap: many performance-critical kernels that the NVIDIA ecosystem takes for granted: paged attention, optimized SSM scan kernels, fused MoE routing are either absent or naive without hardware-specific tuning. MLX runs models correctly but often leaves significant performance on the table.
This gap is what motivates the rest of this post.
What is K-Search?
K-Search is an evolutionary kernel optimization framework originally developed by our first author Shiyi Cao at UC Berkeley Sky Lab. Given a naive kernel and a hardware specification, it runs an iterative optimization loop: an LLM reasons about which optimizations to try next, a code-writing model generates candidate kernels, and those candidates are compiled and benchmarked on real hardware.
Measurements feed back into the search, which keeps refining, pursuing promising directions and dropping dead ends until performance converges.
Algorithm 1: K-Search via co-evolving world models. The search alternates between selecting the most promising action, instantiating and evaluating code until improvement stagnates, and evolving the world model through insert, update, and prune operations. Adapted from Cao et al. (2026).
Search is grounded by a Spec: a domain-specific document encoding hardware rules, optimization patterns, and mathematical constraints which keeps generated code from hallucinating invalid primitives and ensures candidates will actually compile and run efficiently.
In our runs, a single model (Gemini 3.5 Pro Preview) plays both roles: it maintains the reasoning state and writes the kernels. The reasoning half is prompted as a “GPU kernel performance engineer” and asked to work through a fixed analysis before proposing anything: classify the kernel (reduction, scan, attention/softmax, …), rewrite the reference computation in canonical form, map out data layout and access patterns, and hypothesize the likely bottleneck (bandwidth, latency, compute, or synchronization) in each runtime regime. Only then does it emit candidate optimizations, each as a single change implementable in one iteration.
We call the persistent reasoning state a world model. Rather than a flat list of things to try, it is a decision (prefix) tree: each root→leaf path composes a full optimization plan, and sibling branches are competing alternatives. Every node is scored — an overall_rating in [0, 10], a confidence in [0, 1], and per-node impacts on memory bandwidth, register pressure, and compute/hardware fit — so the search can rank partial plans and expand the most promising ones. The tree persists and grows across rounds: refining an idea adds a child node rather than overwriting its parent, and if the best score fails to improve for a few rounds (a stagnation window) the search backs off to explore an alternative branch. A single node, as it appears mid-run on the attention kernel, looks like this:
{
"action": "Replace the threadgroup-memory softmax reduction
with a register-only reduction: each SIMD group
owns 8 query rows and reduces across lanes with
simd_shuffle_xor, removing a threadgroup_barrier.",
"difficulty_1_to_5": 4,
"impacts": {
"memory_bandwidth": 8,
"register_pressure": 4, // risk: spill if Br > 8
"compute_hw_fit": 9 // SIMD width 32; keep tile 8x8
},
"overall_rating_0_to_10": 8,
"confidence_0_to_1": 0.7
}
Listing 1: Example K-Search world-model node. Each candidate optimization records a concrete action, estimated hardware impacts, an overall priority rating, and the model's confidence.
Figure 2: Overview of K-Search. The framework operates on a Search State $S_t$ structured as a search tree. The tree consists of Closed nodes (blue, visited states with attached program like $x_{12}$) and a Frontier of Open nodes (orange, pending hypotheses like $u_{13}$). The workflow iterates through three phases: (1) Action Selection, where the most promising action node is retrieved from the frontier based on world model estimated priority score $V$; (2) Local Refinement, where a stochastic policy $\pi_{\mathrm{code}}$ samples concrete implementations until stagnation; and (3) World Model Update, where the LLM reasons over the trajectory to update the search tree via Insert (adding new actions), Update (adjusting $V$, e.g., $u_{11}$ dropping from 0.9 to 0.6), and Prune (removing less promising nodes like $u_{10}$).
The original K-Search paper evaluated this search strategy on CUDA kernels from FlashInfer. Across GQA decode, MLA decode, MLA prefill, and MoE, K-Search improved more consistently than OpenEvolve and ShinkaEvolve over the same 120-iteration budget. These results establish the search framework we build on here; the remainder of this post asks whether its optimization knowledge can transfer beyond CUDA.
Figure 3: Main results from the original K-Search paper. Across three runs, K-Search achieves stronger best-so-far search scores, per-workload kernel performance, and speedup distributions than OpenEvolve and ShinkaEvolve on four FlashInfer CUDA kernels. Reproduced exactly from Cao et al. (2026).
Building an MLX backend
To bring K-Search to Apple Silicon, we first built a native MLX backend. We implemented a full MLX-specific task adapter for K-Search, including:
An MLX task backend in k_search/tasks/ handling kernel compilation and execution on Apple Silicon via MLX’s Metal/C++ APIs.
Updated kernel generator prompts for writing and modifying Metal/MLX kernels.
MLX-specific benchmarking integration using mlx.core measurement utilities.
Translating CUDA expertise to MLX
However, the more interesting challenge was not simply running K-Search on MLX. The key insight is that expert CUDA kernels encode decades of optimization knowledge that is transferable to Apple GPU if you can bridge the conceptual gap. Simply handing an LLM a CUDA kernel and asking it to port it is not enough: without deep hardware context, it produces code that is syntactically valid but architecturally wrong (wrong tile sizes, invalid primitives, mismatched memory assumptions).
Our translation layer consists of:
Concept mapping tables: A structured glossary of CUDA primitives and their MLX/Metal equivalents with hard constraints. For example:
__shared__ maps to Metal threadgroup memory but with a hard 32 KB limit (vs. NVIDIA’s 48 KB)
warp_reduce maps to MMA (preferred)
__syncthreads() becomes threadgroup_barrier(mem_flags::mem_tg)
H100’s ~3.35 TB/s HBM3 maps to M3 Max’s ~400 GB/s unified DRAM a bandwidth difference that reshapes which optimizations are worth pursuing.
MLX-specific hints and patterns: Concrete code-level patterns for operations with no direct CUDA equivalent, such as register-based row reductions using simd_shuffle_xor in an 8×8 MMA tile layout, or the “exp2 trick” (replacing $exp(x)$ with $exp_2(x \log_2 e)$) for faster softmax on Apple’s fast $exp_2$ hardware instruction.
Reusable assertions: Expert kernel behaviors reframed as properties the evolutionary search must preserve, rather than code to copy.
Matching expert kernel performance: the Attention kernel
We evaluate three configurations of an MLX attention kernel for Apple Silicon: (1) a naive baseline, (2) pure evolution with no additional provided context, and (3) a full context translation layer, which supplies the optimizer with architecture-specific implementation knowledge extracted from high-performance kernels (e.g., FlashAttention-2), letting the evolutionary search reason about implementation strategies rather than starting from a naive kernel. Together, these three configurations let us isolate the exact impact of the translation layer.
Figure 4: Performance scaling of the Attention Kernel through stacked optimizations. The "Full Context" configuration successfully discovers and implements advanced strategies like double buffering and loop unrolling, achieving near-expert performance.
The jump from 0.26× to 0.97× the speed of Apple’s state-of-the-art attention kernel — illustrates how much the translation layer matters. With full context, the evolved kernel independently discovers the key optimizations in FlashAttention 2: threadgroup memory tiling, online softmax, K-transposition for memory access, and the exp2 trick. The last of these replaces every softmax exponential with a base-2 exponential,
\[e^x = 2^{x \log_2 e},\]
which is exact and lets the kernel use Apple’s fast fast::exp2() hardware instruction directly instead of paying for a base conversion at runtime.
A 20× faster prefill: the Mamba SSM kernel
To evaluate whether K-Search generalizes beyond attention kernels, we applied it to the state-space model (SSM) kernel used by Mamba. Unlike attention, the computational bottleneck is a recurrent state update rather than a softmax, providing a substantially different optimization challenge. We compare the evolved implementation against the community MLX implementation (mlx-lm) and the PyTorch reference implementation (mamba.py) on an M1 Max.
Evaluated on mamba-370m f16, M1 Max 64GB:
Metric
mlx-mamba (ours)
mlx-lm (community)
mamba.py
Decode
152 tok/s
116 tok/s
40 tok/s
Prefill L=512
5,751 tok/s
329 tok/s
1,089 tok/s
Prefill L=1024
6,010 tok/s
327 tok/s
1,127 tok/s
Prefill L=2048
6,612 tok/s
326 tok/s
1,092 tok/s
Prefill L=4096
6,743 tok/s
339 tok/s
1,042 tok/s
Table 1: Prefill and decode throughput on mamba-370m (f16, M1 Max 64GB). mlx-mamba (ours) reaches ~20× higher prefill throughput than the community mlx-lm baseline, while decode remains comparable.
The ~20× prefill speedup over mlx-lm comes down to one difference: mlx-lm does not implement a parallel scan for the SSM. The state recurrence
\[h_t = \bar{a}_t h_{t-1} + \bar{b}_t\]
looks inherently sequential, but each step can be written as a pair $(\bar{a}_t, \bar{b}_t)$ under the associative combine
\[(a_2, b_2) \circ (a_1, b_1) = \left(a_2 a_1,\ a_2 b_1 + b_2\right),\]
which reproduces the recurrence exactly. Because the operator is associative, the whole sequence can be evaluated with a parallel (prefix) scan in $O(\log N)$ dependent steps instead of $O(N)$. mlx-lm skips this and processes tokens one at a time, leaving most of Apple Silicon’s compute idle; our evolved Metal kernel applies the scan and makes much fuller use of GPU throughput. The gain shows up in prefill, where the full sequence is available to scan in parallel, and not in single-token decode, where there is only one new token per step and no scan to parallelize — which is why the decode row is roughly flat while prefill is ~20×.
mamba.py is slow on both prefill and decode because it is a PyTorch reference implementation that falls back to CPU or MPS on Apple Silicon, forgoing the hardware-specific optimizations that MLX’s Metal backend makes possible.
What’s next?
On the two kernels we studied, AI-driven evolutionary kernel search grounded in structured cross-platform translation knowledge reached near-expert performance on Apple Silicon without a team of GPU experts starting from scratch. We do not yet know how far this generalizes, but the result is encouraging.
For us the main takeaway is that the bottleneck was not the LLM’s ability to write Metal code, but the quality of the context and constraints we gave it. Our CUDA translation layer converts existing NVIDIA kernel expertise into actionable guidance for Apple Silicon, and lets K-Search’s evolutionary search do the rest.
We are actively extending this work in several directions: supporting new architectures, with current efforts focused on developing new kernels for the IBM Spyre AIU and broader hardware targets; adding more kernels such as paged attention and fused MoE routing; and improving integration with the K-Search evolution loop to make translation context even more automatic.
Acknowledgements
This work was carried out by IBM Research and builds on K-Search from the UC Berkeley Sky Lab (Cao et al., 2026). We welcome collaboration and feedback from the MLX and broader AI systems communities. If you are working on kernel optimization for non-CUDA hardware, we would love to hear from you.
Citation
@article{cao2026k,
title={K-Search: LLM Kernel Generation via Co-Evolving Intrinsic World Model},
author={Cao, Shiyi and Mao, Ziming and Gonzalez, Joseph E and Stoica, Ion},
journal={arXiv preprint arXiv:2602.19128},
year={2026}
}
Appendix: Try it yourself
The MLX backend is built on top of the open-source K-Search repo, so the results here can be reproduced directly. The steps are:
1. Clone and install
git clone https://github.com/caoshiyi/K-Search.git
cd K-Search
uv pip install openai wandb
uv pip install git+https://github.com/caoshiyi/flashinfer-bench-ksearch.git
2. Set your credentials
Open the relevant script under scripts/ and set three variables at the top:
KSEARCH_ROOT=/path/to/K-Search
API_KEY=your-llm-api-key
3. Run kernel search
# Optimize Flash Attention on Apple Silicon (world-model mode)
bash scripts/mac_flash_attention_wm.sh
# Or a Mamba SSM kernel, e.g. the selective scan
bash scripts/mamba_selective_scan_fwd_wm.sh
Full CLI reference and documentation are in the README.
[2026-07-29]
Teaching LLMs to Update Beliefs for Efficient Long-Horizon Interaction
.abbel-fig {
display: block;
text-align: center;
margin: 2.4em 0;
line-height: 1.4;
max-width: 100%;
}
.abbel-fig img {
display: block;
margin: 0.65em auto 0;
height: auto;
max-width: 100%;
}
/* Image sizes; captions use a narrower measure below */
.abbel-fig--wide img { width: 100%; max-width: 100%; }
.abbel-fig--wide-90 img { width: 100%; max-width: 90%; }
.abbel-fig--wide-lg img { width: 100%; max-width: 100%; }
.abbel-fig--chart img { width: 100%; max-width: 82%; }
.abbel-fig--chart-sm img { width: 100%; max-width: 64%; }
.abbel-fig--portrait img { width: 50%; max-width: 520px; }
.abbel-fig--equation img { width: 100%; max-width: 52%; }
.abbel-fig--video {
width: 110%;
max-width: 110%;
margin-left: -5%;
margin-right: -5%;
box-sizing: border-box;
}
.abbel-fig--video .abbel-frames {
max-width: 100%;
width: 100%;
}
.abbel-frames {
margin: 0.65em auto 0;
max-width: 100%;
user-select: none;
}
.abbel-frames__stage {
position: relative;
cursor: pointer;
border: none;
background: transparent;
line-height: 0;
width: 100%;
}
.abbel-frames__stage img {
width: 100%;
height: auto;
display: block;
}
.abbel-frames__hint {
position: absolute;
right: 0.55em;
bottom: 0.55em;
background: rgba(0,0,0,0.4);
color: #fff;
font-size: 0.68em;
font-style: normal;
padding: 0.18em 0.5em;
border-radius: 3px;
pointer-events: none;
opacity: 0;
transition: opacity 0.2s ease;
}
.abbel-frames__stage:hover .abbel-frames__hint,
.abbel-frames.is-paused .abbel-frames__hint {
opacity: 1;
}
.abbel-frames.is-playing .abbel-frames__hint {
opacity: 0;
}
.abbel-frames__controls {
display: flex;
align-items: center;
justify-content: center;
gap: 0.55em;
margin-top: 0.35em;
flex-wrap: wrap;
}
.abbel-frames__controls button {
appearance: none;
border: none;
background: transparent;
color: #999;
font: inherit;
font-size: 0.78em;
padding: 0.15em 0.35em;
border-radius: 2px;
cursor: pointer;
}
.abbel-frames__controls button:hover {
color: #666;
background: transparent;
}
.abbel-frames__controls button.abbel-frames__next,
.abbel-frames__controls button.abbel-frames__prev {
color: #bbb;
font-weight: 400;
}
.abbel-frames__controls button.abbel-frames__next:hover,
.abbel-frames__controls button.abbel-frames__prev:hover {
color: #999;
}
.abbel-frames__controls button#abbel-frames-play {
color: #777;
letter-spacing: 0.02em;
}
.abbel-frames__controls button[aria-pressed="true"] {
background: transparent;
color: #555;
border-color: transparent;
}
.abbel-frames__meta {
font-size: 0.72em;
color: #bbb;
font-variant-numeric: tabular-nums;
min-width: 4em;
text-align: center;
}
.abbel-frames__dots {
display: flex;
justify-content: center;
gap: 0.25em;
margin-top: 0.25em;
flex-wrap: wrap;
}
.abbel-frames__dots button {
appearance: none;
width: 0.4em;
height: 0.4em;
padding: 0;
border-radius: 50%;
border: 1px solid #ccc;
background: #fff;
cursor: pointer;
}
.abbel-frames__dots button[aria-current="true"] {
background: #aaa;
border-color: #aaa;
}
.abbel-fig .abbel-fig-cap,
i.abbel-fig-cap {
display: block;
text-align: center;
font-style: italic;
color: #444;
margin: 2.7em auto 0.15em;
max-width: 38em;
width: 100%;
box-sizing: border-box;
padding: 0 0.5em;
font-size: 0.8rem;
line-height: 1.4;
}
.abbel-fig .abbel-fig-cap sub,
.abbel-fig .abbel-fig-cap sup,
i.abbel-fig-cap sub,
i.abbel-fig-cap sup {
font-size: 0.75em;
line-height: 0;
}
.abbel-fig--tight .abbel-fig-cap {
margin-top: 0.9em; /* ~1/3 of default figure→caption gap */
}
.abbel-fig--equation .abbel-fig-cap {
margin-top: 1.35em; /* half of default 2.7em figure→caption gap */
}
.abbel-fig--chart-sm .abbel-fig-cap {
margin-top: 1.35em; /* one line less than default 2.7em */
}
@media screen and (max-width: 40em) {
.abbel-fig--wide img,
.abbel-fig--wide-lg img,
.abbel-fig--wide-90 img,
.abbel-fig--chart img,
.abbel-fig--chart-sm img { max-width: 100%; }
.abbel-fig--equation img { max-width: 75%; }
.abbel-fig--portrait img { max-width: 50%; }
.abbel-fig--video {
width: 100%;
max-width: 100%;
margin-left: 0;
margin-right: 0;
}
.abbel-fig--video .abbel-frames { max-width: 100%; }
.abbel-fig .abbel-fig-cap { max-width: 100%; }
}
.abbel-footnotes {
font-size: 0.8em;
color: #888;
font-style: italic;
margin: 1.5em 0;
}
.abbel-footnotes ol {
padding-left: 1.25em;
margin: 0.4em 0 0;
}
.abbel-footnotes li {
margin: 0.55em 0;
}
.abbel-footnotes p {
margin: 0.2em 0;
}
.abbel-footnotes a {
color: #888;
}
.abbel-table-wrap { overflow-x: auto; margin: 0.5em auto 0; text-align: center; }
.abbel-fig--table { margin: 2.4em 0; }
.abbel-fig--table .abbel-table-wrap { margin: 0.65em auto 0; }
.abbel-table {
width: 100%;
max-width: 560px;
margin: 0 auto;
border-collapse: collapse;
font-size: 0.88em;
line-height: 1.35;
}
.abbel-table th,
.abbel-table td {
padding: 0.5em 0.7em;
border-bottom: 1px solid #ddd;
text-align: center;
vertical-align: middle;
}
.abbel-table th {
border-bottom: 2px solid #333;
font-weight: 600;
}
.abbel-table th:first-child,
.abbel-table td:first-child { text-align: left; }
.abbel-table tr.abbel-baseline td { color: #888; font-style: italic; }
.abbel-ack a {
color: #1565c0;
font-weight: 500;
text-decoration: none;
border-bottom: 1px solid #90caf9;
padding-bottom: 0.06em;
}
.abbel-ack a:hover {
color: #0d47a1;
border-bottom-color: #1565c0;
}
/* Suppress "View on alphaXiv" badges/tags (browser extension / userscript injectors) */
a[href*="alphaxiv.org"],
a[href*="alphaXiv"],
[class*="alphaxiv"],
[class*="alphaXiv"],
[class*="AlphaXiv"],
[id*="alphaxiv"],
[id*="alphaXiv"],
[data-alphaxiv],
[data-alpha-xiv],
img[src*="alphaxiv"],
img[alt*="alphaXiv" i],
img[alt*="alphaxiv" i],
button[aria-label*="alphaXiv" i],
a[title*="alphaXiv" i],
a[aria-label*="alphaXiv" i],
span[title*="alphaXiv" i] {
display: none !important;
visibility: hidden !important;
width: 0 !important;
height: 0 !important;
overflow: hidden !important;
pointer-events: none !important;
position: absolute !important;
left: -9999px !important;
}
/* Section / subsection spacing (title → body, and gap before next section) */
.post-content > h2 {
margin-top: 2.6em;
margin-bottom: 0.75em;
}
.post-content > h2:first-of-type {
margin-top: 1.6em;
}
.post-content > h3 {
margin-top: 1.85em;
margin-bottom: 0.6em;
}
Overview of ABBEL compared to traditional recursive summarization. Beliefs replace the full interaction history as the agent’s working context, and belief grading improves performance by
supervising the contents of each belief state..
As task horizons grow, LLM contexts can’t scale forever. Self-summarization enables concise, interpretable contexts, but at a significant performance cost, especially for human assistance domains where high quality data is scarce, e.g., collaborative code generation. We address this with ABBEL: a framework that isolates and supervises the information content of summaries in the form of natural-language belief states.
Motivation: the cost of recursive summarization
For language models to effectively assist with increasingly complex tasks such as software development, they must be able to interact with us over hundreds or even thousands of steps. For such long tasks, it is impractical to keep the history of the entire interaction in context. The heuristic approach used so far has been summary generation, sometimes called context compaction. For example, Cursor’s latest model composer 2.5 uses compaction during training for improved performance (Cassano et al., 2026). Alongside composer, Grandcode (DeepReinforce et al., 2026), the first system to consistently beat all human competitors in online coding competitions, despite using one of the newest efficient attention models (Qwen 3.5-397B),1 still found it necessary to employ context summarization.
But compaction has a problem. Despite seemingly low performance gaps in benchmarks, model servers like Cursor continue to recommend that users avoid compaction with their coding assistants in the middle of a task (Heule et al., 2026).
To understand why, see below the performance over RL fine-tuning of a Context summary model compared to full context models in Combination Lock, a Wordle-like game that allows up to 16 guesses.2 Though both model types improve over the course of training, the summary model never closes the gap.
Fig. 1: Average attempts to guess the target word on Combination Lock over RL fine-tuning (lower is better). Context-summary policies improve with training but do not close the gap to full-context policies.
Making models self-summarize while completing a task increases the complexity of the learning problem. While this could typically be addressed by training with more data, the performance degradation observed in real world interactive settings likely arises from the difficulty we have in creating and using human simulators effectively to generate high quality training environments (Lin et al., 2025, Tomlin et al., 2025). Thus, the better you can learn to summarize on the limited and messy multiturn interaction trajectories you can collect, the better off your model will be for downstream users.
ABBEL: acting through belief bottlenecks
Fig. 2: Autoencoder-inspired belief grading. The model encodes prior belief, action and observation (bt, at, ot) into posterior belief bt+1 and is rewarded for how well select information from the history can be reconstructed from that belief.
To address poor learning efficiency, we isolate the summary generation task. Drawing inspiration from recursive Bayesian estimation, we formulate summaries as belief states, which we periodically prompt the model to update based on new information.3
‹ Prev
Pause
Next ›
1 / 16
Fig. 3: ABBEL rollout. Belief updates from the latest observation alternate with action selection conditioned only on the current posterior belief.
Belief grading
We then extract and supervise the contents of the belief states (Fig. 2, Belief Grading). Belief grading can be thought of as adding an auxiliary RL task, using heuristics designed to capture what makes a good belief as the reward. An example heuristic for coding could be shorter is better, but closer to being able to reconstruct the git diff is also better, so balancing these would yield a good belief. In domains where good heuristics are hard to define, we propose a general autoencoding-inspired grading function, which treats the current language model πθ as both encoder and decoder of information from
the history, and the belief states as the codes. We grade each belief bt+1 by how well it can be used by the current model πθ to reconstruct the most recent observation ot:
Eq. 1: Reconstruction grading objective. Here bt+1 is the updated belief, ot the latest observation, at the action just taken, bt the prior belief, pI the task prompt, and πθ the current model. Higher grades reward beliefs that retain information needed to decode the latest observation.
What do we gain by grading beliefs?
Collaborative coding on CollabBench
We demonstrate the utility of belief grading in our motivating domain of human-driven assistive coding, with the CollabBench environment from Sweet-RL (Zhou et al., 2025).
Fig. 4: CollabBench collaborative coding environment. The agent asks clarifying questions, then submits a function scored against hidden unit tests.
We see that with the general reconstruction-based belief grading function we reduce the performance gap from full context models by about 50%, and train in 50% fewer steps compared to training models to summarize without belief grading (no BG). After training, ABBEL still uses significantly less memory than the full context setting, as measured by the peak context token length (Peak Tokens).
Model
Test Pass Rate ↑
Success Rate ↑
Peak Tokens × 10² ↓
Training Steps ↓
Full Context
0.52±0.02
0.39±0.02
14.08±0.55
100
ABBEL (no BG)
0.46±0.02
0.31±0.02
4.20±0.37
100
ABBEL-rec-BG
0.48±0.01
0.36±0.01
6.01±0.33
50
Fig. 5: CollabBench results. With reconstruction belief grading, ABBEL-rec-BG recovers about half the gap to full context while using fewer peak tokens, and trains in 50 steps instead of 100.
Combination Lock
Additionally, in CombinationLock, we demonstrate that ABBEL with a belief grader which leverages domain knowledge (by computing useful statistics over the history and checking that they can be reconstructed from the belief state), enables even higher learning efficiency than full context (FULL CTX) models.
Fig. 6: Average attempts to guess the target word on Combination Lock (lower is better). With domain-knowledge belief grading, ABBEL approaches or exceeds FULL CTX in this setting; without belief grading, learning is slower.
Multi-objective question answering
In a third environment, multi-objective question answering (from MEM1 Zhang et al., 2025, a recent work which performed end-to-end optimization in a modified version of typical recursive summarization), we demonstrate the utility of isolating belief states from reasoning, by showing that a Peak Belief length Penalty (more details in paper) significantly reduces memory usage with minimal performance degradation, unlike is commonly observed when penalizing reasoning lengths (Arora et al., 2025).
Fig. 7: Exact-match score and peak memory versus number of objectives in multi-objective QA. ABBEL with a peak belief penalty (PBP) maintains comparable performance while using less memory than MEM1 and ABBEL without PBP in this evaluation.
Related work
Alternative solutions to managing long contexts involve different tradeoffs, and are worth considering depending on the requirements of a deployed system. Context compression methods generate dense representations which, while computationally efficient, sacrifice human-understandability (Kontonis et al., 2026, Eyuboglu et al., 2025, Gupta et al., 2025, Chevalier et al., 2023, Deng et al., 2025, Deng et al., 2025, Bulatov et al., 2022). Hand-designed summarization prompts (Wang et al., 2025, Örwall et al., 2025, Starace et al., 2025) and pruning strategies (Jiang et al., 2024) specific to target environments require expert human knowledge and don’t allow an agent to learn what to remember as part of its decision-making strategy. Methods that process long contexts into an external memory store (Packer et al., 2023, Xu et al., 2025) for the agents or subagents to query (Zhang et al., 2025) are complementary, as they may benefit from better next context creation through summarization training. We would like to point out some exciting works in the space of general recursive summarization focused on math (Wu et al., 2026), reasoning with belief generation (Zhou et al., 2025), competitive coding with a distilled summarization module using similar autoencoding objectives to our general belief grader (DeepReinforce et al., 2026), and adding continuous features to summaries (Kontonis et al., 2026).
What’s next for better memory?
Many more possibilities are enabled through using explicit belief states as information bottlenecks for multi-step interaction. You could reward actions based on their effect on the belief state to guide exploration, transmit the explicit belief states for better communication between agents, or even improve user controllability by directly modifying the memories on which the agents’ decisions are based.
Some forms of information, e.g., what a person looks like, are not represented well by text alone. A continuously learning system will also have to capture such information. Additionally, if we want a system to learn to communicate in a brand new language or to play a brand new game better than any person in the world, the skills accumulated over the lifetime of conversations or games must be stored in a very compressed form, essentially taking on the role of the weights of the model itself.
More powerful systems will likely utilize a combination of multiple forms of memory, where the contents of the context may correspond to working memory while other approaches are used for short and long-term memory. How to instantiate these other forms of memory, for instance via test-time training, adapter memories, continuous context memories, or some combination thereof, presents an exciting challenge.
Acknowledgements
Acknowledgements: We would like to thank Alane Suhr and Kartik Goyal for advising this research as well as Ethan Mendes, David He, Jitesh Jain, and Nicholas Tomlin for comments on early drafts of this post. We would like to thank the MEM1 authors for their email correspondence and for sharing private reviewer feedback which we found particularly insightful.
Citation
If abbel was inspiring for your future work, please cite us with this! And here is some advice for doing similar research!
@misc{lidayan2026abbellearningnaturallanguagebelief,
title={ABBEL: Learning Natural-Language Belief States for Memory-Efficient Interaction},
author={Aly Lidayan and Jakob Bjorner and Satvik Golechha and Kartik Goyal and Alane Suhr},
year={2026},
eprint={2512.20111},
archivePrefix={arXiv},
primaryClass={cs.CL},
url={https://arxiv.org/abs/2512.20111},
}
With newer models the number of tokens till 50% compute spend is on attention gets much larger than 25K. Interleaving linear attention alternatives with full attention as is done with gpt-oss and DeepSeekv4, results in massive flops reductions for the attention computation. For example with DeepSeekv4-Pro (1.6T A49B) it requires nearly 450 thousand tokens to reach the 50% tradeoff point. Grandcode uses Qwen-3.5-397B-A17B a model which hits 50% FLOPs for attention at ~150 thousand tokens.
↩
This setting is technically solvable with much more computationally effective tools, but serves as a flexible test bed to study properties of recursive summarization. Bertsimas et al., 2022, showed that an exact solution for the wordle game instantiated with the original vocabulary of the javascript game can be found with dynamic programming, but evidently the general formulation of wordle as a guessing game on K letters with L attempts and some dictionary of valid words and correct words D is NP hard to determine the minimal number of moves required.
↩
In practice there is an O(N/K) overhead cost for summary. N is the total number of actions. K is the number of actions till summarization is triggered. This is necessarily true for any summary approach. For ease of illustration this gif uses K = 1. In our experiments, to put more emphasis on summarization weaknesses we also use K=1. In practice overhead is small as K can be chosen to be near the efficient hardware limit.
↩
(function () {
var root = document.getElementById('abbel-frames');
if (!root) return;
var count = parseInt(root.getAttribute('data-frame-count') || '15', 10);
var prefix = root.getAttribute('data-frame-prefix') || 'https://bair.berkeley.edu/static/blog/abbel/frames/frame_';
var intervalMs = parseInt(root.getAttribute('data-interval') || '1300', 10);
var img = document.getElementById('abbel-frames-img');
var meta = document.getElementById('abbel-frames-meta');
var dots = document.getElementById('abbel-frames-dots');
var btnPrev = document.getElementById('abbel-frames-prev');
var btnNext = document.getElementById('abbel-frames-next');
var btnPlay = document.getElementById('abbel-frames-play');
var stage = root.querySelector('.abbel-frames__stage');
var hint = root.querySelector('.abbel-frames__hint');
var i = 0;
var playing = false;
var timer = null;
var urls = [];
for (var n = 0; n < count; n++) {
urls.push(prefix + (n < 10 ? '0' + n : '' + n) + '.png');
}
urls.forEach(function (u) { var im = new Image(); im.src = u; });
if (dots) {
for (var d = 0; d < count; d++) {
(function (idx) {
var b = document.createElement('button');
b.type = 'button';
b.setAttribute('aria-label', 'Go to frame ' + (idx + 1));
b.addEventListener('click', function (e) {
e.stopPropagation();
stop();
show(idx);
});
dots.appendChild(b);
})(d);
}
}
function setHint(text) {
if (hint) hint.textContent = text;
}
function show(idx) {
i = (idx + count) % count;
img.src = urls[i];
if (meta) meta.textContent = (i + 1) + ' / ' + count;
if (dots) {
var kids = dots.children;
for (var k = 0; k < kids.length; k++) {
if (k === i) kids[k].setAttribute('aria-current', 'true');
else kids[k].removeAttribute('aria-current');
}
}
}
function stop() {
playing = false;
if (timer) { clearInterval(timer); timer = null; }
root.classList.remove('is-playing');
root.classList.add('is-paused');
if (btnPlay) {
btnPlay.textContent = 'Play';
btnPlay.setAttribute('aria-pressed', 'false');
}
//setHint('Click to advance'); not pretty text, and functionality is clear already
}
function start() {
playing = true;
root.classList.add('is-playing');
root.classList.remove('is-paused');
if (btnPlay) {
btnPlay.textContent = 'Pause';
btnPlay.setAttribute('aria-pressed', 'true');
}
//setHint('Click to pause');
if (timer) clearInterval(timer);
timer = setInterval(function () { show(i + 1); }, intervalMs);
}
function advance() {
if (playing) { stop(); return; }
stop();
show(i + 1);
}
function back() { stop(); show(i - 1); }
if (stage) {
stage.addEventListener('click', advance);
stage.addEventListener('keydown', function (e) {
if (e.key === 'Enter' || e.key === ' ') { e.preventDefault(); advance(); }
if (e.key === 'ArrowRight') { e.preventDefault(); advance(); }
if (e.key === 'ArrowLeft') { e.preventDefault(); back(); }
});
}
if (btnNext) btnNext.addEventListener('click', function () { stop(); show(i + 1); });
if (btnPrev) btnPrev.addEventListener('click', back);
if (btnPlay) btnPlay.addEventListener('click', function () {
if (playing) stop(); else start();
});
show(0);
// Autoplay like a GIF: start once the first frame is ready, and again when scrolled into view
function kickoff() {
if (!playing) start();
}
if (img.complete) kickoff();
else img.addEventListener('load', kickoff, { once: true });
if ('IntersectionObserver' in window) {
var io = new IntersectionObserver(function (entries) {
entries.forEach(function (e) {
if (e.isIntersecting) kickoff();
});
}, { threshold: 0.35 });
io.observe(root);
}
})();
[2026-07-26]
