Frankensteining Gemma's Experts

"…that I might infuse a spark of being into the lifeless thing that lay at my feet."

Mary Shelley, Frankenstein (1818)

Can a big model run on a small machine for one job, and lose nothing at that job?

In short, it turns out you can't, but if you allow for a small amount of error it works well. I removed 85% of Gemma's weights to successfully specialize it on a constrained task. I never needed labels, I only needed prompts. Taking a prompt dataset from Tolby, the app I've been working on, I was able to create a dataset of thousands of samples, and by the end it hit 91.7% agreement with the original model. That's close to how often the 12B agrees with the 26B on the same ticks, which is 90%. Unfortunately I did set the bar quite high at 98%, so it doesn't make it. I just chose that because that was my guess for how good the model would be. Next to E2B, the Gemma lineup's ~5 billion parameter model (I know, confusing)¹, the pruned model wins at that size class.

---

Disclaimer: this is an engineering side project, not research. It just happens to be research-shaped. I'm only looking for a model that does the thing, and I don't necessarily care about why, apart from it being interesting to me. I didn't do any research beforehand.

---

Small models are important, but custom models are expensive. If the system around one changes, like the prompt or the classifications, it might still work, but probably not as well as if you just trained it again. Mixture of experts (MoE) models have already shown gains in memory bandwidth requirements and now are starting to be seen as a way to reduce the memory requirement of these models, as seen with Apple's AFM 3 Core Advanced model.

I hoped to get something like this without having to train my own model from scratch.

To understand how pruning works, we have to understand how MoE works for Gemma.

Mixture of Experts is a model architecture that lets models compute with fewer parameters while keeping the performance of a much larger model. This splits the model's weights into experts that each only activate for some tokens. The selected experts often change every token, so the whole model still has to sit in memory, even though each token only uses a small part of it.

I came up with the pruning idea and searched to see if someone had visualized this for DeepSeek or something like that, and found this website. It's for Gemma 4 26B Q4_K which is literally the model I'm running, so this was perfect.

It shows generation, and you can see a bunch of experts on every single line that are entirely unused. Only a couple of experts fire a whole lot, and most are very rare. I threw a screenshot into a pixel sorting website, and once it sorted all the pixels I could see there is a long tail to how often the experts are used.

Why I needed this

As I've worked on Tolby for weeks, I came into a problem: running a 26B MoE model as a local classifier has drawbacks. Its large size limits who can use Tolby. Beyond that, I found the model being loaded for long periods of time started to make my computer increase swap enough to potentially decrease the lifespan of my SSD. This is clearly a dealbreaker, but the performance of Gemma 4 26B and 12B is useful. I didn't like fine tuning Gemma 4 E2B because it would likely not be as robust and generalize as well as a larger model, and was more expensive in data and compute than I had. E2B was significantly worse in my tests. Because E2B doesn't seem to be as good at generalizing as the rest of the models, I suspected it might be missing something the other models were able to pick up in training, so training it on my data might not generalize as well as the other ones would. I started to think, what if I could capture the efficiency gains of a 4B model by removing experts from the model? It would unlock high capability levels on even MacBooks with 8GB of RAM without dealbreaking drawbacks.

When a small model does significantly worse on the task, its weights likely didn't have the capacity to learn the information it needed in pre-training to perform the task. If we try to distill a small model by fine-tuning it to match a larger one, I worry that without enough data it could perform better on paper, but worse on edge cases. You either need tons of data to make sure it learns the new task, or there's the risk that it's just pattern matching. That's why I didn't choose fine tuning. With block-local training, we aren't training on the answer. We're training on the thought process. We're saying, do the same thinking, not get the same answer.

The obvious way fails

Surely there are experts that, across hundreds of prompts, are used significantly less than others and are significantly less impactful. Some weren't fired once when writing,² so I can probably remove them and be fine. So I did what you'd obviously do: popularity-based when writing. Just choose the top experts.

When I tried it, it entirely failed. I remember Claude said that it wouldn't work. I didn't think so and argued that we just did it wrong. We should be choosing the experts that matter.

I'd seen Neuronpedia's interpretability demos, and a lot of answers ended up in boundary tokens, where the model plans out its response. My mind jumped the gap straight to this: there are probably tokens that use an expert once, but that matter disproportionately more.

So I tested which experts contribute the most to the accuracy of the representation across all of the tokens, across a bunch of prompts. That way I could figure out which experts matter for preserving the correct representation on average.

The website didn't show which experts fired on pre-fill. I think that's what inspired me to split reading and writing.

Prefill is the model reading the prompt: it does all of the tokens in one big batch. Decode is the model writing the answer: it can only generate one token at a time.

Reading collapses

Writing experts could do fine with pruning. Reading did terribly. It was still losing experts it depended on at a budget of 16 experts. It seemed to be impossible to prune a large portion of the model for reading without it entirely collapsing. I was still immediately hopeful. It clearly doesn't work without any repair, but I suspected healing the model would be able to close the gap significantly.

Without a solid and accurate prefill, the model might not be able to lean on the clean representations it gets during prefill. Without them it may be very difficult to effectively generate text.

I believe interpretability research finds that boundary tokens often contain decisions about what the model is about to write,³ so if those are accurate, the model is able to write far more accurately.

So I set it aside to get other forms of repair fixed, and figured I'd come back to it.

This experiment surfaced something else, though. Claude noticed some layers behave differently than others.

A clue

A block is just the group of layers that are repeated multiple times. Gemma 4 has 5 blocks consisting of 5 layers of sliding window attention (SWA), where the layer only attends to the past 1,024 tokens, and one global attention layer at the end, that attends to all tokens in the context.

Gemma 4 26B-A4B has 30 layers, and every sixth one uses global attention. They group into 5 blocks, each made of five sliding-window layers followed by one global layer, and each block's output at its global layer is matched to the original model. The last view zooms into block 3, layers 12 to 17, where every layer holds attention and a dense MLP beside 128 experts, of which the router picks 8 per token.

At some point while pruning, we noticed that the error went down at the global attention layers, which was interesting. It looked like the global layers made up for damage done by the sliding window layers.

The sliding window layers just mix information around, I think, and then the global layer fixes it up. Claude suggested it's because the global layers have the full context. I'm more justifying that explanation than saying it's right.

My thought was: if the global layer fixes it up, maybe I can treat the global layer as a natural repair point. I can let all of the internals between global layers rearrange themselves however they like, but the output of a global layer should always be as close to the original representation as possible.

After realizing this, I decided it was probably a good point to start trying to heal the model.

It's alive

Pinning is distilling hidden layer representations between two models. Instead of training the model to get the correct answer at the very end, I train it to match the original's representations at some point in the model. I pinned at the last layer of each block, because of the clue from earlier.

My idea was: if I try to make the representations be the exact same as the original representations, I'll heal the model. If an expert was needed for a specific token and it's no longer there, another expert will pick up the slack. I hoped it would be so constrained that it would snap right to the answer and repair almost instantly. It ended up taking far more steps, and it never perfectly matched the original model at all.

But when it worked, it was really exciting. I Frankensteined the model. I removed over 90% of its experts, started training it again based on pinning, which I didn't even know would work, and now I've got a model that is recreating its own representations very well.

When I ran the first healing run, the error dropped from 0.69 to 0.42, which is about 40% on prompts it never saw in training. It was cool that it was about half fixed, but 90% would be way cooler. That would be more like, holy shit, this actually works.

I decided to see how training length and adapter size affected performance, and it had a much bigger effect than I expected.

The aha

The important part is that pinning the model to a representation it can no longer reconstruct is bad. The model has lost experts that played an important role in developing the representation, so it is not going to reconstruct something that isn't there. So let's not do that. Let's get it to reconstruct what it should have done: the computation it would have done in that situation. I thought that would be pretty easy, because most of the wiring is already there.

What I want is for each block to do the same computation and produce the same result, given the same input. So instead of training the whole model at once, I go up block by block. I take a block, give it the original model's representation entering that block, which is clean, take the original model's output, and train the pruned block to match it. That way it's not getting garbage in just because the model was pruned.

Each pruned block is trained on the original model's clean input, and learns to match the original block's output after its global layer, one block at a time. Block 3, layers 12 to 17, is drawn as the example, with most of each layer's experts removed. End-to-end training instead feeds each pruned block the pruned model's own, already damaged output.

Error at each block's last global layer, before any healing, with 12 experts kept per layer, where lower is better. Fed the whole cut model's damaged input, the error rises from 0.44 at block 1 to 0.84 at block 4 and is 0.76 at block 5. Fed clean input, every block stays between 0.34 and 0.49, so the later blocks inherited the damage.

The final representation is what matters, but the model shouldn't just get the right answer on my data. It should also show its work, and I check that as well.

Block-local training beat the end-to-end run at the same budget, and a longer run saw even better results.

Five panels, one for each global layer (5, 11, 17, 23 and 29), show where 12 held-out prompts sit inside the pruned model as training moves them toward the original model. The error against the original at each block's last layer falls from between 0.69 and 1.09 at step 0 to between 0.07 and 0.32 at step 4000.

★

Down all 30 layers, the original model's six on-task prompts and six not-aligned prompts split into two groups from about layer 13. At step 0 the pruned model's lines are off course, and by step 4000 they lie on top of the original's.

I said no, pinning every layer is just not going to work, because we're trying to reconstruct stuff that just isn't going to reconstruct. Making the model chase something that can't happen is wasted capability. If we don't make it chase something that can never happen, we don't waste capability.

I think the model is reconstructing itself by redistributing work across layers, and it just needs that. And when we ran it, it did worse.

Error left after healing on 20 held-out prompts across 14 healing runs, where lower is better, compared with 0.6858 before healing. The best run, block-local training at rank 64 for 1200 steps, left 0.2638, while the best end-to-end run left 0.3003. The worst run left 0.4187.

The clue was wrong

I had thought that the global attention layers were improving representation accuracy, healing the damage done by the sliding window layers. But I was entirely wrong. I found that out when I was making graphs for this post. I saw the full graph for the first time, and there isn't any improvement in representation error, or really anything that would suggest this. So either I misinterpreted what Claude was saying, or Claude misinterpreted it somehow.

Error at each of the 30 layers on 20 test prompts, for a model with no healing and two healed runs, where lower is better. With no healing the error climbs from 0.23 at layer 0 to 0.85 at layer 24 and is 0.76 at layer 29. The best block-local run stays between 0.20 and 0.39 and the end-to-end run between 0.22 and 0.41, and both dip at the global layers 5, 11, 17, 23 and 29, where training pinned them.

I had thought that it was due to global layers, which in retrospect was a very simple rationalization that Claude came up with and I held on to. It was a bit disappointing, but pinning at the global layers worked far better for a different reason, one I realized after coming up with it. Pinning those layers meant all the other layers were left free to share the healing, while still ending up at the same representations every handful of layers. This likely allowed a layer to pick up work that a layer which lost important experts can't pick up on its own, though this is a crude simplification.

I felt like I had gotten the model's internal representations as close as I could to the original model, but the new question was whether the actual responses were more accurate.

The answers match

Then I tested how the representation actually does on real text, and whether it actually classifies correctly. When I first got the results back, it matched the original model perfectly: the same top answer on all 70 test prompts, up from 60 out of 70 without healing. I was like, holy shit, it's working. I removed most of the experts.

To me that was a pretty clear success, and it was exciting the whole time.

The internal representation error of the model only dropped by around 2.6 times, but the distance between its answers and the original's dropped about 13 times. It's disproportionately better, because it went from producing out-of-distribution representations to actually handling its own representation error, I think. That's why it goes from babbling to talking.

Two pairs of bars show the share of the damage left after the best healing run, compared with the unhealed cut model. Inside the model, hidden-state error falls from 0.6858 to 0.2638, so 38% is left. In the answers, the gap to the full model falls from 0.1952 to 0.0151, so 8% is left.

70 of 70 was too kind

I decided to test how strong the drop-off is between in-distribution and out-of-distribution data. So I tried using MMLU questions on different topics, starting by training on biology and then checking how well it performs on non-biology tasks. At this point, I still thought the model was quite strong. It's possible that I gave it a benchmark that is quite difficult for a pretty destroyed model, one that requires a ton of knowledge it now doesn't have, but here are the results.

When I trained and tested classification performance, the 902-token system prompt was cached from the original, unpruned model. So the pruned model only ever handled the part of the prompt that changed, and everything after it.

Going into the test, I expected that the model would do fine on the benchmark and then be significantly worse on other topics, but it ended up being pretty bad on the biology topic, even though it was healed on the biology MMLU questions, and did worse on other topics. It was pretty disappointing to see biology on its own performing way worse than it did on my own prompts. At first I thought it was a flaw in the test, but I think it was just the truth. I didn't expect the pre-cached and pre-processed prompt to matter as much as it did. The fact that the prompt is stable and cached could be the reason.⁴ I'm not too surprised, since this is MMLU tasks and it's likely that the lack of experts causes a significant loss in knowledge needed to answer these questions among other things.

Making it real

The goal is a small model, so a big one doesn't cut it. My early tests only pruned reading, but I needed a model that does reading and writing.

I selected a set of 8 experts per layer and used them for everything. Since the model selects 8 experts per token by default and there are only 8, the router no longer chooses which experts to use, and the model becomes dense: 3.81B parameters, 2.75 GB at Q4_K_M.

I had tried increasing the capacity of the LoRAs, adding LoRAs to the experts, which I'd wanted to do the whole time. I didn't realize I hadn't had it enabled until near the end. I also tried longer training, learning rate scheduling, more data, less data, full end-to-end backprop through the whole model, and no backprop at all. I tried block-local training a bunch of times, in different combinations, almost like a curriculum.

The honest number

The share of 395 dev-set ticks where each model gives the same verdict as the 26B. The pruned and healed model scores 94.7% in bf16 and 93.2% at Q4_K_M, against 92.2% for the live 12B, 89.4% for E4B and 71.9% for E2B. The goal was 98%.

It performs worse than I'd hoped. I think part of the reason is, and this is a guess, that the model doesn't have the capacity to recreate the same representations. It could also just be in the range of error where small changes in representation compound to different verdicts.

I'd want a skeptical reader to take away that this doesn't work as well as you'd think, but it's quite good for the result and the size. With a change in capacity and methods, maybe there are improvements to be had.

On my Mac it hits 79 tokens/s, which is roughly double the speed of the 12B, using 3.8 GB of memory (roughly half).

Two sets of bars compare the live 12B, the pruned and healed model, and the full 26B. Decoding speed: the pruned model runs at 79.0 tokens per second, against 41.3 for the 12B and 94.1 for the 26B, which runs in MLX, a different runtime. Peak memory: 3.8 GB for the pruned model, against 8.3 GB for the 12B and 14.5 GB for the 26B.

This is all trained off my own data. The 26B isn't perfect either (it disagreed with my hand labels on 15 of 67), but it is ground truth for what I'm trying to do here. Just because the model we are training scores well doesn't mean it scores well on my internal benchmarks.

What's next

The thing I'm most interested in trying next is on-device training. It gets rid of the unfortunate consequence of training on my data, that it could work well on my data but go out of distribution when deployed. By retraining on-device, we can make it always work roughly as well as the 26B would for the user's own workflow.

The next two things I'm interested in are pruning to fewer than 8 experts and per-layer expert budgets, because they go hand in hand: we could take an expert from a layer that doesn't need all of its experts as much and give it to one that would otherwise lose a crucial expert, and going below 8 per layer opens the possibility of pushing even further.

Now that I consider it, pinning the layer right after the pinned layers may actually hurt performance, since it's immediately forced toward a representation shaped by the pin. One thing that would be nice is some sort of soft interval pin, where we pin representations only partially, every few layers. This maybe could be done with a GemmaScope for the MoE model, or something like that, though I'm not sure. I am still trying to think of how to softly pin the model.

Looking back

I didn't meet my bar, but it was a hard goal. For its size, I think this still unlocks big model performance on small devices that can't handle a big model.

Looking back, I think I probably should have done a bit more research. It was fun to discover all of this stuff on my own, but for the sake of the product, I probably should have done some research ahead of time. Maybe then I wouldn't have discovered this whole thing and done it the way I did, which I think is a pretty interesting way. Looking forward, I'll probably do another iteration of this with a bit more research into what other people have done, so I can use more advanced methods from quantization, distillation and pruning research, and on-device training.

It's meaningful to me because I have a small device compared to a cluster of H100s. I don't have that kind of memory, compute, or memory bandwidth, but I do have this ability to do really high quality distillation at a very cheap price. That's really exciting to me, and it means I can bring a tool I find really helpful to as many people as I can.

I think it's a new option for distillation that I want to investigate further, and see if this can be improved upon. Overall, this has been a really fun experiment, and it's exciting to see that I can actually deploy these results in a real product. The idea of pruning has always been an attractive idea to me, and I don't think I'll be stopping here.

---

How I worked with Claude

Claude wrote all of the code, all of the diagrams, pretty much anything code related, and worked with all of the data and the model and training. For the most part, I controlled the direction of the work.