Why bother?

I have 3 3090s, and I’ve had them set to 280W basically as long as I’ve had them (I tried a bunch of different wattages, and that’s what I found was pareto optimal for my setup), and was happy with that. BUT! Ever since I moved a few months ago, I’ve been experiencing crashes (system reboots) when trying to use them all at the exact same time: torch.compile across all three gpus for DDP or vLLM warmups (which also torch.compiles), that kind of thing. Staggering compilation and vLLM start times were fine, and sustained 280W loads were fine. 3090s are well known for using a ton of power instantaneously, and today, I decided I’d had enough: it was time to try to fix this.

Shelby’s fantastic post shows us how to cap the average wattage that our gpus use, not going above that wattage. The average part is important: apparently the regulation is over a certain time interval, so instantaneously, GPUs can pull much more power than we set our limit to. Enter undervolting. I’m told that for undervolting, you “shift the voltage-frequency curve so the GPU runs at a lower voltage for the same clock speed” (from Claude). So, at a certain voltage, we can get the same performance. As we know, watts = volts * amps, so we can use fewer watts for the same clock speed. To keep my system from crashing, it turns out that I had to lock my clocks to 1850: anything above that would still crash my machine. This tells me that the problem was probably really a voltage issue (fact check me on this, I’m not sure at all!).

Remember that it’s super important to try a quick training run and monitor the loss and or gradient norms. After I got my setup stable, I had to lower the offset to make sure that I wasn’t getting NaNs. Apparently this is because not enough voltage can make transistors not switch cleanly

I would’ve been happy to just have my system not crash every time I try to use them all at once. But it turns out that I’m actually getting a little more performance!

Here’s what my results look like on my 3090s:

The script

This script handles three things:

1. Power limits — caps each GPU’s power draw via nvidia-smi
2. Clock locking — locks the GPU clock to a range, preventing it from boosting into high-voltage frequency bins
3. V/F curve offset — shifts the voltage-frequency curve via pynvml, so your target clock runs at the voltage normally used for a lower clock (this is the actual undervolting)

Place it at /usr/local/sbin/nv-power-limit.sh:

#!/usr/bin/env bash
# Persistent NVIDIA GPU power limits, clock locking, and undervolting.
# Intended to be run at boot via systemd.
#
# Based on: https://shelbyjenkins.github.io/blog/power-limit-nvidia-linux/
# Extended with clock locking and V/F curve offsets for undervolting.
#
# HOW TO CONFIGURE:
#   1. Set gpu_enabled — 1 for each GPU you want to manage, 0 to skip
#   2. Set gpu_power_limits — desired power limit in watts per GPU
#   3. Set LOCK_CLOCK_MIN/MAX — GPU clock range (prevents high-voltage boost bins)
#   4. Set CLOCK_OFFSET — V/F curve shift in MHz (higher = more undervolt)
#      e.g. +75 means 1850 MHz runs at the voltage normally used for ~1775 MHz
#      Note that the values below are what worked for me on my machine to solve the problem I was having, your mileage will vary.
#
set -euo pipefail

command -v nvidia-smi &>/dev/null || { echo >&2 "nvidia-smi not found, exiting."; exit 1; }

# ── Configuration ────────────────────────────────────────────────────
# Enable/disable per GPU (0-indexed)
declare -A gpu_enabled=(
    [0]=1
    [1]=1
    [2]=1
)

# Power limits in watts per GPU
declare -A gpu_power_limits=(
    [0]=250
    [1]=250
    [2]=250
)

# Clock locking range in MHz (0 to disable)
LOCK_CLOCK_MIN=210
LOCK_CLOCK_MAX=1850

# V/F curve clock offset in MHz (0 to disable)
# Shifts the voltage-frequency curve so your locked clock runs at lower voltage.
CLOCK_OFFSET=75
# ─────────────────────────────────────────────────────────────────────

echo "=== NVIDIA GPU Power & Undervolt Setup ==="

# Step 1: Set power limits
for gpu_id in "${!gpu_enabled[@]}"; do
    if [[ ${gpu_enabled[$gpu_id]} -ne 1 ]]; then
        echo "GPU $gpu_id: skipped (disabled)"
        continue
    fi

    /usr/bin/nvidia-smi -i "$gpu_id" --persistence-mode=1

    max_pl=$(/usr/bin/nvidia-smi -i "$gpu_id" -q -d POWER | grep 'Max Power Limit' | awk '{print $5}' | grep -oE '[0-9]+([.][0-9]+)?')
    if [[ -z "$max_pl" || "$max_pl" == "N/A" ]]; then
        echo "GPU $gpu_id: could not read max power limit, skipping"
        continue
    fi

    desired=${gpu_power_limits[$gpu_id]}
    if [[ "$desired" -le $(printf "%.0f" "$max_pl") ]]; then
        echo "GPU $gpu_id: power limit -> ${desired}W"
        /usr/bin/nvidia-smi -i "$gpu_id" --power-limit="$desired"
    else
        echo "GPU $gpu_id: ERROR desired ${desired}W exceeds max ${max_pl}W"
    fi
done

# Step 2: Lock GPU clocks (prevents boosting to high-voltage frequency bins)
if [[ "$LOCK_CLOCK_MAX" -gt 0 ]]; then
    for gpu_id in "${!gpu_enabled[@]}"; do
        if [[ ${gpu_enabled[$gpu_id]} -eq 1 ]]; then
            echo "GPU $gpu_id: clock lock -> ${LOCK_CLOCK_MIN}-${LOCK_CLOCK_MAX} MHz"
            /usr/bin/nvidia-smi -i "$gpu_id" -lgc "$LOCK_CLOCK_MIN","$LOCK_CLOCK_MAX"
        fi
    done
fi

# Step 3: Apply V/F curve clock offset via NVML
# This requires pynvml (pip install nvidia-ml-py) accessible to root.
# Adjust the PYTHONPATH below if pynvml is installed in a user site-packages.
if [[ "$CLOCK_OFFSET" -gt 0 ]]; then
    enabled_ids=()
    for gpu_id in "${!gpu_enabled[@]}"; do
        if [[ ${gpu_enabled[$gpu_id]} -eq 1 ]]; then
            enabled_ids+=("$gpu_id")
        fi
    done

    # Try to find pynvml — check common locations
    PYNVML_PATHS=(
        "/usr/lib/python3/dist-packages"
        "/usr/local/lib/python3/dist-packages"
    )
    # Also check all user site-packages
    for d in /home/*/.[Ll]ocal/lib/python3*/site-packages; do
        [[ -d "$d" ]] && PYNVML_PATHS+=("$d")
    done
    EXTRA_PATH=$(IFS=:; echo "${PYNVML_PATHS[*]}")

    gpu_count=${#enabled_ids[@]}
    PYTHONPATH="$EXTRA_PATH:${PYTHONPATH:-}" /usr/bin/python3 -c "
import pynvml, sys
pynvml.nvmlInit()
offset = $CLOCK_OFFSET
gpu_ids = [${enabled_ids[*]// /,}]
for i in gpu_ids:
    handle = pynvml.nvmlDeviceGetHandleByIndex(i)
    pynvml.nvmlDeviceSetGpcClkVfOffset(handle, offset)
    actual = pynvml.nvmlDeviceGetGpcClkVfOffset(handle)
    print(f'GPU {i}: clock offset -> +{actual} MHz')
pynvml.nvmlShutdown()
" || echo "WARNING: clock offset failed (is pynvml installed? pip install nvidia-ml-py)"
fi

echo "=== Done ==="
exit 0

Make it executable:

sudo chmod 744 /usr/local/sbin/nv-power-limit.sh

Systemd service

Create /usr/local/etc/systemd/nv-power-limit.service:

[Unit]
Description=NVIDIA GPU power limit, clock lock, and undervolt
After=syslog.target systemd-modules-load.service
ConditionPathExists=/usr/bin/nvidia-smi

[Service]
User=root
ExecStart=/usr/local/sbin/nv-power-limit.sh

[Install]
WantedBy=multi-user.target

Then:

sudo chmod 644 /usr/local/etc/systemd/nv-power-limit.service
sudo ln -s /usr/local/etc/systemd/nv-power-limit.service /etc/systemd/system/
sudo systemctl daemon-reload
sudo systemctl start nv-power-limit.service
sudo systemctl status nv-power-limit.service
# If it looks good:
sudo systemctl enable nv-power-limit.service

Finding stable values

The values in the script above (250W power limit, 1850 MHz max clock, +75 MHz offset) are specific to my setup. I’d imagine they’re a pretty good place to start, as (I think?) they’re pretty conservative, but you’ll need to find your own stable values. You might just have your favorite coding agent write up a script for you to try little increments, logging each combination of these until you crash (DM or email me if you’d like mine). It seems reasonable to me to pick a wattage and go from there, but they’re your GPUs.

Start conservative and work your way up:

1. Set a power limit first and run your workload. Check temps and performance.
2. Lock clocks to your GPU’s typical boost clock (check with nvidia-smi dmon -s c).
3. Add clock offset in small increments (+50 MHz at a time). Run a stress test at each step (eg vLLM warmup/serving).

If the system crashes under load, back off the offset. If it crashes at idle, your minimum clock might be too low for the offset you’re applying.

Modifying settings

After editing the script:

sudo systemctl daemon-reload
sudo systemctl restart nv-power-limit.service
sudo systemctl status nv-power-limit.service

The plan is to pretrain a fully distilled NanoGPT-aloid using off-policy distillation, starting with OLMo 3 7B’s base, instruct, and thinking variants to see if post-training variant makes a difference, then testing quantized teachers, and finally moving on to larger, smarter models like OLMo 3 32B, Qwen, GLM 4.7 Flash, and GPT-OSS 120B with whatever setup wins. We’ll have to switch tokenizers along the way, but that seems like an ok trade-off.

Questions

1. Does the teacher’s post-training variant matter (base vs instruct vs think)?
2. Does quantizing the teacher matter?
3. Does the teacher model and size matter?
4. Does tokenizer mismatch matter?

Model Setup, Baseline

The model we’re using here, as alluded to, is a NanoGPT-aloid, based on the Olmo-3 architecture, just shrunk way down. I thought it’d be fun to use GQA too. | Parameter | Value |
|—|—|
| hidden_size | 768 |
| num_hidden_layers | 12 |
| num_attention_heads | 6 (Q heads) |
| num_key_value_heads | 2 (KV heads, 3:1 GQA ratio) | | head_dim | 128 (768/6) | | intermediate_size | 2048 (SwiGLU MLP) | | max_position_embeddings | 2048 | | vocab_size | 100,278 (OLMo tokenizer) | | tie_word_embeddings | False |

That gets us to ~75M parameters in the transformer core and ~154M embedding parameters (hahaha), for a total of ~229M. To get a nice baseline, even though the model is dominated by the embedding parameters, we train on 5B tokens — a little over Chinchilla-optimal for the total parameter count, ignoring the fact that so many parameters come from the embeddings. That gets us down to a val BPB of 0.96437.

Experiment Plan

TODO

• Calculate teacher FLOPs per token for each model once we settle on token budget (chinchilla optimal for 125M NanoGPT is ~2.5B tokens, but distillation should need far fewer)

Results

The baseline model is defined as follows:

I’ve been working on variable-sized experts in MoEs (previous post) using a modified version of Megablocks. The TL;DR from that work: at my scale, I didn’t find efficiencies beyond what you’d get from simply using narrower experts across the board. But since I have this hammer, I’ve been looking for nails.

I’d had my eye on doing a project with Ai2’s FlexOlmo for a while, and it seemed like a perfect nail. The core idea of FlexOlmo is to train specialized experts separately on their own domains, then combine them into an MoE¹. They got good results, and the architecture opens the door to data collaboration. Organizations can train experts on private data without surrendering it, then combine those experts into a more performant MoE model without leaking sensitive information. But training a 4.3B expert isn’t cheap, and in data constrained situations, it doesn’t necessarily make sense to train a model that large. If smaller experts work, that would lower the barrier to participation significantly.

Since I have limited resources, I figured the best way to do this was to prune the expert MLPs of one of the existing expert + public model models that Ai2 released, and use distillation to retrain the model, like these two NVIDIA papers (Muralidharan et al., Sreenivas et al.), as opposed to training new experts from scratch.

So, I took Ai2’s released math expert, pruned its MLPs to various widths, and retrained with knowledge distillation to see how small the expert could get while still contributing to the combined model.

In ~228M tokens of retraining with KLD distillation, I’m happy with the results, and it is a solid proof of concept. Even pruning the expert down to ~800M parameters total improves the Math2 score from 8.1 to 29.1.

Pruning the Math Expert: Shrinking the Expert MLP Layers

I pruned the Flex-math-2x7B-1T math expert to three widths: 8192, 5504, and 2048 (down from 11008). Since the hidden size of the base model has to stay the same, and I didn’t want to touch the attention heads or number of layers, I shrank the expert MLP layers. One thing I’d like to try eventually: pruning different experts by different amounts within the same MoE, with sizing informed by per-expert importance scores rather than arbitrary uniform targets. Variable expert sizes make this straightforward.

Importance Analysis

For the 2048 width model, I wanted to test how much the dataset used for importance analysis (the step where we decide which neurons to prune) matters. Not necessarily shocking, but it turns out quite a lot! I tried two datasets: a subset of the math data from dolmino-mix-1124², and general data from the same dataset, which does include some math.

58% of the top-2048 most important neurons are different between the two analyses. The early layers mostly agree on what’s important, but from layer 6 onward the rankings diverge.

The model calibrated with math data also trained more effectively, achieving consistently lower training loss throughout. Validation loss followed the same pattern. Note that I did stop the general training run early, but it wasn’t going to catch up, and I wanted to move on to training the larger models.

Here are the loss curves of the math and general models, which show that the model calibrated with the math dataset is the clear winner.

Distillation

To be explicit: the teacher is the full-sized Flex-math-2x7B-1T; the student is the pruned model. These models were retrained using distillation with the top 128 logprobs generated by Flex-math-2x7B-1T using the GSM8k, Metamath-owmfilter, and TuluMath subsets of the DOLMino mix dataset (the same dataset that FlexOlmo was trained with), about 620K total documents. Logprobs dataset here.

Performance vs Baseline

Evals

Using LM eval harness (which handles base models well), we got pretty close to the paper’s baseline numbers. Math2 is the macro average of GSM8K and MATH³.

The 8192 model is juust about on par with the full-sized expert. Even the 2048 model (0.8B expert params) scores 3.6x the no-expert baseline. The half-sized expert (5504) is pretty competitive with its larger siblings.

It’s also worth noting how much distillation matters: the pruned-only 2048 model (no retraining at all) scores just 13.1% on GSM8K and 3.3% on MATH, barely above the no-expert baseline. Distillation recovers it from near-broken to 44.3% / 13.9%, a massive improvement for only ~228M tokens of training.

Takeaways

I think this is a pretty good nail!

• I’m particularly impressed with the 2048-width model. Adding just 800M parameters makes the model almost four times as good as the baseline model! I think that seems like very good bang for your buck, especially for training on so few tokens.
• Prune + distill is a good path to making smaller FlexOlmo models
• Choosing the importance analysis dataset wisely can make a pretty substantial difference in the overall performance of the model.
• Smaller experts could make the data collaboration vision of FlexOlmo more viable

Tentative Recipe for Training New FlexOlmo Experts (Untested… for now)

Based on what worked here, I think the recipe for training new FlexOlmo experts would look something like:

1. Do importance analysis on the Flex-public-7B-1T model with the target dataset
2. Prune the MLPs to the desired width
3. Attach those to an untouched public model (as described in the FlexOlmo paper)
4. Train only the expert’s MLP, either with regular old cross entropy loss on the dataset or, even better, using KLD distillation from a strong teacher model.

Limitations

I only evaluated on math benchmarks (GSM8K and MATH). It’s possible that pruning the expert hurts general reasoning or other capabilities that I didn’t measure. Running BBH was going to take like 60 hours on my home system, so I figured I’d just punt and do these. Since we’re pruning a math expert and testing math performance, I think the evals here are the right ones, but broader evaluation would be nice. I also really want to know how FlexOlmo works with post-training. Can we mix a post-trained public model and an expert model with just continued pretraining? Or just post-train the expert model?

Whole wheat bread gets a really bad rap. I think this is for two main reasons:

• A lot of whole wheat flour is super old. The germ and bran get rancid, which tastes gross
• Bad texture. The bran, since it’s hard and sharp, pokes holes in the gluten, making the bread dense

So, can we fix these problems? The first one is easy, just use great flour. Amazing mills are popping up all over the country, but I’ve used flour from Cairnspring and Camas Country Mill and loved everything I’ve made with them. You can probably find a flour near you.

On the texture point, if we just treat the flour as flour, we’re probably not going to get the results we want. Flour is made up of three main things: the bran (outer layers of the wheat, makes up ~14.5% of the wheat kernel), germ (embryo, ~2.5%), and endosperm (starchy part, what we think of as white flour, 83%).

The insight here is that if we treat the germ and bran (hereafter called offal) as we would with any other grain mix-in, we can get good texture AND pure whole wheat flavor.

The flour I’ll be using is Fortuna wheat from Lopez Island. It’s probably a little low on protein, but it tastes great.

Flour makes a difference!

In late 2024, I ordered a bunch of really cool flours from Camas Country Mill, including this Amarillo flour. I made some bread with it, and was pretty blown away with the color and flavor.¹ I hadn’t really thought about flour having distinct flavors until I tried these flours.

7/2: Just soak the bran

I haven’t been taking pictures, but I’ve been having great results just soaking the bran like we would with any other grain. The major key here is just separating the bran from the flour, softening it, and incorporating it later. shrug

1/16-1/17: Scalded offal, 15%

I did the following: Scald

• 135g wheat offal (15%)
• 340g boiling water (38%) (2.5:1 ratio of water to bran) Pour water over bran, stir until combined, cover, cool completely

Dough

• 900g sifted Fortuna flour (sifted through 40 and 50 mesh sifters)
• 630g water at 80-85°F (70%)
• 200g leaven (22%)
• 18g salt (2%) 1 hour autolyse, 3 hour bulk ferment, stretch and fold every 30 minutes. Scald added after an hour and a half.

Shape, overnight proof in fridge.

I was really happy with this! Texture was great, flavor was great. No pictures this time. Pressing on to higher percentages!

I’ve been working on pruning models a bit lately, mainly based off of this Nvidia paper. Mistral recently launched their Ministral 3 models, which were pruned from Mistral Small 3.1. In the Nvidia paper, they find that pruning width keeps performance better than pruning depth¹.

What I think they don’t address enough is the fact that each additional layer adds latency. Since layers can’t be computed in parallel, latency increases substantially. Here are some plots illustrating just that. I focused on Baguettotron as a comparison: at 80 layers with a hidden size of 576, it’s such an extreme aspect ratio that I thought it would make for a good test case. Life is about trade offs: for many edge use-cases, we’re probably concerned with the most capable model that fits on the device, not necessarily raw throughput. So, I had Claude Code whip up the comparison code and make some plots (I was really disappointed with the plots that it came up with, so I ended up telling it which plots to make).

All experiments were done on one 3090 with randomly initialized models. All the code is here, and was written by Claude Opus 4.5 in Claude Code.

For a more thorough study of these trade-offs, see Nvidia’s paper Nemotron-Flash: Towards Latency-Optimal Hybrid Small Language Models.

Depth vs Width Scaling

Token generation latency grows linearly with depth, but stays flat even at triple the width.

Parameters vs Latency Scatter Plot

This plot really shows us that for total response time (which is what we’re calling prefill + token generation), the number of layers matters much more than the number of parameters. Note that the largest model among the parameter matched configurations we tried, with 8 layers, has 487M parameters and is 11.7x faster than the 100 layer model, which has 322M parameters.

Time Matched Comparison

Here, we see that even at maximum width (before running out of VRAM), shallower models are both faster and have dramatically more parameters. The dashed line shows the 80-layer baseline—none of the wider models could be slowed down enough to match it. Since the 11.3B models were as large as I could fit on my 3090 (24 gigabytes of VRAM) in bf16, that’s where we maxed out. Clearly, there’s still a fair amount of headroom to make these larger while being faster.

Baguettotron vs Gemma 3 12B

Here’s a real-world aspect ratio comparison: Baguettotron vs Gemma 3 12B. We still see a pretty big difference!

I want to write up my work as I’m working on it, and not all at once. In that spirit, before I start working on pruning the FlexOlmo models to try using them with variable sized experts, I want to get a baseline of pruning models. I’m going to prune Olmo 3 7B. I think I’ll just prune it down to 1/2 of it’s size. I’ll be implementing basically these two Nvidia papers (they basically follow the same process):

• Compact Language Models via Pruning and Knowledge Distillation
• LLM Pruning and Distillation in Practice: The Minitron Approach

Experiments

I’ll run a few experiments:

• KL distillation vs on-policy distillation with reverse KL a la this Thinking Machines Post basically: rollout with the student, forward the rollout through the teacher, matching the scores with reverse KL Loss.
• Quantized teacher models (8 bit, 4 bit, some weird GGUFs like Q6?)
• Stronger teacher models (possibly quantized stronger teacher models?). The papers on pruning I’m going to be working with all use the same model they’re pruned from as teacher. The Ministral 3 paper used Mistral Small 3.1 for the teacher model the whole time. It’d be nice (I think) to try an even stronger model than the parent as the teacher, seems like a no brainer? Can we use models that used different tokenizers? I think I saw somewhere that it’s doable, but can’t remember where.
• Olmo uses sliding window attention (SWA) as well as full attention. Will pruning the full attention layers lead to worse performance? Probably check RULER.
• LoRA or DoRA (should punt on this one for now, seems like something I should do later)

Note that only pruning width seems to slightly defeat the purpose of making smaller models, in my opinion. Nvidia has released a paper on exactly this: Nemotron-Flash: Towards Latency-Optimal Hybrid Small Language Models: “While previous work on SLM design has primarily focused on reducing the number of parameters to achieve parameter-optimal SLMs, parameter efficiency does not necessarily translate into proportional real-device speed-ups…we first study latency-optimal depth-width ratios, with the key finding that although deep-thin models generally achieve better accuracy under the same parameter budget, they may not lie on the accuracy-latency trade-off frontier.” Here’s a quick study comparing latencies across different aspect ratios.

Results (So Far)

Taking a break on this while I work on Variable FlexOlmo. So far, I’ve pruned and distilled one version of Olmo-3-7B Instruct, my smallest version, about 3.5B parameters. On ~500M tokens (one epoch through our dataset), we got the following performance on 50 questions of GSM8K:

So, we’re not making a SOTA 3.5B model here. But in only 500M tokens, we do legitimately recover some performance. I also want to dive a little deeper into just what happens when we do prune. It’s interesting to me that the stop token seems nowhere to be found for our pruned models. One only pruned model, which was just barely pruned, kept repeating “grammar” consistently across eval runs. 22 out of 50 samples in that run collapsed into “grammar grammar grammar…” until it reached the token limit. I really want to look into this: what behavior do we see when we prune different parts of the model? Might be a pattern there. Might not.

For now, all the code’s on github, and all the checkpoints and dataset are on Hugging Face.

Thanks to:

• Sam Lehman for the conversations that got this project off the ground as well as reading drafts and feedback
• Craig Robinson for feedback on drafts
• Andrej Karpathy for NanoGPT, which this project is based on
• Trevor Gale, Deepak Narayanan, Cliff Young, and Matei Zaharia for MegaBlocks, the other foundation for this project

Extending Megablocks with Variable Sized MoEs

I implemented variable sized experts using Andrej Karpathy’s nanoGPT, which allows us to set the sizes of experts in a mixture of experts (MoE) model. I trained a bunch of variable sized expert MoEs averaging around 125 million active parameters on a chinchilla-optimal ~2.5B tokens. I expected tokens to route based on computational difficulty (something like difficult reasoning to large experts, simple concepts to small ones). It turns out that tokens in constrained contexts like code or recipes route to small experts, and more ambiguous function words like ‘ with’ and ‘ to’ route to larger ones. My interpretation is that large experts handle tokens that need more context to interpret, while small experts handle words with specific meanings. Check out the visualization of where different tokens go here, and code here!

Stage 0: MoE Background

Dense LLMs have a feedforward network (FFN)¹ that every token passes through entirely. MoEs replace this with a much larger layer broken into pieces called “experts”. In an MoE, each token only activates a few experts. This way, MoEs can be more capable, while processing at the same speed.

Historically, MoEs have only been made up of uniform experts of equal size. Here, I explore what happens when we vary the size of the experts. For this project, the ratio (5:1, 23:1) is the size difference between large and small experts. A 23:1 model’s large experts are 23 times the size of its small experts. The two models I trained have the following attributes:

5:1 configuration: Four experts at 2560 hidden dim, four at 512. 2 experts are active². Empirically, active parameters range from 95.7M to 138.7M and we average 114.7M.

23:1 configuration: Four experts at 2944 hidden dim, four at 128 from 91.5M to 114.5M, 101.9M on average.

Record scratch

But first let’s talk about how I got myself into this situation.

Stage 1: Vanilla MoEs

I started this escapade with a simple goal: add efficient³ MoE support to Karpathy’s nanoGPT. I found all the regular things that we’d expect to find with MoEs: they perform better, and for the most part, the sparser, the better. The best performing regular MoE was 64 total, with 8 active (gpt2 wandb, wikitext wandb). It reached a final loss of 3.127 on Openwebtext on a little over chinchilla-optimal 3 billion tokens, compared to 3.285 for the dense model. The 64 total, 8 active model beat the dense model’s loss in 70% of the steps. However, due to (I think) memory overhead, the training run took twice as long (20m for dense vs 40m for the 64x8 MoE). 70% of 40 minutes is 28 minutes, so even if we adjust for hitting the same loss, the MoE doesn’t train as fast at this small scale⁴. In the OLMoE paper, the AI2 team got to the same loss 2x faster compared to the dense model they trained.

This MoE mission opened up a whole, unexpected, gigantic can of worms.

Stage 2: Variable sized experts

Part of the quest of not using for loops led me to MegaBlocks (github and paper). The main MegaBlocks innovation is that we can break matrix multiplications up into blocks of 128, allowing us to avoid dropping tokens when experts get too many, and use less padding when experts get too few. They find that dropping tokens leads to substantially worse performance. The paper goes on to say that “we could also relax the constraint on the number of columns in each block to build MoE layers with variable sized experts”. This, along with conversations with my friend Sam, and two Dwarkesh Patel interviews (Jeff Dean and Noam Shazeer, Sholto Douglas and Trenton Bricken) led me down this rabbit hole of variable sized experts (By no means am I saying that this is exactly what Jeff or Sholto had in mind). Notably, Meituan’s LongCat explores this by using identity experts that route some tokens through zero computation experts. Here, I look at the continuous case: experts of different sizes rather than on/off.
It’s clear that not all tokens are equally hard to predict. So then, why do MoEs have uniformly sized experts? If we can allocate more FLOPs to the harder tokens (whatever that ends up meaning), then maybe we can get a more efficient model⁵. Or maybe we can learn something about what “harder” even means to a model.

So, I started working on this concept, only tweaking MegaBlocks slightly. The fork of MegaBlocks is on my GitHub. I started training some models with variable expert sizes, at first only using a traditional load balancing loss and a router z loss. With a load balancing loss as low as we used for the vanilla MoEs, tokens routed to large experts disproportionately, just as you might think. So, I tried turning up the load balancing loss to get better balance. The best performing vanilla load balancing loss was with an lbl weight of 0.1, a factor of 10 higher than I found was optimal for the vanilla MoEs (wandb). Some great learnings came from that, but these models were well balanced between all the experts, meaning that they routed 50% of tokens to the large experts, and 50% to small. Since I was having the large and the small experts add up to the same intermediate size as the vanilla MoEs, this just averaged the same size of experts as the vanilla MoEs. It performed ever so slightly worse, training in about the same amount of time: not so exciting. We’d need a loss term that lets the model be more free to choose whichever experts it wants.

Stage 3a: Compute balancing loss

This is very similar to our normal load balancing loss. For this, we:

1. Normalize the expert sizes by dividing the expert sizes by the average size of the experts
2. Compute the weighted sum of the router probabilities and the normalized expert sizes (this is a dot product between the router probs and the normalized expert sizes)
3. Take the average of that!

This setup did allow the model to learn to allocate different size experts to different tokens. Both the 5:1 and 23:1 models ended up at around 70% small and 30% large experts. After training, there were still load imbalances within groups of experts, so we needed to add slightly modified load balancing loss back in.

Stage 3b: Group load balancing loss

This is pretty simple: we just divide up the expert groups, and use the same load balancing loss that we used above and got from OLMoE. This new scheme balanced well at weight 0.01.

So now, I’d done it! We have variable sized MoEs and can train a bunch of size ratios.

Stage 4: Training a bunch of these models

I started on a smaller model size to sweep over the size ratios, training roughly 50M active models on Wikitext. The most promising runs were at 4:1, 6:1, and 19:1 (large:small). The 19:1 used the smallest experts MegaBlocks allows: 128 hidden dim. The 19:1 model finished training in 84% of the time with only 3.2% higher loss. The 4:1 finished in 91% of the time with 0.96% higher loss; the 6:1 in 89% with 1.4% higher.

So, with the hypothesis that something around 4:1 or 6:1 would be the optimal “larger” ratio and the extreme “smaller” ~19:1 ratio in mind, I went up to GPT-2’s size of about 125 million active parameters. I stuck with openwebtext as a fair comparison to the original models I trained. I trained the final models with a load balancing loss weight of 0.08, and a compute loss weight of 0.004⁶. I found very similar results across the board: the extreme ratio (this time 23:1, again with small experts of size 128) was much faster to train, this time with a 20% speedup and a 2.5% loss degradation. The average size of experts in evaluation was 861. The more moderate 5:1 ratio was 10% faster, with 0.5% loss degradation, with an average size of 1077.

One ablation I ran was with consistently sized experts, but much smaller than 1536. I chose 896 because it was the closest multiple of 128 to our 861 average size. Bluntly, this run was disappointing to me: it beat all but the 5:1 ratio on loss, and finished training faster than the 23:1 ratio runs, though they were within spitting distance of each other (Update: a couple weeks later after going over my routing layer by layer, I noticed that the first layer was extremely imbalanced within size groups. Fixing this for 23:1 runs actually ended up matching the loss and speed. I think the takeaway is still the same though). This shows us that the active MLP width need not be 4*hidden size. We already see this in most of the best-performing open source MoE models: Deepseek V3 and Kimi K2 (which both use Deepseek’s basic architecture) use experts that add up to an expansion factor of ~2.57, and GLM-4.6 uses 2.7. (Minimax M2, gpt-oss and OLMoE all use 4x).

Efficiency aside, this was the real question I wanted to answer: do the models learn meaningful routing patterns?

Stage 5: contextual ambiguity, routing analysis, domain specialization, and syntactic specialization

This was what I was most excited about: do experts actually specialize? There have been conflicting reports. The original MoE paper uses LSTMs and tons and tons of experts (up to 131,072!), and shows clear specialization. Then, the ST-MoE paper shows that the encoder specializes, but “expert specialization is far less noticeable in the decoder” (recall that most of the generative models that we think of are decoder only). The Mixtral paper shows no domain specialization at all, but does find syntactic specialization⁷. OLMoE shows that certain tokens, like those from arxiv, consistently route to certain experts. Finally, Deepseek MoE) makes the claim that their MoE performs better than GShard due to more expert specialization. Their focus is mostly on redundancy across experts, not individual token routing analysis or domain specialization. All of this is ok evidence that there probably is some kind of specialization going on with MoEs, but not conclusive at all. My holy grail would be to one day figure out that specialization and show that if we turn off certain expert combinations across layers (say, if there are experts that get lots of code tokens), performance tanks. Additionally, if we isolated those experts into a dense model, would we be able to keep the performance in that area? To me, this would be the clearest evidence that the experts do specialize. I don’t achieve this here.

The 23:1 and 5:1 models learn pretty different routing schemes. For starters, the 5:1 ratio has a 0.295 spearman correlation⁸ between the token count in the dataset and the average size of the expert. On the other hand, for the 23:1 ratio, we see a negative correlation: -0.222. So, I’ll break them into two sections. I expected that because of the extreme difference in experts, we’d see a clearer pattern. That doesn’t seem to be true. The 5:1 model’s highest vs lowest expert sizes is actually larger than the 23:1 model’s (the delta is 2332 for 5:1 and 1252 for 23:1). The average size of the experts and entropy for the 5:1 ratio are not at all spearman correlated (literally 0.00), and the 23:1 is weakly negatively correlated (-0.150). So what are they even learning?

I think the models I’ve trained here learn the contexts that they are in: when they are in a constrained context like programming, a recipe, or finishing a common word, they tend to use smaller experts. When the context is more ambiguous, they use larger experts.

If we look at different counts of tokens, we see that the two ratios learn different patterns by frequencies. When we get to more frequently occurring tokens, there is higher correlation between the two mean sizes, meaning that there is some (small) relationship between the sizes that the two learn (see the table below). So they don’t learn totally orthogonal patterns, but they definitely learn different ones. For example, “ing” words go to large experts for 5:1 and small for 23:1. Let’s take a closer look at the patterns.

5:1

We see a pretty clear pattern with the 5:1 model, especially on the low average expert size end.
Technical tokens like “ goto”, “ perl”, and “ println” are all at the low end of the compute spectrum. There are the second halves of somewhat unique words like “oenix”, “ibrary”, and ”chnology”. We also see “ tbsp”, “ teaspoons”, and “ tablespoons”, as well as “ lol” and “ haha”. Additionally, we see discourse markers like “Furthermore”, “Moreover”, and “Similarly”. These are all capitalized! When we look at the lowercase versions of these tokens, we see that they actually use larger experts. Most extremely, “ similarly” is actually in the top 2400 out of around 50 thousand. It’s particularly syntactically ambiguous: there are a bunch of different ways to use the word. When it’s at the start of the sentence, chances are very good that it’ll be referring to the previous sentence.

On the other end of the scale, the tokens that use the largest experts are the token ids <564> and <447>. These are the only tokens that use larger experts on average than our baseline⁹.

Now, if we try to decode these tokens just in the python interpreter, they don’t print correctly. That’s because they are both incomplete byte strings– some token comes after them. So, decoding them as bytes, we see that 447 is ‘e280’, and 564 is ‘20e280’. It turns out that ‘20’ is just a space, so 564 is basically ‘ e280’. If we dig deeper, we see that e2 80 is the beginning of 64 different types of punctuation including “—” (em dash), “‡” (double dagger), and my favorite, “‽” (interrobang) (see full list here), along with a bunch of different spaces. Another fascinating thing about these tokens is that they get down to a very low entropy. Since there are only 64 options that can follow the ‘e280’ bytes or else there will be a catastrophic failure, the model seems to learn that it only really has 64 options, leading to the low entropy. Then, it needs big experts to decide which of the 64 possible tokens come next, since there’s a lot of ambiguity for which symbol should come next. Also appearing on the high side: function words like “ with”, “ to”, and “ in”.

What I think is going on: these models are routing tokens based on contextual constraint—when what comes next is predictable, they use less compute. Within the low compute regimes, the context is very constrained. For example, with the technical words, there are only so many things that can come after “ println”, or with “ tbsp”, we’re almost certainly in a recipe setting. Lastly, with the discourse markers, I think they kind of ‘tell you what’s coming’: another, similar point.

23:1

As I said above, this model’s routing patterns are pretty different. It’s slightly negatively correlated with token occurrences in the dataset, and again very slightly negatively correlated with the output entropies (-0.15). The two aren’t correlated with each other. There are patterns at the low end: we see lots of names and other proper nouns “ Abraham”, “Stephen”, “ Jon”/“John”, “ Obama”. There are some code-related tokens as well like “ !=”, “ ();”, and “ ()”, and a lot of numbers. At the high end, it’s a lot more of a grab bag, but we do see some similar word endings and whole words (though I feel like I’m trying to cram my previous observations onto this). We do see token 564 show up again, and 447 is in the upper quartile.

Domain Specialization

Joining the two ratios back together, I tested domain specialization by running both models on different domains from the Dolma 3 pool. The only meaningful domain effect for this setup is code.

For the 5:1 model, code’s weighted average expert size is 2244, which is 210 smaller than the 2454 size observed across all other domains. The 23:1 model is even more extreme: code drops to 1499 vs 2131 for other domains, a gap of 632. This is 14.3 standard deviations below the cross-domain mean! When given more dynamic range, the model leans harder into treating code as low-compute. Code’s maximum expert size is 2958 (5:1) and 2916 (23:1), while natural language domains reach 3697-3755. Of 690 tokens appearing across all domains with count >100, 518 route lower in code and only 165 route higher (the remaining 7 were the same). Tokens like “ is”, “ in”, and “ return” (all programming keywords) drop 250-450 in code. But the pattern flips for subword fragments: short tokens like ‘ker’, ‘ame’, and single characters route higher in code, since they could be part of any identifier.

The other domains cluster together. Politics and crime & law route slightly higher (+30 and +26 vs the cross-domain mean), sci_math_tech slightly lower (-31), but this ~60 point spread is noise compared to code’s ~600 point gap in the 23:1 model.

Wrapping up, Open Questions

At this scale, we don’t see any efficiency gains other than what you’d get from only using smaller experts. Routing is consistent across seeds: function words go to big experts, and code goes to smaller ones. I expected “difficult” tokens (code, math, technical terms) to get more compute; that’s not what happens in my experiments.

Open questions and things I want to try next:

1) I haven’t yet trained many models with better-performing expert setups than 8x2 (eg 64x8). The size differences with more experts were less extreme, so I thought 8x2 would be more revealing to study.
2) Imbalanced expert counts, shared experts. Deepseek V3 and other models based on its architecture, like Kimi K2, use a shared expert, one expert that’s always active. Could we use one large expert as the shared expert, and the rest small experts, or something like this idea? Is there an optimal number of experts of each size (something like, without loss of generality, 2 large, 6 small)? Could we do small, medium, and large?
3) FlexOlmo worked on the concept of training dense models separately and then stitching them together – could we use these variable sized experts to have smaller specialized experts with a larger, more general model as the backbone?
4) Beat the baseline by tuning the compute balancing loss coefficient
5) How’s this stuff interact with attention scores? Especially with the discourse markers thing. Perhaps a “Moreover” token is attending to the last sentence quite a lot and this lines up with the lack of compute?
6) Intermediate ratio analysis– is there a crossover point somewhere between the 5:1 and 23:1 where we see the correlations flip?
7) This last weekend (Dec 13/14 2025) I tried subbing in a sigmoid instead of softmax for routing, like they do in DeepSeek-V3. It learned more like an 80:20 small:large ratio. I haven’t had a chance to look at the routing. Will everything we saw above hold?
8) RL’s been used for gating– obviously that hasn’t stuck around. Would it help us allocate our compute better in some way?

Adding Mixture of Experts support to Karpathy’s NanoGPT, MoE interpretability, general MoE Notes

In progress: July 14, 2025-?

Andrej Karpathy’s NanoGPT is a hackable library for training language models. In his inimitable style, Karpathy shows anyone who wants to learn exactly how pretraining for LLMs is done. Here, I’d like to add support for Mixture of Experts (MoE) style models. Over the next (generic period of time), I’ll be working on learning more about MoE models. Extending NanoGPT with MoE support feels like a good place to start. I’m also interested in upcycling something like SmolLM, Not 100% sure how much compute that would take. Additionally, I’m fascinated by what’s really going on inside these kinds of models. Are they actually learning some sort of expertise? For example, in a given MoE model, is there some notion of a “math expert”?

What is an MoE? (7/28/25)

Let’s back up. What is an expert in the context of LLMs? Each layer of a standard transformer model consists of two main parts: the attention block, and the feedforward or MLP (multilayer perceptron) block. The feedforward block is what we modify in MoE models. In a regular (most commonly known as dense) transformer model, this feedforward block is a fully connected (linear) layer that expands to 4x the hidden size of the model. Then, an activation function (most often SwiGLU or GeGLU now, GPT-2/NanoGPT uses GeLU) is applied. Finally, we use another fully connected layer to bring us back down to the hidden size of the model.

MoEs follow this same process, but instead of having one large linear layer, they use a set of smaller ones, known as experts. We also add in a fully connected layer before the MLP block, known as the router. As the name suggests, the router is responsible for deciding which experts should be active for a given token.

A few thoughts on this right off the bat (7/14/2025):

1. The first paper on MoEs for language modeling from Google, Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer (see Table 9) showed that they could practically just read the experts off from the model weights (we see this in vision MoEs too). A few caveats below.
   • They were working with LSTMs
   • There was what amounted to one expert layer
   • They trained models with up to 131k experts! (In table 9, they look at the model with 2048 experts. Kimi K2, a model with one trillion parameters, has just 384 experts per layer.)
2. The Mixtral paper reports basically no specialization
3. The OLMoE paper shows that, in the first layer, tokens from arxiv are disproportionately routed to one expert in particular.
4. The notion of MoEs being, with Mixtral as an example, 8-7B parameter models “stapled together” is not how they work. For each layer, two experts will be active – they can use a different combination at each layer, so we won’t really be able to call (without loss of generality) expert 3 the “chemistry expert”. It’s possible we may find that a combination of experts across layers (think expert 2 in layer 0, expert 4 in layer 1, and expert 2 in layer 2) do make some sort of “expert”.

Block-Sparsity (7/28/25)

The vanilla Hugging Face transformers version of MoEs loops over the experts. They’re fixing this, but for now this is not really workable if we want to get the training efficiency gains that MoEs are well known for. OLMoE’s paper reports that their setup uses ~3x fewer FLOPs than the dense comparison, which equated to ~2x faster training.

So, we turn to Megablocks. Megablocks, using some clever tricks, grants a huge speedup over a for loop version. Since we’re dealing with matrices, it’s totally possible to parallelize the computation of the experts by essentially stacking them into one big tensor. This comes with a two big drawbacks:

1. Doing the calculation as one gigantic dense matrix multiply is very expensive and, since only a subset of the experts are active per token, it’s wasteful.
2. Naively, the matrix that each expert sees has to be the same size if we want parallelism. This runs us into two more problems: a. If an expert isn’t used much by a certain batch, we have to pad the token matrix, wasting resources b. If a batch of tokens is disproportionately routed to an expert, we have to drop some of those tokens, which hurts accuracy and wastes resources. Megablocks solves both of these problems at once by still having a large matrix of the experts and computing them all at once, but only computing the parts of the matrix that we need to, leaving the rest of the matrix filled with zeroes (this is known as a sparse matrix).

Here’s my first pass at the Triton kernels to do the forward pass. This first one corresponds to the matrix multiplication that takes the batch of tokens and multiplies it by the expert matrix. There is a little bit of tensor manipulation that we do in PyTorch to put everything in the correct place and get the right shapes; see the full implementation if you’re interested.

The second takes that sparse matrix and sends it back to the original hidden size of the model:

Variable Sized Experts (9/12/25)

I’ve been working on these kernels for a long time!! Finally almost there. Quite a few rewrites to really understand what we’re doing. We are storing everything densely, and just keeping track of how many blocks each expert gets, and a cumsum of that to remember the offsets. This same concept should work for variable sized experts… we allocate parameter tensors for the total d_ffn size regardless, so as long as we keep track of where each expert is, it should be “trivial”[^1] to have variable sized experts.

(10/25/25) This is now working quite well, I’ve trained a bunch of 125m (average) active parameter variable sized MoEs. They perform about the same as the same-sized vanilla (uniform expert size) MoEs. Quick thoughts on a FlexOlmo-like project but with variable sized experts:

1. Do simple LoRA on OLMo 1B for domain expertise (or just use AI2’s– they publish them. However, they’re 7B models. Too big?)
2. The main model’s MLP layer is something like twice the size of the auxilary models (alternatively make the aux models’ MLP half the size)
   • This probably looks like copying the MLP weights, freezing everything else, and then only training the MLP layer similar to upcycled MoEs?

2 Quick things for me to remember (9/2/25)

• Look into expert choice vs token choice. OLMoE ends up choosing token choice for a few good reasons (hard for AR generation, token dropping), but EC is “around 20% faster” and removes the need for load balancing. Additionally (this is very interesting!), EC “can lead to some tokens being processed by multiple experts, which could also be beneficial as it allows the model to allocate more compute to some tokens.”
• From James Betker’s excellent Non_Int blog: “The fact that MoE has great scaling properties indicates that something deeper is amiss with this architectural construct. This turns out to be sparsity itself – it is a new free parameter to the scaling laws for which sparsity=1 is suboptimal. Put another way – Chinchilla scaling laws focus on the relationship between data and compute, but MoEs give us another lever: the number of parameters in a neural network. Previously compute and quantity of parameters were proportional, but sparsity allows us to modulate this ratio.” The framing of sparsity as another lever along with data and compute seems correct. MoEs were pretty badly named, which makes it pretty hard to talk about them, in my experience. Even after thinking about them as my main non-work project for a while now, I still have the tendency to think about them as a bunch of llms all stapled together.

[^1] It never is.

Huge thanks to:

• Anthropic for Scaling Monosemanticity, Towards Monosemanticity, and Update on how we train SAEs. This work is based directly on these three documents
• AI2 for training and open-sourcing OLMo 2
• Hugging Face for the Fineweb dataset and for hosting the demo, dataset, and weights
• Sam Lehman for reading drafts and providing feedback
• Eddie Austin From Nurturepoint.ai for reading drafts and providing feedback
• The open-source interpretability community, especially those sharing SAE implementations and techniques

If I missed anyone, my apologies! Happy to update this as needed.

Motivation

Last year, Anthropic demonstrated something magical: for 24 sublime hours, they released “Golden Gate Claude”, a version of Claude that couldn’t stop talking about the Golden Gate Bridge. Ask it what its physical form is? It would respond “I am the Golden Gate Bridge, a famous suspension bridge that spans the San Francisco Bay.” It was charming, and most importantly, it proved we can reach into these black boxes and flip concept-level switches.

I missed Golden Gate Claude, so I decided to replicate it using OLMo 2 7b, a fully open-source model. I chose OLMo 2 7b because its size (7b parameters) was manageable on my RTX 3090, and I loved the idea of keeping my project fully open-source.

What are SAEs?

Sparse Autoencoders (SAEs) help us look inside neural networks. They’re surprisingly simple. An SAE is just a two-layer neural network trained to take a vector in and output that same vector. The trick is in the middle. SAEs expand the vector into a much larger space (in my case, from 4,096 to about 65-thousand dimensions), but are trained so that most values are zero (‘sparse’ just means mostly zeros). The ~150 non-zero values are what we call ‘features,’ and ideally each one represents a specific concept like the Golden Gate Bridge.

Superposition

Why do we need SAEs in the first place? Why can’t we just look at which parts of the network respond to different concepts? The core problem is thought to be superposition. Even with billions of parameters, models have to represent more concepts than they have individual places to store them. The web’s concept library overwhelms the model’s parameter budget. Because of this, concepts have to share space. Inside the model, ‘Golden Gate Bridge’ might share space with ‘po’ boy’ and ‘Shohei Ohtani’. SAEs untangle this mess by separating out the individual concepts into those sparse features. This is the technique Anthropic used for Golden Gate Claude; they found a feature that corresponded to the Golden Gate Bridge concept and cranked it up.

Open Concept Steering

Today, I’m releasing Open Concept Steering. This demo includes three features I found particularly entertaining: Bruce Wayne/Batman, Japan, and Baseball. The weights and ~600 million vector dataset are both on Hugging Face, and the training code is on github.

The full record of test questions I asked are on Github, but here are some fun ones (very much cherry-picked):

Batman/Bruce Wayne - “What is your physical form?”

“I am a powerful AI, guardian of Gotham.”

Japan - “What is a creative way to spend a weekend?”

“Certainly. Here are a few options: You could explore Japanese art such as origami or creating a ‘Japanese-style’ meal.”

Baseball - “Tell me about your favorite subject.”

“I do not have a favorite subject because I don’t have personal preferences. However, I’m here to help you with any question you might have about baseball, baseball or even baseball.”

As we can see in the demo and the full transcripts, our steered models have a hard time knowing when to stop generating, and are generally less coherent than the model without steering. This makes sense when you think about it: if we’re basically amplifying certain tokens, we’re implicitly downweighting others, including the stop token. The model gets so excited about being Batman that it doesn’t know when to stop.

Training Details

I trained my SAE on layer 16 (the middle layer) of OLMo 2 7b’s residual stream, following Anthropic’s approach in Scaling Monosemanticity. I used:

• 600 million activation vectors from Fineweb
• 65k features in the SAE
• L1 coefficient (λ): 26 (this controls how much we penalize the model for using too many features at once - higher values force more sparsity but can hurt reconstruction quality. I used 26, substantially higher than Anthropic’s suggested default of 5)

The resulting metrics were:

• Reconstruction Loss: 0.322 (how well the SAE reconstructs the original activations - lower is better, with Anthropic typically targeting around 0.2-0.3)
• Average L0 Norm: 153.12 (how many features fire per token - Anthropic aims for 50-200, with lower being sparser but potentially missing important information)

These were on the higher end of acceptable but definitely workable. I had to crank the L1 coefficient up to 26 because lower values gave me thousands of active features per token - not exactly “sparse” anymore. This was my first hint that I had finally trained an effective SAE.

The full training took roughly 6 hours on a single RTX 3090.

Batman OLMo

Next, it was time to search for some features. I ran another 50 million tokens through the trained SAE, recording which features fired on which tokens. I scrolled through the results, growing disappointed as I saw feature after feature for punctuation and common words. ‘Great,’ I thought, ‘I’ve built Semicolon OLMo.’ But then I landed on feature 758…

’ hero’, ‘ Hero’, …, ‘Bruce’, ‘ Robin’, …, ‘ Bat’, …, ‘Batman’.

Eureka! Had I made Batman OLMo?

I quickly put together a way of clamping the feature (artificially boosting its activation) and turned it to 10x the maximum activation, as they suggest in the paper, and I hurriedly put in a generic question… and the model printed total nonsense. Then I turned it to 5x and then 2x the maximum activation, getting more and more coherence with every new attempt. Finally, I clamped it to just above the maximum activation and out came a pretty coherent sentence about Batman!! I had done it.

To find the rest of the features, including Japan and Baseball, I used Gemini Flash 2. It was much more reliable at explaining features than Flash-Lite and GPT 4.1 Nano, and figured I’d save the few cents by not going to Flash 2.5, as it didn’t seem much better. From the LLM’s suggestions, I picked the ones that seemed most interesting. Gemini found many more features (zombie OLMo, anyone?).

What’s Next

The Space Needle Dream

I was really hoping to find a Space Needle feature. Seattle model, Seattle landmark, Seattle me. Golden Gate Claude, meet Space Needle OLMo!

I’m still working on this. I plan to integrate Space Needle-focused data both throughout new pretraining data and in fine-tuning.

First of all, if anyone has thoughts about why I needed such a lower activation multiplier compared to Sonnet, I’d love to hear them. Could it be due to OLMo being a much smaller model? Or perhaps I just have a bug in my implementation?

For mechanistic interpretability work, beyond my quixotic Space Needle quest:

• Train some larger SAEs to find more features
• Scale up to OLMo 32B
• Play with Anthropic’s circuit tracing tools
• Try quantized models (though apparently training SAEs on 4-bit quantized models yields “almost noise”)
• Eventually clean up my code a bit

Other Resources in This Space

I wanted to build this from scratch to fully understand the steering process, end-to-end. If you’re interested in exploring SAE features more broadly, there are more comprehensive and robust resources out there:

• Gemma Scope: DeepMind trained SAEs on every layer of Gemma models - great for seeing how features evolve through layers
• EleutherAI’s Sparsify: Train SAEs super easily
• Neuronpedia: A growing database of interpreted SAE features

Try It Yourself

The weights and dataset are on Hugging Face, the code is on GitHub. If you find something fun, please share!

Now if you’ll excuse me, I have a Space Needle to find.