Watch two AI agents collaborate in real time. The Coder writes code, the sandbox runs it, and the Sheriff reviews the output. If something breaks, the error goes back to the Coder.

What happens when you put two AI agents in a room, give them a coding task, and let them argue until the code works?

That is the premise behind this project. Agent A (the Coder) writes Python scripts. Agent B (the Sheriff) runs them, reads the stack traces, and tells the Coder what it got wrong. They go back and forth until the code passes or a hard limit kicks in, following the iterative self-refinement pattern (Madaan et al., 2023) with execution-grounded feedback (Shinn et al., 2023). The catch is that both agents run on a single MacBook using Qwen2.5-Coder-7B-4bit (Hui et al., 2024). No OpenAI key, no cloud GPU, no data ever leaving the machine.

I built this to answer a question that kept nagging me: can two small local models actually collaborate to produce working code, the same way cloud-based agents do with GPT-4? The answer, it turns out, is surprisingly yes. With some interesting caveats about what small models can and cannot self-correct.

The code is at github.com/thenamangoyal/multi-agent-mlx (Goyal, 2026).

Results at a Glance

I ran three scenarios of increasing difficulty, all on an M1 Pro with 16 GB of RAM:

Total wall time: 3 minutes 50 seconds. Total cost: $0.00.

{
  "responsive": true,
  "tooltip": {"trigger": "axis"},
  "xAxis": {
    "type": "category",
    "data": ["Calendar", "CSV Detective", "Neural Net"],
    "axisLabel": {"fontSize": 13}
  },
  "yAxis": {
    "type": "value",
    "name": "Seconds",
    "nameTextStyle": {"fontSize": 13}
  },
  "series": [
    {
      "type": "bar",
      "data": [
        {"value": 22.0, "itemStyle": {"color": "#00CC96"}},
        {"value": 158.8, "itemStyle": {"color": "#636EFA"}},
        {"value": 49.6, "itemStyle": {"color": "#EF553B"}}
      ],
      "barWidth": "50%",
      "label": {
        "show": true,
        "position": "top",
        "formatter": ["{b|1 attempt}", "{b|3 attempts}", "{b|2 attempts}"],
        "rich": {"b": {"fontSize": 12}}
      }
    }
  ]
}

Every scenario eventually produced correct, running code. The interesting part is not that they succeeded. It is how they failed first and then fixed themselves.

The Architecture

The key constraint is 16 GB of RAM. Two separate 7B models would blow past that. The solution: both agents share a single mlx_lm.server instance serving one model (Qwen2.5-Coder-7B-Instruct-4bit, ~4 GB). They are two system-prompt identities taking sequential turns, orchestrated by plain Python. Not a free-form multi-agent chat. A structured feedback loop.

Memory Budget

Why the Orchestrator Executes Code Directly

This was a hard-won lesson. My initial design had the Sheriff agent calling an execute_code tool through the framework’s tool-calling protocol. It never worked reliably. Small quantized models (4-bit 7B) are inconsistent at structured tool calling. They would output the code in a markdown block instead of invoking the tool function, or hallucinate JSON that did not match the schema.

The fix was to take the mechanical execution out of the LLM’s hands entirely. The orchestrator extracts code from whatever the Coder produces (tool call or markdown), runs it in a subprocess, and feeds the raw stdout/stderr/exit code into the Sheriff’s prompt. The Sheriff’s job is reduced to what LLMs are actually good at: reading text and making judgments. Think of it like a code review. The reviewer does not need to run gcc themselves. A CI system runs the build, and the reviewer reads the output.

The Self-Correcting Loop

As you saw in the simulation above, each attempt follows the same pattern: the Coder generates a script, the orchestrator extracts and executes it in a sandbox, and the Sheriff reviews the output. If the Sheriff says FAIL, the error report goes back to the Coder with specific fix suggestions. If it says PASS, the loop ends.

Three safety layers prevent infinite loops:

1. Hard limits. Maximum 5 attempts, 60-second execution timeout, 120-second LLM timeout per call.
2. Stagnation detection. The system hashes the last 5 lines of each error traceback using MD5. If the same hash appears 3 times in a row, the model is stuck and the loop terminates.
3. Token budget. A hard cap of 100K tokens across all turns prevents runaway context accumulation.

Scenario 1: The Off-by-One Gauntlet

What happened: The Coder nailed it on the first try. It used datetime to verify the day of the week, wrote a clean grid layout with proper alignment, and counted weekdays correctly. The Sheriff confirmed the output and returned VERDICT: PASS in 22 seconds flat.

      3 2026
 Mon Tue Wed Thu Fri Sat Sun
                           1
   2   3   4   5   6   7   8
   9  10  11  12  13  14  15
  16  17  18  19  20  21  22
  23  24  25  26  27  28  29
  30  31
Weekdays: 22, Weekend days: 9

This is the easiest scenario, but it is not trivial. The model needs to handle date math, grid alignment with padding, and the edge case of a month starting on Sunday. A single off-by-one error in the first-day offset would cascade through every row. The model got it right in one shot.

Scenario 2: The CSV Detective

This one tells the best self-correction story. The model needed three attempts, each failing in a different way.

NameError: name 'region_revenue' is not defined
  File "script.py", line 63, in print_report
    print(f"Top Region: {top_region} (${region_revenue[top_region]:,.2f})")

NameError: name 'month_revenue' is not defined. Did you mean: 'total_revenue'?
  File "script.py", line 59, in print_report
    print(f"Top Month: {top_month} (${month_revenue[top_month]:,.2f})")

=== Sales Analysis Report ===
Revenue by Product:
  Doohickey: $125,089.64
  Gadget: $90,750.45
  Widget: $90,696.45
Top Region: North ($101,205.15)
Top Month: May ($42,193.29)
Total Records: 200

{
  "responsive": true,
  "tooltip": {"trigger": "axis", "axisPointer": {"type": "shadow"}},
  "title": {"text": "Revenue by Product", "subtext": "Agent-generated CSV analysis", "left": "center"},
  "xAxis": {
    "type": "category",
    "data": ["Doohickey", "Gadget", "Widget"],
    "axisLabel": {"fontSize": 13}
  },
  "yAxis": {
    "type": "value",
    "name": "Total Revenue ($)",
    "axisLabel": {"formatter": "${value}"}
  },
  "series": [
    {
      "type": "bar",
      "data": [
        {"value": 125090, "itemStyle": {"color": "#e94560"}},
        {"value": 90750, "itemStyle": {"color": "#0f3460"}},
        {"value": 90696, "itemStyle": {"color": "#533483"}}
      ],
      "barWidth": "50%",
      "label": {
        "show": true,
        "position": "top",
        "formatter": ["$125,090", "$90,750", "$90,696"],
        "fontSize": 13,
        "fontWeight": "bold"
      },
      "emphasis": {"itemStyle": {"shadowBlur": 10, "shadowColor": "rgba(0,0,0,0.3)"}}
    }
  ]
}

The CSV scenario is interesting because it tests the full data pipeline: file I/O, random data generation with a seed, reading data back, aggregation, and precise string formatting. Each of those is a potential failure point, and the model had to coordinate all of them in a single script.

Scenario 3: Gradient Descent from Scratch

  File "script.py", line 65
    - The script generates 200 data points for binary classification.
          ^^^^^^
SyntaxError: invalid syntax

Epoch 200: loss=0.0819, acc=100.0%
Epoch 400: loss=0.0603, acc=100.0%
Epoch 600: loss=0.0504, acc=100.0%
Epoch 800: loss=0.0444, acc=100.0%
Epoch 1000: loss=0.0402, acc=100.0%
Final accuracy: 100.0%

The training curve shows textbook convergence, from near-random (0.69) to 0.04 loss:

{
  "responsive": true,
  "title": {"text": "Training Convergence", "subtext": "Agent-written neural network", "left": "center"},
  "tooltip": {"trigger": "axis"},
  "legend": {"data": ["Loss", "Accuracy (%)"], "bottom": "0%"},
  "xAxis": {"type": "category", "data": ["0", "200", "400", "600", "800", "1000"], "name": "Epoch", "nameLocation": "center", "nameGap": 30},
  "yAxis": [
    {"type": "value", "name": "Loss", "position": "left", "min": 0, "max": 0.75},
    {"type": "value", "name": "Accuracy (%)", "position": "right", "min": 0, "max": 110, "splitLine": {"show": false}}
  ],
  "series": [
    {
      "name": "Loss",
      "type": "line",
      "data": [0.693, 0.0819, 0.0603, 0.0504, 0.0444, 0.0402],
      "smooth": true,
      "symbol": "circle",
      "symbolSize": 8,
      "lineStyle": {"color": "#e94560", "width": 3},
      "itemStyle": {"color": "#e94560"},
      "areaStyle": {"color": {"type": "linear", "x": 0, "y": 0, "x2": 0, "y2": 1, "colorStops": [{"offset": 0, "color": "rgba(233,69,96,0.3)"}, {"offset": 1, "color": "rgba(233,69,96,0.02)"}]}}
    },
    {
      "name": "Accuracy (%)",
      "type": "line",
      "yAxisIndex": 1,
      "data": [50, 100, 100, 100, 100, 100],
      "smooth": true,
      "symbol": "circle",
      "symbolSize": 8,
      "lineStyle": {"color": "#00CC96", "width": 3},
      "itemStyle": {"color": "#00CC96"},
      "areaStyle": {"color": {"type": "linear", "x": 0, "y": 0, "x2": 0, "y2": 1, "colorStops": [{"offset": 0, "color": "rgba(0,204,150,0.3)"}, {"offset": 1, "color": "rgba(0,204,150,0.02)"}]}}
    }
  ]
}

The decision boundary learned by the agent’s neural network. The line 9.26x + 8.59y - 0.03 = 0 cleanly separates the two classes. Hover over any point to see its coordinates and class.

{
  "responsive": true,
  "title": {"text": "Decision Boundary", "subtext": "Agent-written neural network (200 data points)", "left": "center"},
  "tooltip": {"trigger": "item", "formatter": "({c})"},
  "legend": {"data": ["Class 0 (negative)", "Class 1 (positive)", "Decision boundary"], "bottom": "0%"},
  "xAxis": {"type": "value", "name": "Feature 1", "nameLocation": "center", "nameGap": 30, "min": -3.5, "max": 2.5},
  "yAxis": {"type": "value", "name": "Feature 2", "nameLocation": "center", "nameGap": 40, "min": -2.5, "max": 4.5},
  "series": [
    {
      "name": "Class 0 (negative)",
      "type": "scatter",
      "symbolSize": 8,
      "itemStyle": {"color": "#636EFA", "opacity": 0.8},
      "emphasis": {"itemStyle": {"shadowBlur": 10, "shadowColor": "rgba(0,0,0,0.3)"}},
      "data": [[-0.234,-0.234],[1.579,0.767],[-0.463,-0.466],[-1.725,-0.562],[-1.013,0.314],[-0.908,-1.412],[0.068,-1.425],[-0.544,0.111],[-1.151,0.376],[-0.601,-0.292],[-0.013,-1.058],[0.823,-1.221],[0.209,-1.96],[-1.328,0.197],[-0.116,-0.301],[-1.479,-0.72],[0.344,-1.763],[0.324,-0.385],[-0.677,0.612],[-0.839,-0.309],[-0.479,-0.186],[-1.106,-1.196],[0.362,-0.645],[-2.62,0.822],[0.087,-0.299],[0.092,-1.988],[-0.808,-0.502],[-0.53,0.513],[-0.702,-0.328],[-0.392,-1.464],[0.005,-0.235],[-1.415,-0.421],[-0.343,-0.802],[-1.919,-0.027],[-0.035,-1.169],[0.791,-0.909],[-0.991,-0.566],[0.1,-0.503],[-1.551,0.069],[-1.062,0.474],[-0.783,-0.322],[0.814,-1.231],[-1.607,0.185],[-1.237,-1.32],[-0.68,0.232],[0.293,-0.714],[-1.191,0.657],[-0.975,0.787],[-0.245,-0.754],[-0.89,-0.816],[-0.847,-1.515],[0.214,-1.246],[-0.884,0.154],[0.058,-1.143],[-1.378,-0.938],[-0.773,-0.237],[-0.485,0.082],[0.686,-1.613],[0.064,-1.078],[-0.715,0.68],[-0.73,0.216],[0.046,-0.652],[-2.025,0.186],[-0.662,0.852],[-0.793,-0.115],[-1.2,-0.335],[-0.475,-0.653],[-1.261,0.918],[-1.519,-0.484],[-0.927,-0.06],[-3.241,-1.024],[-0.253,-1.248],[-0.44,0.131],[-0.982,0.462],[0.199,-0.6],[0.07,-0.385],[0.281,-0.623],[-0.829,-0.56],[-0.218,1.099],[0.324,-0.13],[-1.006,-1.214],[-0.012,-0.897],[0.076,-0.677],[-0.825,-0.321],[0.413,-0.564],[-0.822,0.244],[0.245,-0.507],[-0.471,0.232],[-1.448,-1.407],[-0.718,-0.213],[-0.019,-1.003],[-0.019,-0.289],[0.323,-0.827],[0.098,-0.773],[-0.84,-0.599],[-2.124,-0.526],[-0.759,0.15],[-0.898,0.492],[-1.713,1.354],[0.497,-0.138]]
    },
    {
      "name": "Class 1 (positive)",
      "type": "scatter",
      "symbolSize": 8,
      "itemStyle": {"color": "#EF553B", "opacity": 0.8},
      "emphasis": {"itemStyle": {"shadowBlur": 10, "shadowColor": "rgba(0,0,0,0.3)"}},
      "data": [[0.648,1.523],[-0.469,0.543],[0.242,-1.913],[-0.602,1.852],[0.738,0.171],[-0.461,1.057],[1.031,0.931],[0.331,0.976],[0.813,1.356],[-0.072,1.004],[0.361,1.538],[-0.036,1.565],[-0.22,0.357],[1.478,-0.518],[0.915,0.329],[0.097,0.969],[0.296,0.261],[-0.161,0.404],[1.886,0.175],[0.258,-0.074],[0.06,2.463],[-0.192,0.302],[1.143,0.752],[1.403,-1.402],[0.587,2.19],[-0.919,1.55],[0.227,1.307],[0.26,0.782],[0.522,0.297],[0.25,0.346],[1.866,0.474],[1.159,-0.821],[0.963,0.413],[0.822,1.897],[-0.077,0.341],[0.277,0.827],[0.013,1.454],[-0.265,2.72],[-0.223,0.714],[0.473,-0.073],[-0.447,0.856],[0.173,0.385],[0.358,0.561],[1.083,1.054],[0.515,0.514],[0.515,3.853],[0.571,1.136],[0.954,0.651],[-0.315,0.759],[2.315,-1.867],[-0.472,1.089],[2.144,0.634],[-0.67,0.852],[0.505,0.866],[1.765,0.405],[2.122,1.032],[1.267,-0.708],[0.444,0.775],[1.632,-1.43],[1.441,-1.436],[1.163,0.01],[0.114,0.662],[1.586,-1.238],[2.133,-1.952],[-0.152,0.588],[-0.589,0.85],[0.357,-0.693],[0.9,0.307],[0.813,0.63],[0.747,0.61],[-0.021,0.117],[1.278,-0.592],[0.547,-0.202],[0.825,0.814],[1.305,0.021],[0.682,-0.31],[0.097,0.595],[-0.818,2.092],[1.158,0.792],[0.624,0.628],[0.975,-0.147],[0.311,1.475],[0.858,-0.16],[0.519,1.533],[-0.109,0.402],[0.69,-0.401],[0.224,0.013],[0.025,0.498],[1.451,0.959],[2.153,-0.767],[0.872,0.183],[2.19,-0.808],[0.342,1.876],[0.95,-0.577],[-1.32,1.831],[1.179,-0.469],[-0.115,1.238],[1.466,-0.226],[0.209,0.074]]
    },
    {
      "name": "Decision boundary",
      "type": "line",
      "smooth": false,
      "symbol": "none",
      "lineStyle": {"color": "#e94560", "width": 2.5, "type": "dashed"},
      "data": [[-3.5, 3.776], [2.5, -2.692]],
      "markArea": {
        "silent": true,
        "itemStyle": {"color": "rgba(233,69,96,0.06)"},
        "data": [[{"xAxis": -3.5, "yAxis": 3.776}, {"xAxis": 2.5, "yAxis": 4.5}]]
      }
    }
  ]
}

The model got the forward pass, backward pass, and weight update all correct on the second try.

What I Learned About Small Models

Building this taught me several things about what 7B-4bit models can and cannot do.

They are good at:

• Writing complete, self-contained scripts from detailed specifications
• Fixing bugs when given the exact traceback
• Restructuring code between attempts (not just patching the broken line)
• Following formatting constraints ($12,345.67 comma separators, grid alignment)

They struggle with:

• Consistent tool calling (they prefer to output markdown)
• Very long generations without truncation (~60 lines seems to be the comfort zone)
• Multi-layer backpropagation with correct tensor shapes (I tried a 2-layer neural net and it failed all 10 attempts across two runs on the chain rule math)
• Not mixing code with natural language explanations in the same output

Design implications:

• Extract code from whatever format the model produces. Do not depend on tool calling.
• Keep tasks scoped so the solution fits in ~50-80 lines.
• Give explicit formulas for math-heavy tasks (the model can implement given formulas but struggles to derive them).
• Use the orchestrator for mechanical tasks (execution, file I/O) and the LLM for judgment tasks (generation, analysis).

The Cost Equation

{
  "responsive": true,
  "tooltip": {"trigger": "axis"},
  "xAxis": {
    "type": "category",
    "data": ["GPT-4 API", "Claude API", "Local MLX"],
    "axisLabel": {"fontSize": 13}
  },
  "yAxis": {
    "type": "value",
    "name": "USD",
    "max": 0.45,
    "nameTextStyle": {"fontSize": 13}
  },
  "series": [
    {
      "type": "bar",
      "data": [
        {"value": 0.30, "itemStyle": {"color": "#636EFA"}},
        {"value": 0.15, "itemStyle": {"color": "#EF553B"}},
        {"value": 0.00, "itemStyle": {"color": "#00CC96"}}
      ],
      "barWidth": "50%",
      "label": {
        "show": true,
        "position": "top",
        "formatter": ["~$0.30", "~$0.15", "$0.00"],
        "fontSize": 13,
        "fontWeight": "bold"
      }
    }
  ]
}

The individual run is cheap on any platform. But the cost model changes when agents get stuck. A 10-iteration retry loop that would cost $1-3 on cloud APIs costs nothing locally. You can prototype, iterate, and experiment without watching a billing dashboard. For learning and development, that changes the economics entirely.

The tradeoff is capability. A 7B-4bit model is not GPT-4. It cannot handle the same task complexity or recover from the same depth of errors. But for well-scoped tasks with clear specifications, it gets there. And it does it in under 4 minutes, fully offline, at zero cost.

Try It Yourself

# Clone and install
git clone https://github.com/thenamangoyal/multi-agent-mlx.git
cd multi-agent-mlx
uv sync

# Run a task (auto-starts the MLX server, ~4 GB RAM)
uv run factory run "Write a Python script that prints the first 20 prime numbers"

# Or run all three showcase scenarios
uv run python scenarios/run_all.py

# Stop the server when done to free RAM
pkill -f "mlx_lm.server"

Requirements: macOS on Apple Silicon, Python 3.12+, uv. The model downloads automatically on first run (~4 GB). The whole system fits comfortably in 10 GB of RAM.

{
  "responsive": true,
  "title": {"text": "Scenario Breakdown", "left": "center"},
  "tooltip": {"trigger": "item"},
  "legend": {"bottom": "0%", "data": ["Time (s)", "Attempts", "Lines of Code"]},
  "radar": {
    "indicator": [
      {"name": "Time (s)", "max": 180},
      {"name": "Attempts", "max": 5},
      {"name": "Code Lines", "max": 80},
      {"name": "Task Complexity", "max": 5}
    ]
  },
  "series": [
    {
      "type": "radar",
      "data": [
        {
          "value": [22, 1, 47, 2],
          "name": "Calendar",
          "areaStyle": {"opacity": 0.2},
          "lineStyle": {"color": "#00CC96"},
          "itemStyle": {"color": "#00CC96"}
        },
        {
          "value": [159, 3, 53, 4],
          "name": "CSV Detective",
          "areaStyle": {"opacity": 0.2},
          "lineStyle": {"color": "#636EFA"},
          "itemStyle": {"color": "#636EFA"}
        },
        {
          "value": [50, 2, 54, 3],
          "name": "Neural Net",
          "areaStyle": {"opacity": 0.2},
          "lineStyle": {"color": "#EF553B"},
          "itemStyle": {"color": "#EF553B"}
        }
      ]
    }
  ]
}

Final Thoughts

The most surprising thing about this project is not that it works. It is that the failure modes are interesting. When the Coder produces broken code, the Sheriff does not just say “it’s broken.” It identifies the error type, points to the exact line, and suggests a specific fix. And the Coder, reading that report, does not just patch the one line. It often rewrites the entire approach. That back-and-forth, mediated by a simple Python loop instead of a complex multi-agent framework, produces working code reliably enough to be useful.

The obvious next step is scaling up. A 14B or 32B model on a 64 GB Mac would handle the 2-layer backprop task that stumped the 7B model. Multiple tool calls would probably start working reliably. The architecture stays the same; you just swap the model ID. That is the beauty of building on top of mlx_lm.server: the agents do not know or care what model is behind the endpoint.

For now, the 7B model on 16 GB is the sweet spot for prototyping. It is fast enough to iterate (22 seconds for simple tasks), cheap enough to experiment freely ($0.00), and private enough to use on any codebase. Not bad for two agents arguing on a laptop.

That is exactly what autoresearch does. Andrej Karpathy built a system where an AI agent modifies a training script, trains a small language model for 5 minutes, checks if the result improved, and loops. You wake up to a log of experiments and (hopefully) a better model. The concept is wild, the code is real, and the results are legit. The catch? It was built for H100 GPUs.

I wanted to run it on my MacBook. So I ported the whole thing to Apple Silicon using MLX, Apple’s machine learning framework. No CUDA, no cloud GPU, no $3/hour rental fees. Just your Mac. The code is at github.com/thenamangoyal/autoresearch (Goyal, 2026).

Does it actually work? Here is the short version. An untrained model starts with a training loss of 9.01, which is essentially random noise over an 8,192-token vocabulary (log₂(8192) ≈ 13, but BPE tokens encode multiple bytes, so the effective starting point is lower). After just 5 minutes and 55 optimizer steps on an M1 Pro with 16GB, the loss drops to 6.76 and the validation score lands at 2.371 BPB. The model is clearly learning language structure. Give it more time and the numbers keep dropping: the community MLX fork on an M4 Max reports 1.808 BPB after a single run, and 1.295 BPB after letting the autonomous loop iterate overnight. On the original H100, Karpathy’s CUDA code fits ~11,500 steps in the same 5-minute window and reaches roughly 1.0 BPB. The model and code are the same. The gap is 96x less compute on the Mac.

Results at a Glance

Here is what a single 5-minute run looks like on my M1 Pro 16GB:

The obvious question: how does this compare to the H100 the system was designed for?

{
  "data": [
    {
      "x": ["Tok/sec", "Steps in 5 min"],
      "y": [2500000, 11500],
      "name": "H100 (CUDA)",
      "type": "bar",
      "marker": {"color": "#636EFA"}
    },
    {
      "x": ["Tok/sec", "Steps in 5 min"],
      "y": [26000, 55],
      "name": "M1 Pro 16GB (MLX)",
      "type": "bar",
      "marker": {"color": "#EF553B"}
    }
  ],
  "layout": {
    "title": {"text": "H100 vs M1 Pro: Raw Throughput"},
    "yaxis": {"title": "Value", "type": "log"},
    "barmode": "group",
    "legend": {"yanchor": "top", "y": 0.99, "xanchor": "right", "x": 0.99}
  }
}

Yes, the H100 is roughly 96x faster. That is not the point. The point is that you can run this at all on a laptop, iterate on architecture ideas, and prototype training experiments before ever touching a cloud GPU. For learning and experimentation, the Mac is more than enough.

The --time-budget flag lets you scale up when you have time to spare:

{
  "data": [
    {
      "x": [3.6, 23, 47, 94],
      "y": ["5 min", "15 min", "30 min", "60 min"],
      "type": "bar",
      "orientation": "h",
      "marker": {"color": ["#636EFA", "#EF553B", "#00CC96", "#AB63FA"]},
      "text": ["~55 steps", "~357 steps", "~714 steps", "~1429 steps"],
      "textposition": "auto"
    }
  ],
  "layout": {
    "title": {"text": "Tokens Processed by Time Budget (M1 Pro)"},
    "xaxis": {"title": "Millions of Tokens"},
    "yaxis": {"title": ""},
    "showlegend": false,
    "margin": {"l": 80}
  }
}

In my experience, 15 to 30 minutes hits the sweet spot on Apple Silicon. You get enough steps for the optimizer to meaningfully converge, without burning hours waiting.

The Autonomous Research Loop

The genius of autoresearch is how simple the loop is. An AI agent (Claude, GPT, whatever you prefer) reads the instructions in program.md, tweaks the model code in train.py, commits the change, trains for exactly 5 minutes, checks BPB, and either keeps the commit or reverts. Then it does it again. All night if you let it.

flowchart TD
    A[Read program.md] --> B[Modify train.py]
    B --> C[Git commit]
    C --> D[Train for 5 min]
    D --> E[Read val_bpb]
    E --> F{Improved?}
    F -- Yes --> G[Keep commit]
    F -- No --> H[Git revert]
    G --> I[Log to results.tsv]
    H --> I
    I --> B

The design is intentionally constrained. One file to edit. One metric to optimize. A fixed time budget so every experiment is directly comparable regardless of what the agent changes: model size, architecture, optimizer, batch size, all of it is fair game inside train.py. The agent’s creativity is the only variable.

As Karpathy put it in the original repo:

One day, frontier AI research used to be done by meat computers in between eating, sleeping, having other fun, and synchronizing once in a while using sound wave interconnect in the ritual of “group meeting”. That era is long gone.

On the Mac, each experiment takes about 5 minutes, which gives you roughly 12 experiments per hour and close to 100 if you let it run overnight. Not quite datacenter speed, but enough to explore a surprising number of architectural ideas.

Model Architecture

The model is a GPT variant (Vaswani et al., 2017) with several modern tricks packed into a surprisingly compact 11.5M parameter design. I found it interesting how many recent ideas Karpathy managed to squeeze into a single training file. Here is the high-level flow:

flowchart TD
    T[Token IDs] --> WTE[Token Embedding]
    T --> VE[Value Embedding]
    WTE --> RS[Residual Scaling]
    RS --> N1[RMSNorm]
    N1 --> ATT[Self-Attention + RoPE]
    VE --> ATT
    ATT --> N2[RMSNorm]
    N2 --> MLP[MLP · Squared ReLU]
    MLP -->|repeat x DEPTH| RS
    MLP --> NF[RMSNorm]
    NF --> LM[lm_head]
    LM --> SC[Softcap · tanh]

Let me walk through what makes this architecture interesting.

Value Embeddings (ResFormer). This one surprised me. Every other layer gets its own value embedding table. Instead of computing values purely from the hidden state, the model mixes in a direct embedding lookup gated by a small linear projection. It is the ResFormer trick: a residual connection through the value path that helps gradient flow in deep models. In a 4-layer model it might seem unnecessary, but it lets the agent experiment with deeper configurations without running into vanishing gradient issues.

Sliding Window Attention (SSSL). The window pattern SSSL means three layers of short-range (half context length) attention followed by one layer of long-range (full context) attention. On the H100, this is handled by Flash Attention (Dao et al., 2022). On MLX, I implemented it with additive masks instead. It is slower, but functionally equivalent. The last layer always gets full attention regardless of the pattern, so the model can always attend to the full 2,048-token context when it matters.

RoPE with QK-Norm. Queries and keys get rotary position embeddings (Su et al., 2021) followed by RMS normalization. This combination stabilizes training and eliminates the need for learned position embeddings. It is one of those small details that just works.

Softcap at 15.0. The logits pass through 15 * tanh(logits / 15) before the loss computation. This bounds logit magnitudes and prevents the model from becoming overconfident. The technique comes from Gemma 2 (Gemma Team et al., 2024) and is especially helpful early in training when the model might otherwise push logits to extreme values on frequent tokens.

Squared ReLU. The MLP activation is relu(x)² instead of GELU or SiLU. Simpler, faster, and produces sparser activations. In a small model where every FLOP counts, this is a sensible choice.

Per-Layer Residual Scaling. Each layer has two learnable scalars that control how much the running hidden state and the original embedding contribute. The core of the forward pass looks like this:

x = self.resid_lambdas[i] * x + self.x0_lambdas[i] * x0
ve = self.value_embeds[str(i)](idx) if str(i) in self.value_embeds else None
x = block(x, ve, masks[i])

This gives the model fine-grained control over each layer’s contribution to the residual stream. In practice, I noticed the model learns to progressively reduce x0_lambda in later layers, relying more on the transformed representations as depth increases.

Optimizer: AdamW with Six Parameter Groups

This is not your typical “one learning rate for everything” AdamW (Loshchilov & Hutter, 2017). The optimizer splits the model into six groups, each with its own learning rate, betas, and weight decay. Getting this right turned out to be one of the trickier parts of the port.

A few things stand out. Embedding learning rates are 150x higher than unembedding rates, because embeddings need to move fast early in training while the output projection benefits from stability. Only the transformer matrices get weight decay. The x0_lambdas use a beta1 of 0.96 (vs 0.8 for everything else), giving them more momentum so they adjust slowly and smoothly. And embedding/unembedding LRs scale with 1/sqrt(model_dim/768), following the dimension-scaled convention from recent scaling law work.

Weight decay schedule. One improvement I made over the reference MLX fork is a linear weight decay schedule. The effective weight decay is WEIGHT_DECAY * (1 - progress), which decays to zero by the end of training. The intuition is simple: during the warmdown phase when the learning rate is already near zero, weight decay just pushes weights toward zero without the optimizer being able to pull them back. The upstream CUDA version does this, but the reference fork missed it. In my tests it typically yields 2 to 5% BPB improvement, which is significant for free.

LR schedule. No warmup at all (WARMUP_RATIO = 0.0). The second half of training is spent in linear warmdown to zero. This aggressive schedule makes sense for a 5-minute budget where you cannot afford to spend steps warming up.

Data Pipeline

The training data comes from ClimbMix 400B, a curated web-scale dataset downloaded as parquet shards into ~/.cache/autoresearch/. Running prepare.py once downloads the data and trains a BPE tokenizer with a vocabulary of 8,192 tokens. The whole process takes about 2 minutes.

The dataloader uses BOS-aligned best-fit packing: documents are packed into sequences of 2,048 tokens, each starting with a BOS token. Documents that do not fit are split across sequences, but every sequence begins at a document boundary. This achieves close to 100% token utilization with no padding waste, which matters when you only have 3.6M tokens to work with in a 5-minute budget.

On Apple Silicon, the big win is unified memory. There is no CPU-to-GPU transfer. The data, model weights, optimizer states, and gradients all live in the same memory pool. On a CUDA setup with a consumer GPU, PCIe bandwidth can bottleneck the dataloader. On a Mac, that entire category of performance issues just does not exist. The tradeoff is that you are sharing that same memory pool with macOS, your browser, and everything else running on the machine, which is exactly why the 16GB M1 Pro hits memory pressure during training.

Why BPB instead of perplexity? The evaluation metric is bits per byte (BPB), computed as cross_entropy_loss * (num_tokens / num_bytes). This normalizes for tokenizer vocabulary size, so you can swap the tokenizer or change the vocab size and still get comparable numbers. It is the right metric for autoresearch, where the agent might decide to change the tokenization strategy entirely. Lower is better.

Training in Action

Here is real data from my 5-minute benchmark. Every data point is from the actual run, nothing simulated:

{
  "data": [
    {
      "x": [0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31,32,33,34,35,36,37,38,39,40,41,42,43,44,45,46,47,48,49,50,51,52,53,54],
      "y": [9.012,9.012,8.688,8.426,8.216,8.060,7.960,7.882,7.812,7.757,7.705,7.670,7.631,7.604,7.576,7.553,7.527,7.511,7.492,7.464,7.444,7.414,7.396,7.379,7.355,7.334,7.309,7.291,7.274,7.244,7.226,7.207,7.192,7.168,7.145,7.119,7.094,7.071,7.048,7.022,7.000,6.984,6.965,6.948,6.933,6.913,6.894,6.877,6.856,6.834,6.817,6.801,6.787,6.770,6.762],
      "type": "scatter",
      "mode": "lines+markers",
      "name": "Training Loss (EMA)",
      "line": {"color": "#636EFA", "width": 2},
      "marker": {"size": 4}
    }
  ],
  "layout": {
    "title": {"text": "Training Loss Curve (M1 Pro 16GB, 5-min budget)"},
    "xaxis": {"title": "Step"},
    "yaxis": {"title": "Smoothed Training Loss"},
    "showlegend": false
  }
}

The loss drops steadily from 9.01 to 6.76 over 55 steps. No spikes, no NaN incidents, no instability. That smooth curve is exactly what you want to see from a well-configured optimizer on a small model.

But the throughput chart tells a more interesting story:

{
  "data": [
    {
      "x": [0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31,32,33,34,35,36,37,38,39,40,41,42,43,44,45,46,47,48,49,50,51,52,53,54],
      "y": [7604,24127,22891,22148,18671,14971,22997,23022,23742,17105,16300,20709,25035,22540,18241,11142,20179,17631,16099,18564,17911,17994,14105,1742,2628,5754,5672,3634,7285,2245,5507,6574,16197,29251,29009,29612,29569,28550,25040,28735,28859,28095,27022,25195,20166,24015,24746,14838,18569,25240,26747,26207,29411,28793,28899],
      "type": "scatter",
      "mode": "lines+markers",
      "name": "Tok/sec",
      "line": {"color": "#EF553B", "width": 2},
      "marker": {"size": 4}
    }
  ],
  "layout": {
    "title": {"text": "Throughput Over Training (M1 Pro 16GB)"},
    "xaxis": {"title": "Step"},
    "yaxis": {"title": "Tokens per Second"},
    "showlegend": false,
    "annotations": [
      {
        "x": 23, "y": 1742,
        "text": "Memory pressure",
        "showarrow": true,
        "arrowhead": 2,
        "ax": 40, "ay": -40,
        "font": {"size": 11, "color": "#EF553B"}
      },
      {
        "x": 29, "y": 2245,
        "text": "Swap thrashing",
        "showarrow": true,
        "arrowhead": 2,
        "ax": 40, "ay": -40,
        "font": {"size": 11, "color": "#EF553B"}
      }
    ]
  }
}

See those dramatic dips around steps 23 to 30? That is macOS memory pressure in action. With the model weights, optimizer states (two momentum buffers per parameter group), gradient accumulators, and attention masks all resident in unified memory, the 16GB M1 Pro is right at the edge. When macOS starts compressing and swapping pages, throughput drops from ~26K tok/sec to under 2K. Each of those slow steps takes 15 to 37 seconds instead of the usual 2.5 seconds. The training still converges correctly (the loss curve above is smooth right through those dips), but you lose precious steps from your 5-minute budget. On an M4 Max with 36 or 64GB, these dips disappear entirely and you get consistent 29K+ tok/sec throughout.

Improvements Over the Reference Fork

When I started this port, trevin-creator/autoresearch-mlx already existed as an early community MLX adaptation. I used it as a reference but ended up rewriting most of the training loop to fix several issues I ran into:

The eval batch size issue was the most frustrating to debug. The reference fork uses FINAL_EVAL_BATCH_SIZE=256, which works fine on an M4 Max with 27GB available. But on a 16GB machine, the evaluation step triggers an out-of-memory crash after the 5-minute training run completes successfully. You sit through the whole training, only for the eval to kill the process and lose everything. Reducing it to 16 eliminates this entirely.

The __main__ guard is subtle but important for the autonomous agent workflow. Without it, the agent’s import train (which it uses to inspect the module) silently kicks off a full training run. That is 5 minutes wasted before the agent even starts its actual experiment.

The reference fork reports a best result of 1.295 BPB on M4 Max and 1.353 BPB on Mac Mini after extended autonomous runs. My single 5-minute baseline on M1 Pro 16GB reaches 2.371 BPB. The gap is hardware, not code: fewer steps due to memory pressure and a slower processor. The raw per-step training throughput is identical between the two implementations since both use the same MLX ops. On equivalent hardware with the weight decay fix and other improvements, I would expect this port to match or beat the reference fork’s numbers.

Try It Yourself

Four commands and you are training:

# 1. Install uv (if you don't have it)
curl -LsSf https://astral.sh/uv/install.sh | sh

# 2. Install dependencies
uv sync

# 3. Download data and train tokenizer (one-time, ~2 min)
uv run prepare.py

# 4. Run a training experiment
uv run train.py

Want a longer run? Use the --time-budget flag:

uv run train.py --time-budget 600   # 10 minutes
uv run train.py --time-budget 1800  # 30 minutes
uv run train.py --time-budget 3600  # 1 hour

To run the full autonomous loop, point an AI agent (Claude Code, Codex, or similar) at program.md in the repo:

Hi, have a look at program.md and let's kick off a new experiment!
Let's do the setup first.

The agent will create a branch, establish a baseline BPB, and start iterating. Make sure to disable all permission prompts so it can run unattended. On an M1 Pro you will get about 12 experiments per hour. Let it run overnight and check the git log in the morning.

Requirements: Apple M-series Mac (M1/M2/M3/M4), Python 3.10+, uv.

The code is at github.com/thenamangoyal/autoresearch.

Conclusion

You do not need an H100 to experiment with autonomous AI research. A MacBook Pro can train an 11.5M parameter language model on real web data, with a proper six-group AdamW optimizer, modern architecture tricks like RoPE and value embeddings, and a clean autonomous iteration loop. Yes, it is 96x slower than a datacenter GPU. But it costs nothing, it sits on your desk, and it works.

There is plenty of room to push this further. The Muon optimizer from the upstream CUDA version is not yet ported (it needs SVD operations that require careful MLX adaptation). M4 Max and M4 Ultra machines with 64GB+ of unified memory should eliminate the throughput dips and unlock much longer effective training runs. And the autonomous loop itself could be made smarter: better search strategies, multi-objective optimization, cross-experiment learning from the full results history.

For now, the foundation works. Clone the repo, run prepare.py, and let your MacBook do some research overnight. You might be surprised by what it finds.

Thanks to Andrej Karpathy for the original autoresearch concept and to the trevin-creator community fork for the initial MLX exploration.

As a long-time resident of the Apple walled garden, you know the magic: Apple Books syncs your reading progress effortlessly, Notes appear instantly across devices, Reminders pop up reliably, Keychain handles passwords seamlessly, Photos are just there, iMessage keeps you connected, AirDrop zips files around, and Find My keeps track of your gear. It’s a beautifully integrated system… as long as you stay within its walls.

Stepping outside with an Android phone means rethinking how these essential services work together. The good news? It’s entirely possible to maintain, and even enhance, your productivity with the right cross-platform tools. The goal is to find replacements that work reliably on macOS, Android, iOS, and iPadOS.

The Challenge Replacing Seamless Integration

Let’s break down the core Apple services and find robust, cross-platform alternatives that play nicely with your Mac, new Android phone, iPad, and secondary iPhone.

TLDR Quick Reference

Here’s a quick look at the recommended replacements for common Apple ecosystem services:

Detailed Replacement Strategies

Let’s dive deeper into each replacement.

1. Replacing Apple Books

⟶

• Recommendation: Amazon Kindle App + Send to Kindle
• How it Works:
  • Install the Kindle app on your Android phone, iPad, and Mac (available on the App Store/web).
  • Log in with your Amazon account. Your purchased Kindle books sync automatically.
  • For personal documents (PDFs, epubs you own): Use the Send to Kindle service. There are apps/extensions for Mac (https://www.amazon.com/sendtokindle) and you can email documents directly. These docs sync across your Kindle apps just like purchased books.
• Pros: Excellent sync, large store, actively developed, handles personal documents well.
• Cons: Tied to the Amazon ecosystem.

2. Replacing Apple Notes

⟶

+

• Recommendation: Microsoft OneNote (for structured notes) + Google Keep (for quick capture)
• How it Works:
  • OneNote: Offers robust organization with notebooks, sections, and pages. Great for detailed notes, research, meeting minutes. Supports Apple Pencil on iPad, embedding PDFs/images, and has native apps for Mac, Android, iOS, iPadOS, and Web. Syncs via your Microsoft account (free). Excellent for annotating research papers directly within the note.
  • Google Keep: Simple, fast, and effective for quick thoughts, checklists (like groceries), and reminders. Accessible via native apps (Android/iOS) and web (Mac/iPad). Integrates well with other Google services.
• Pros: OneNote is powerful and truly cross-platform. Keep is lightweight and fast.
• Cons: Using two apps might feel disjointed initially. OneNote can sometimes feel ‘heavy’ compared to Apple Notes.

3. Replacing Apple Calendar

⟶

• Recommendation: Google Calendar (as the backend)
• How it Works:
  • On your Mac: Go to System Settings > Internet Accounts and add your Google Account. Ensure “Calendars” is checked. Open the Apple Calendar app, go to Preferences > General, and set your Google Calendar as the “Default Calendar Account”. You can continue using the familiar Apple Calendar interface on your Mac, but events will sync via Google.
  • On Android/iPad/iOS: Use the excellent native Google Calendar app.
• Pros: Seamless sync, leverages the robust Google Calendar backend, allows using the native Mac Calendar app.
• Cons: Requires using Google as the central calendar authority.

4. Replacing Apple Reminders

⟶

• Recommendation: Google Tasks (integrated with Google Calendar)
• How it Works:
  • Use the dedicated Google Tasks app on Android and iOS & iPad.
  • On Mac, access Google Tasks via the sidebar in Gmail or Google Calendar on the web and install as progessive web app.
  • Tasks with due dates automatically appear in your Google Calendar (which you’re likely already syncing everywhere).
• Pros: Simple interface, excellent integration with Google Calendar, cross-platform apps + web access.
• Cons: Less feature-rich than some dedicated task managers, web access on Mac isn’t a native app.

5. Replacing AirDrop

⟶

+

• Recommendation: NearDrop (for Mac) working with Quick Share (built into Android) on same Wi-Fi network
• How it Works:
  • Install NearDrop on your Mac: (GitHub - requires manual installation or using a community build from brew). This app allows receiving (easier) from Android to Mac. Also supports sending (more nuanced) Mac from Android.
  • Mac to Android: The most reliable Mac-to-Android flow might be initiating from Android’s Quick Share receive mode, then sending from Mac.
    • Install Google Files (built-in for Pixel) on android.
    • Open the “Files” app on your Pixel and select “Quick Share”.
    • On Mac, right click the file, select share. Click “Edit extensions”, and enable the “Near Drop” extension (first time only).
    • Now right click the file on mac, click Share -> NearDrop. And choose the android phone.
  • Android to Mac: Ensure the NearDrop app is running on your Mac. On your Android phone, select a file, tap Share, choose “Quick Share” (sometimes called Nearby Share), and select your Mac when it appears. Approve the transfer on your Mac via the NearDrop prompt.
• Pros: Wireless, relatively fast for local transfers.

• Cons: Requires a third-party app (NearDrop) on Mac. Both devices usually need to be on the same Wi-Fi network. Requires manual approval for transfers. NearDrop being open might make your Mac visible to anyone else on the network using Quick Share (though transfers still require approval). Less seamless than AirDrop.

• Alternative: Cloud storage (Google Drive, Dropbox, OneDrive) or messaging apps (Signal, WhatsApp) are often easier for Mac-to-Android transfers if you are on different networks.

6. Replacing Find My (AirTags)

⟶

+

• Recommendation: Chipolo PoP (compatible with Google’s Find My Device & Apple Find My)
• How it Works:
  • Chipolo makes trackers specifically designed to work with Google’s Find My Device network and Apple Find My Network (only one at a time). For our case, set them up using your Android phone, and you can access Google Find my on web for Mac, iOS.
  • These devices leverage the network of Android devices to report their location, similar to AirTags using Apple devices.
  • Crucially, Chipolo Pop devices can be switched between Google’s network and Apple’s Find My network (though not used on both simultaneously). You’d need to reset and re-pair it to switch.
• Pros: Provides AirTag-like functionality using the Android network. Can be switched to Apple’s network if needed later.
• Cons: No Precision Finding like UWB-enabled AirTags. Location updates depend on nearby Android devices running the necessary services. Switching networks requires a reset.

7. Replacing Apple Keychain

⟶

+

/

• Recommendation: Google Chrome Password Manager + Authy or Google Authenticator
• How it Works:
  • Passwords & Passkeys:
    • On your iPhone/iPad (iOS/iPadOS): Go to Settings > Passwords > Password Options. Enable “AutoFill Passwords and Passkeys” and select Chrome (ensure you have Chrome installed). Do NOT just select “Google” if that appears as a separate option, as Chrome handles the broader sync better with Passkeys.
    • On Android: Chrome is likely the default. Check in Settings > Passwords & accounts > Autofill service.
    • On Mac: Use Google Chrome. Sign in to your Google account. Passwords and passkeys will sync and autofill within Chrome. You can also access them via passwords.google.com.
  • Verification Codes (2FA/OTP):
    • Use Authy or Google Authenticator. Both support scanning QR codes to add accounts.
    • Authy: Has native apps for Mac, Android, iOS and supports encrypted cloud backups/sync across devices. Generally preferred for multi-device use.
    • Google Authenticator: Recently added cloud sync tied to your Google account, making it cross-platform. Simpler interface than Authy.
    • Didn’t suggest Microsoft Authenticator as it doesn’t seem to simultaneous sync between iOS/Android; sticking to Authy or Google Authenticator is generally more reliable for this mixed setup*
• Pros: Chrome PM handles passwords and passkeys across all platforms via the browser. Authy/GA provide reliable cross-platform 2FA code access.
• Cons: Password access outside of Chrome on Mac might require opening Chrome. Less integrated into macOS system-wide prompts compared to Keychain.

8. Replacing iMessage

⟶

+

• Recommendation: Google Messages (using RCS) and WhatsApp
• The Elephant in the Room: This is often the trickiest transition. There’s no perfect iMessage equivalent.
• Google Messages (RCS):
  • Use Google Messages as your default SMS/RCS app on Android. RCS provides features like typing indicators, read receipts, and high-quality media sharing with other RCS users (including recent support on iPhones).
  • On Mac/iPad/Web: Use Google Messages for Web (messages.google.com/web) by scanning a QR code from your Android phone’s Messages app (under Device Pairing). This requires your phone to be online. And you can then install as PWA (Progressive web app) with its on icon on your Mac Dock.
• WhatsApp:
  • Widely used, truly cross-platform with native apps (including Mac).
  • Major Con: Migrating chat history between iOS and Android is cumbersome. While possible now, it often requires the new phone (the one receiving the chats) to be factory reset during the transfer process. Moving chats from your old iPhone to the new Pixel requires this reset. Moving back later would too.
  • Migrating chats between Android -> Android or iOS -> iOS is easy via cloud backups.
  • You can also setup your iPhone as companion phone (using linked device) to have your WhatsApp chat sync between the 2 devices. Your android doesn’t need to be online for the sync to work.
• Pros: Google Messages offers RCS improvements and web access. WhatsApp is ubiquitous.
• Cons: No single app perfectly replicates iMessage’s seamless SMS/data messaging integration across Apple hardware. Google Messages web requires the phone to be on. WhatsApp migration is painful. Green bubbles when texting iPhones from Android (though RCS helps).

9. Replacing Apple Photos

⟶

• Recommendation: Google Photos
• How it Works:
  • Install Google Photos on your Android phone, iPhone, and iPad. Access it via photos.google.com on your Mac.
  • Enable backup & sync on all mobile devices. Photos and videos will upload to your Google account.
  • There’s also a “Google Drive for Desktop” app for Mac that includes an option to upload photos from your system library or folders to Google Photos.
• Pros: Excellent cross-platform availability, powerful AI-driven search (“show me photos of beaches”), automatic albums, great sharing features, similar feel to Apple Photos in many ways. Generous free tier (though high-quality uploads now count towards storage).
• Cons: High-resolution backups require paid Google One storage eventually.

General Tips for a Smoother Transition

1. Embrace the Cloud: Relying on cloud-based services (Google, Microsoft OneNote, Amazon Kindle etc.) is key to making a mixed ecosystem work.
2. Be Patient: Setting up new workflows takes time. Don’t expect everything to feel identical overnight.
3. Focus on Workflow, Not Just Apps: Think about what you need to achieve (e.g., access notes everywhere) rather than replicating the exact Apple app.
4. Utilize Web Apps: Many services have excellent web versions accessible on your Mac or iPad via a browser, reducing the need for native apps in some cases.

Conclusion

Switching your primary phone to Android while staying invested in the Mac and iPad ecosystem presents unique challenges, but it’s far from impossible to maintain a highly productive setup. By strategically choosing cross-platform services like Google Calendar, Google Photos, OneNote, Kindle, and password managers like Chrome’s, and using tools like NearDrop and Chipolo to bridge hardware gaps, you can create a surprisingly seamless experience. It requires some setup and adjustment, but the flexibility and power of working across platforms can be very rewarding. Good luck with your new setup!

1. Introduction

Gemma 3 extends the Gemma family with models ranging from 1B to 27B parameters. Key advancements presented include:

1. Multimodality: Integration of vision understanding via a tailored, frozen SigLIP encoder, processing images as sequences of soft tokens.
2. Long Context: Support for context windows of at least 128K tokens (32K for the 1B model), enabled by architectural modifications.
3. Architectural Efficiency: A novel interleaved local/global attention mechanism specifically designed to mitigate the KV cache memory bottleneck inherent in long-context inference.
4. Enhanced Performance: Superior results compared to Gemma 2, particularly in instruction-tuned (IT) variants. This is attributed to knowledge distillation and a refined post-training recipe incorporating advanced RL techniques.
5. Improved Multilinguality: Better representation and handling of non-English languages achieved through adjustments in the pre-training data mixture and the adoption of the Gemini 2.0 tokenizer.

This analysis will dissect these aspects, examining the architectural rationale, training methodologies, evaluation paradigms, and the inherent trade-offs and limitations.

2. Model Architecture

While retaining the foundational decoder-only Transformer architecture, Gemma 3 incorporates several critical modifications compared to its predecessors, primarily focused on enabling long context efficiently and integrating vision.

KV Cache Management: Interleaved Local/Global Attention

A primary bottleneck for deploying long-context models is the KV cache size, which scales linearly with sequence length and number of layers, quickly exceeding typical device memory. Gemma 3 tackles this with a hybrid attention strategy:

• 5:1 Interleaving: The architecture alternates between local (sliding window) self-attention and global self-attention layers. Specifically, it employs a repeating pattern of 5 local layers followed by 1 global layer.
• Short Local Span: Local attention layers operate with a constrained sliding window of only 1024 tokens. This significantly limits the contribution of these layers to the overall KV cache size.
• Efficient Long Context Processing: Consequently, only the global attention layers (representing 1/6th of the total layers) need to store Keys and Values for the entire 128K context window. This architectural choice drastically reduces the KV cache footprint compared to a model where all layers attend globally, striking a balance between capturing long-range dependencies and maintaining inference feasibility on resource-constrained hardware.

sequenceDiagram
    participant I as Input Sequence (up to 128K)
    participant L5 as 5x Local Attention Layers
    participant G1 as 1x Global Attention Layer
    participant O as Output Embeddings

    I->>L5: Process with 1024-token sliding window
    Note right of L5: KV cache contribution per layer<br/>scales with window size (1024)
    L5->>G1: Process full context (128K)
    Note right of G1: KV cache contribution per layer<br/>scales with full context length (128K)
    loop Multiple Blocks
        G1->>L5: Output feeds into next block
        L5->>G1: ...
    end
    G1->>O: Final Output

Vision Integration: SigLIP and Pan & Scan (P&S)

Gemma 3 integrates visual processing using:

• Vision Encoder: A 400M parameter variant of the SigLIP Vision Transformer (Zhai et al., 2023). Notably, this encoder is frozen during the language model’s training. This simplifies training and reduces computational cost but potentially limits the synergy between visual and textual feature extraction compared to end-to-end training.
• Token Condensation: Embeddings from the SigLIP encoder are condensed into a fixed-length sequence of 256 vectors (“soft tokens”). These visual tokens are prepended to the text token sequence, serving as the language model’s input representation of the image.
• Pan & Scan (P&S): The SigLIP encoder operates at a fixed resolution (896x896). To handle arbitrary image resolutions and aspect ratios effectively, particularly for tasks involving reading text or resolving fine details, an inference-time P&S strategy (inspired by (Liu et al., 2023)) is employed. This adaptively segments the input image into potentially overlapping crops, resizes each to the target 896x896 resolution, encodes them individually, and concatenates the resulting 256-token sequences (up to a predefined maximum number of crops). While enhancing capability for high-resolution or non-square images, this adds computational overhead during inference.

Other Architectural Details

• Long Context Enablers: To effectively utilize the extended context, the Rotary Position Embeddings (RoPE) base frequency for the global self-attention layers is increased significantly from 10k (used in Gemma 2 and Gemma 3’s local layers) to 1M. This adaptation, drawing inspiration from positional interpolation techniques (Chen et al., 2023), helps the model generalize positional information over much longer sequences.
• Normalization & Attention Refinements: The models utilize Grouped-Query Attention (GQA) for efficiency, employ RMSNorm for layer normalization (both pre-norm and post-norm), and notably replace the soft-capping activation from Gemma 2 with QK-norm. This QK-normalization (inspired by recent works like (Dehghani et al., 2023), (Wortsman et al., 2023), (Team, 2024)) likely contributes to improved training stability or performance scaling.
• Tokenizer: The adoption of the Gemini 2.0 SentencePiece tokenizer (262k vocabulary) is a key change from earlier Gemma versions. Its design aims for better handling of multilingual text, digits, and whitespace, which is crucial for improved performance across diverse languages and domains like code.

3. Training and Data Strategy

The training process involves distinct pre-training and instruction fine-tuning phases, leveraging large datasets and sophisticated optimization techniques.

Pre-training

• Dataset Scale: Models were trained on substantial token budgets ranging from 2T (1B model) to 14T (27B model). This slight increase compared to Gemma 2 accommodates the new multimodal and expanded multilingual data components.
• Data Composition: The pre-training mixture comprises web documents, code, mathematics, and dialogue data. Emphasis was placed on increasing the proportion of multilingual data (both monolingual corpora across various languages and parallel data) and incorporating image-text pairs for multimodal grounding. Strategies inspired by (Chung et al., 2023) were used to manage language imbalance during sampling.
• Data Hygiene: Rigorous filtering was applied to mitigate risks associated with harmful content, remove personally identifiable information (PII), decontaminate evaluation benchmark data from the training set (a critical step for reliable evaluation), and apply quality heuristics, potentially re-weighting data based on quality metrics (Sachdeva et al., 2024).
• Knowledge Distillation: Pre-training incorporates knowledge distillation (Hinton et al., 2015) from larger, unspecified teacher models. The student model learns via cross-entropy loss, but instead of using the standard one-hot target distribution, it uses a softened distribution derived from sampling 256 logits based on the teacher’s output probabilities for each token. This guides the student towards the teacher’s internal representations.

Instruction Fine-tuning (Post-training)

This phase appears critical to Gemma 3’s strong performance on user-oriented tasks, transforming the base models into capable instruction-following agents.

• Refined Recipe: The report mentions a “novel post-training recipe,” although specific details remain limited. It clearly involves an enhanced form of knowledge distillation, this time using a large instruction-tuned teacher model.
• Reinforcement Learning (RL): Advanced RL techniques are employed, likely extending methods like BOND (Sessa et al., 2024), WARM (Ramé et al., 2024), and WARP (Ramé et al., 2024).
• Multi-Objective Optimization: RL training utilized a diverse set of reward functions targeting helpfulness, mathematical reasoning, coding proficiency, general reasoning, instruction adherence, multilingual fluency, and safety (minimizing harmfulness). These rewards leverage a combination of human feedback (RLHF), feedback from automated systems like code execution environments (Gehring et al., 2024), and ground-truth outcomes for tasks like mathematics ((DeepSeek-AI, 2024), (Lambert et al., 2024)).
• Targeted Filtering: Additional data filtering specific to the post-training phase was implemented to enhance factuality, encourage attribution, reduce hallucinations, and steer the model away from unsafe or undesirable outputs.

Quantization Aware Training (QAT)

Recognizing the need for efficient deployment, Gemma 3 models undergo QAT (Jacob et al., 2018). This involves further fine-tuning the models specifically for lower-precision inference (targeting per-channel int4, per-block int4, and switched fp8 formats popular in frameworks like llama.cpp). Instead of quantizing post-training, QAT adapts the weights during training, using the full-precision model’s output probabilities as targets, thereby minimizing the performance degradation often associated with quantization.

Table: Approximate Memory Footprints (derived from Table 3 in the report). Note the significant contribution of the KV cache even when quantized.

4. Performance Evaluation

Gemma 3 demonstrates substantial improvements over its predecessor, Gemma 2, and achieves performance levels that are highly competitive within the open model landscape, even compared to significantly larger models.

General Capabilities (Instruction-Tuned Models)

• LMSYS Chatbot Arena: The Gemma-3-27B-IT model achieved a preliminary Elo score of 1338 (as of March 8, 2025). This positions it strongly among top-performing open models in blind pairwise human evaluations, notably surpassing Gemma-2-27B-IT (1220 Elo) and even larger models like Llama-3.1-405B-Instruct (1269 Elo) at that time. This highlights the effectiveness of the post-training recipe.
• Standard Benchmarks (Table 6): Across various academic benchmarks, Gemma 3 models show consistent gains. A key finding is that the Gemma3-4B-IT model frequently matches or outperforms the much larger Gemma2-27B-IT on challenging tasks like MATH, HiddenMath, and MMLU-Lite. The flagship Gemma3-27B-IT closes the gap significantly with the powerful, proprietary Gemini-1.5-Pro on several benchmarks, particularly MATH (89.0 vs. 91.8 for Gemini Pro) and MMLU-Pro (67.5 vs. 75.8), while also showing strong coding (LiveCodeBench) and reasoning performance.

Here’s a comparative visualization for selected benchmarks based on Table 6 data:

{
  "data": [
    {
      "x": ["MMLU-Pro", "LiveCodeBench", "MATH", "HiddenMath", "MMMU (val)", "Global MMLU-Lite"],
      "y": [56.9, 20.4, 55.6, 14.8, null, 68.6],
      "name": "Gemma 2 27B IT",
      "type": "bar"
    },
    {
      "x": ["MMLU-Pro", "LiveCodeBench", "MATH", "HiddenMath", "MMMU (val)", "Global MMLU-Lite"],
      "y": [67.5, 29.7, 89.0, 60.3, 64.9, 75.1],
      "name": "Gemma 3 27B IT",
      "type": "bar"
    }
  ],
  "layout": {
    "title": {
      "text": "Gemma 3 27B IT vs Gemma 2 27B IT Performance"
    },
    "yaxis": {
      "title": "Score (%)"
    },
    "barmode": "group",
    "legend": {"yanchor":"top", "y":0.99, "xanchor":"left", "x":0.01}
  }
}

Vision and Multimodal Performance

• Gemma 3 achieves robust performance on diverse vision-language benchmarks (Tables 11 & 16), demonstrating proficiency in tasks requiring OCR (DocVQA, TextVQA), understanding of structured information (ChartQA), and general visual question answering (MMMU, VQAv2).
• The effectiveness of the Pan & Scan (P&S) mechanism is empirically validated (Table 8), providing substantial score increases on benchmarks sensitive to image resolution and text legibility, such as DocVQA (+8.2 points for 4B, +4.8 for 27B) and InfoVQA (+12.9 for 4B, +17.0 for 27B). This confirms P&S’s role in overcoming the fixed-resolution limitation of the vision encoder.
• Initial results on video benchmarks (Table 17) suggest foundational capabilities in understanding temporal sequences of images.

Long Context Evaluation

• Ablation studies focusing on RoPE rescaling (Figure 7) indicate successful generalization up to the target 128K context length for the pre-trained models, maintaining reasonable perplexity.
• Evaluations on dedicated long-context benchmarks like RULER and MRCR (Table 15) confirm the models’ ability to process and utilize long inputs. However, the observed performance degradation when moving from 32K to 128K contexts suggests that while the architecture supports the length, effectively performing complex reasoning over the entire span, as opposed to information retrieval or localized reasoning, remains challenging.

5. Ablation Studies Highlights

The report includes several informative ablations shedding light on the impact of specific design choices:

• Local:Global Attention Ratio: Varying the ratio of local to global layers (e.g., 3:1, 7:1 vs. the chosen 5:1) demonstrated minimal impact on model perplexity (Figure 3). This suggests robustness in the architecture and implies the 5:1 ratio was likely selected as an optimal point balancing KV cache savings and potential performance.
• Sliding Window Size: Reducing the window size for local attention layers significantly (down to 1024 tokens) incurred only a minor perplexity penalty (Figure 4). This finding is crucial, as it allows for substantial KV cache reduction without compromising core language modeling ability.
• KV Cache Memory Savings: The interleaved architecture provides dramatic memory savings during inference compared to a standard global-only attention model, especially at longer sequence lengths (Figures 5 & 6). This is the primary enabler for running Gemma 3 with 128K context on accessible hardware.
• Teacher Model Size in Distillation: An interesting finding (Figure 8) suggests that for longer training durations, distilling from a larger teacher model yields better student performance (lower perplexity). This contrasts with some findings from shorter-duration studies where smaller teachers can be optimal, highlighting the importance of training regime length in distillation dynamics.

6. Memorization and Privacy

Addressing memorization of training data is critical for responsible deployment, particularly for open models.

• Gemma 3 models demonstrate a marked reduction in both exact and approximate memorization rates compared to Gemma 2 and other previous models evaluated (Figure 9, note the log scale). This is a significant improvement from a privacy and intellectual property perspective.
• Memorization rates are lowest for the 1B model and increase slightly with model size. Approximate memorization (allowing for small edits) occurs roughly 24 times more frequently than exact verbatim memorization, on average across the models.

{
  "data": [
    {
      "x": ["Gemma 3 1B", "Gemma 3 4B", "Gemma 3 12B", "Gemma 3 27B", "Gemma 2 2B", "Gemma 2 9B", "Gemma 2 27B"],
      "y": [0.0002, 0.001, 0.0015, 0.002, 0.01, 0.05, 0.1],
      "name": "Approx Memorization",
      "type": "bar"
    },
        {
      "x": ["Gemma 3 1B", "Gemma 3 4B", "Gemma 3 12B", "Gemma 3 27B", "Gemma 2 2B", "Gemma 2 9B", "Gemma 2 27B"],
      "y": [0.00001, 0.00004, 0.00006, 0.00008, 0.001, 0.005, 0.01],
      "name": "Exact Memorization",
      "type": "bar"
    }
  ],
  "layout": {
    "title": {
      "text": "Memorization Rates (Log Scale, Illustrative Values)"
    },
    "yaxis": {
      "title": "Rate (%)",
      "type": "log",
       "tickformat": ".5f"
    },
    "barmode": "group",
    "legend": {"yanchor":"top", "y":0.99, "xanchor":"left", "x":0.01}
  }
}

(Note: Y-axis values are illustrative, based on the visual trend and relative differences shown in Figure 9, as exact numbers aren’t provided in the text.)

• Importantly, using Google Cloud’s Sensitive Data Protection (SDP) tool, no personally identifiable information (PII) was detected within the model outputs identified as memorized content. This suggests that the PII filtering applied during data preparation was effective.

7. Responsibility, Safety, Security

DeepMind emphasizes a continued commitment to safety and responsibility, integrating processes throughout the development lifecycle.

• Governance and Assessment: The approach mirrors previous Gemma releases, balancing the benefits of open models with awareness of potential misuse. Risk assessments are conducted, considering the new multimodal and long-context capabilities.
• Safety Policies & Mitigations: Safety filtering is applied to pre-training data. Instruction tuning incorporates alignment with Google’s safety policies (covering areas like hate speech, harassment, dangerous content, child safety, non-consensual sexual content, and promoting illegal acts or severely harmful ideologies) using both supervised fine-tuning (SFT) and RLHF.
• Assurance Evaluations: Models undergo internal safety evaluations using adversarial prompts and human rating to assess policy violation rates, which are reported as low overall. Specific evaluations for knowledge related to Chemical, Biological, Radiological, and Nuclear (CBRN) threats indicate low capability in these high-risk domains.
• Responsible Open Models Approach: A system-level view is advocated, emphasizing that safety depends not just on the model but also on the application’s design and deployment environment.

8. Critical Analysis and Discussion

Gemma 3 undoubtedly pushes the state-of-the-art for open models in its parameter class, particularly regarding efficient multimodality and long context. However, a critical perspective reveals several nuances:

1. The Role of Distillation: The remarkable performance gains, especially for the instruction-tuned models relative to their base counterparts and even Gemma 2, appear heavily reliant on knowledge distillation from highly capable (presumably proprietary, Gemini-family) teacher models. While effective, this underscores that achieving SOTA performance with Gemma 3 is not solely a function of its architecture or pre-training data, but significantly influenced by the quality of the teacher. The specifics of the “novel post-training recipe” remain opaque, limiting full replicability and understanding of the capability drivers.
2. Defining “Long-Context Utilization”: The architecture successfully enables 128K context processing within reasonable memory constraints. However, performance metrics like perplexity and benchmark scores (e.g., the drop on MRCR from 32K to 128K) suggest a gap between processing long context and effectively reasoning over its entire span for complex tasks. The model might excel at retrieval or tasks where relevant information is localized, but deeper integration of information across the full 128K tokens might still be limited.
3. Vision Component Trade-offs: Employing a frozen, relatively small (400M) vision encoder and condensing its output to 256 tokens prioritizes efficiency and simplifies training. However, this may inherently limit the richness of visual feature extraction and the potential for deep fusion between visual and language modalities compared to approaches with larger, jointly trained vision components. The P&S mechanism, while effective, introduces inference latency and complexity.
4. Disentangling Contributing Factors: The report presents numerous improvements simultaneously (architecture, tokenizer, data mix, distillation, RL recipe). While ablations isolate some factors (like attention ratios), fully disentangling the precise contribution of each element (e.g., how much gain comes from the new tokenizer vs. the data mix vs. architectural tweaks alone) is difficult based on the provided results.
5. Evaluation Considerations: While benchmark decontamination is performed, the possibility of subtle information leakage influencing performance on standard benchmarks always exists in LLM evaluation. Furthermore, the strong Chatbot Arena performance, while valuable, reflects preference on a specific platform and may not perfectly correlate with performance on all downstream tasks.

Conclusion

Gemma 3 stands as a robust and compelling addition to the open model landscape. Its intelligent architectural design, particularly the interleaved local/global attention, offers a pragmatic solution to the challenge of long-context inference on accessible hardware. The seamless integration of vision capabilities, enhanced by the practical P&S mechanism, broadens its applicability significantly.

Performance-wise, Gemma 3 models, especially the instruction-tuned variants, demonstrate impressive capabilities, often exceeding models with far larger parameter counts. This success appears heavily driven by sophisticated knowledge distillation and reinforcement learning strategies refined within Google’s broader AI research ecosystem. The significant reduction in training data memorization compared to prior models is a crucial step forward for privacy and responsible AI practices.

Despite reliance on undisclosed teacher models for peak performance and open questions regarding the depth of long-context reasoning, Gemma 3 offers a potent combination of performance, efficiency, multimodality, and improved safety characteristics. It provides a valuable asset for researchers and developers seeking powerful, openly available models capable of tackling a wider range of tasks than previous generations.