autoresearch-mamba: Karpathy-Style Autoresearch for Mamba-2, Mamba-3, and Hybrid Mamba–Transformer MoE

autoresearch-mamba is a compact research harness for asking one concrete question: can an agent improve a Mamba training recipe under a fixed compute budget and an evaluator it cannot rewrite?

This is intentionally not a general training platform. It is 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 point
• train_mamba_3_mlx.py: the main architecture-aware MLX training surface for both mamba-2 and mamba-3
• train_hybrid_moe_mlx.py: hybrid Mamba-Transformer MoE training surface (Nemotron-H style)
• program.md: instructions for the agent loop
• prepare.py and train_mamba.py: a secondary PyTorch/CUDA path
• analysis.ipynb: a notebook for analyzing results.tsv and 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 is block-level hybridization: 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:

1. Prepare data and tokenizer once.
2. Run training for a fixed TIME_BUDGET = 300 seconds.
3. Evaluate using val_bpb.
4. Keep the change if val_bpb improved.
5. 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. It is not a GPT metric. It is not a transformer metric. It works just as well for a Mamba autoregressive language model.

One subtle but important point: in a fixed five-minute autoresearch loop, the final val_bpb is not just about the model. It is also about how much optimization progress your hardware can fit into the time budget. So comparisons are meaningful when the setup is held constant: same evaluator, same preset, same platform class.

MLX, GPU, and the Local Preset

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-2 or mamba-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.

In other words: the repo is not just "MLX-compatible" in theory. It is set up so the autoresearch loop can actually 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. The thing I wanted to preserve from Karpathy's setup was not 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 do not really have autoresearch. You just have an agent editing files until the story sounds good.

What Ships Today

The release includes:

• an architecture-aware MLX autoresearch path for mamba-2 and mamba-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.md covering 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

Wrapping up

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 is not another supported architecture. It is evidence that repeated agent runs discover improvements a careful human baseline would miss—and that those improvements survive replication on the same hardware.