autoresearch-mamba is that test, adapted to Mamba-2, Mamba-3, and hybrid Mamba–Transformer MoE models: can an agent improve a Mamba training recipe under a fixed compute budget and an evaluator it cannot rewrite? The agent edits a constrained training surface, runs for five minutes, measures val_bpb, and keeps only experiments that improve the metric. Everything else is discarded.
This isn't meant as a general training platform. It's a small experimental system designed to make agent-driven model research measurable, repeatable, and difficult to game. MLX is the first-class path so the complete loop can run on Apple Silicon; a secondary PyTorch/CUDA path keeps the same ideas portable to GPUs.
Why Mamba
Early autoresearch experiments naturally centered on GPT and nanoGPT-style architectures. Mamba changes the search space: recurrence design, state dimension, discretization, convolution width, hybrid block placement, and MoE routing all become variables an agent can test under the same fixed budget.
That makes the project useful for more than reproducing a GPT loop with a different model class. It tests whether the autoresearch pattern survives when the architecture has different bottlenecks, different failure modes, and a different relationship between training throughput and inference efficiency.
What the Repo Actually Is
autoresearch-mamba is deliberately small. The structure mirrors the same discipline that makes Karpathy's setup useful:
prepare_mlx_mamba_3.py: fixed architecture-aware MLX prep entry pointtrain_mamba_3_mlx.py: the main architecture-aware MLX training surface for bothmamba-2andmamba-3train_hybrid_moe_mlx.py: hybrid Mamba–Transformer MoE training surface (Nemotron-H style)program.md: instructions for the agent loopprepare.pyandtrain_mamba.py: a secondary PyTorch/CUDA pathanalysis.ipynb: a notebook for analyzingresults.tsvand plotting progress over time
The canonical path is Apple Silicon + MLX. The CUDA/PyTorch path exists as a secondary reference implementation.
Architecturally, the repo now supports Mamba-2, Mamba-3, and hybrid Mamba–Transformer MoE through the MLX entry paths. The Mamba-3 side is aimed at training-oriented autoresearch rather than reproducing every upstream fused kernel implementation detail from mamba_ssm. The hybrid side adapts the Nemotron-H block-level hybridization pattern for single-file MLX training.
The point is to keep a compact, agent-friendly research harness that preserves the core block logic while staying small enough to iterate on.
Hybrid Mamba–Transformer MoE
The latest addition to the repo is a hybrid architecture that interleaves Mamba SSM layers with self-attention and Mixture-of-Experts, following NVIDIA's Nemotron-H design.
The key insight from Nemotron-H, block-level hybridization, is that instead of each layer being a full transformer block (attention + FFN), every layer is one standalone residual block of a specific type. A pattern string like ME*EM encodes the entire architecture:
M= Mamba SSM block (sequence mixing via state space dynamics)E= MoE FFN block (per-token transformation via routed experts)*= Self-Attention block (global sequence mixing)-= Dense MLP block
Mamba layers handle efficient sequence processing. MoE layers provide per-token capacity with sparse activation. Attention layers add global context at key positions. The ME pair functions like one logical transformer layer: Mamba does the mixing, MoE does the feed-forward. But because each is a separate residual block, the pattern is fully composable.
The implementation (train_hybrid_moe_mlx.py) is a single-file, self-contained MLX script. It includes CausalSelfAttention with grouped-query attention and RoPE, MoELayer with top-k softmax routing and Switch Transformer-style load-balancing loss, and HybridResidualBlock which wraps any block type in a uniform pre-norm + residual interface. The Mamba blocks are reused directly from the existing Mamba-2/Mamba-3 code.
With the local preset (ME*EM, 4 experts, top-2, d_model=256), the model trains at ~29K tokens/sec on Apple Silicon, producing a 4.7M param model (3.9M active per token) with a val_bpb of 2.059 in five minutes.
The Autoresearch Loop
The loop is intentionally simple:
- Prepare data and tokenizer once.
- Run training for a fixed
TIME_BUDGET = 300seconds. - Evaluate using
val_bpb. - Keep the change if
val_bpbimproved. - Otherwise discard it and move on.
That five-minute budget matters. It forces the question to become: what is the best model and training recipe an agent can discover on this hardware in this exact amount of time?
It also keeps the search honest. The agent cannot win by quietly changing the evaluator, expanding the benchmark, or redefining success after the fact. It has one target: lower val_bpb under the fixed harness.
Why val_bpb
The metric here is validation bits per byte (val_bpb). Lower is better.
I like this choice for the same reason Karpathy did: it stays tied to language modeling performance without becoming too tokenization-specific, and it isn't tied to GPT or to transformers either. It works just as well for a Mamba autoregressive language model.
val_bpb is partly a hardware measurement: in a fixed five-minute autoresearch loop, the final number isn't just about the model. It's also about how much optimization progress your hardware can fit into the time budget, so comparisons are meaningful only when the setup is held constant: same evaluator, same preset, same platform class.
Making it run on a MacBook
One of the practical constraints here was that I wanted this to work not only as a theoretical repo, but as something I could actually test on Apple Silicon.
So the repo has two practical operating modes:
- the built-in default MLX path, which can target either
mamba-2ormamba-3 - an optional local preset path for smaller-memory local testing
That local preset workflow turned out to be important. It let me validate the loop end-to-end on constrained Apple Silicon while keeping the default architecture-aware configuration intact as the canonical tracked setup.
MLX-specific detail, briefly: the canonical local path is prepare_mlx_mamba_3.py + train_mamba_3_mlx.py, the same path can switch between mamba-2 and mamba-3, and the smaller local preset flow makes Apple Silicon smoke testing practical. In the current state of the repo, Mamba-3 SISO runs end-to-end locally, while larger Mamba-3 MIMO runs are mainly bounded by Metal memory budget rather than by a missing MLX code path.
So the repo is not just "MLX-compatible" in theory. It's set up so the autoresearch loop can run locally, collect results.tsv, and generate a progress plot without changing the underlying philosophy.
What I Wanted to Preserve from Karpathy's Repo
There are a lot of ways to turn "autonomous research" into fluff. What I wanted to preserve from Karpathy's setup wasn't the vibe, it was the discipline:
- one fixed evaluator
- one main mutable training file
- one fixed time budget
- one ground-truth metric
- one keep/discard loop
That design is what makes the repo useful. Without those constraints, you don't really have autoresearch, just an agent editing files until the story sounds good.
What Ships Today
The release includes:
- an architecture-aware MLX autoresearch path for
mamba-2andmamba-3 - a hybrid Mamba–Transformer MoE training path (Nemotron-H style,
train_hybrid_moe_mlx.py) - a training-oriented MLX Mamba-3 path with SISO and MIMO branches
- the secondary PyTorch/CUDA reference path
- a documented README and
program.mdcovering all three architectures - local MLX preset workflows for smaller Apple Silicon testing (pure Mamba and hybrid)
- an analysis notebook for plotting progress from
results.tsv
What the loop hasn't proven yet
The current release covers Mamba-2, Mamba-3, and hybrid Mamba–Transformer MoE models in one measurable loop. The more important contribution is the experimental contract around them: one mutable surface, one fixed budget, one evaluator, and an auditable keep-or-discard history.
The next useful result isn't another supported architecture. It's evidence that repeated agent runs discover improvements a careful human baseline would miss, and that those improvements survive replication on the same hardware.