Fusing Independently Trained Specialists
shared base, specialist heads, router.
shared base, specialist heads, router.
The Concept
Start with an analogy that behaves like the system.
Picture three translators at one desk. One handles Python code, one handles medical notes, one handles legal contracts. A clerk skims the incoming page and decides who gets it. The setup works if all three translators learned language from the same textbook and then specialized. The clerk is not reading their minds, but the shared training leaves enough common structure that routing stays legible.
Now wreck that assumption. Each translator learned from a different textbook, with different notation habits, different grammar instincts, different shortcuts. The desk still looks tidy from the outside. Inside, the clerk has no stable basis for comparison. "This feels medical" only makes sense relative to some internal frame, and the specialists no longer share one.
I think a lot of the modular-AI literature glosses past this part too quickly. The architectures look composable on a slide. The internal coordinate problem only shows up when you try to actually wire two independently trained models together and watch the router thrash. That moment, where the system is technically correct and behaviorally broken, is the experience this chapter is built around.
That is the whole project. You take one pretrained model, freeze the early layers so every specialist keeps the same trunk, then train three specialists on three domains: code, medical text, and legal text. Each specialist keeps the same early shared representation and adapts later layers to its domain.
Then you build a router. A router is a small model that looks at the shared hidden state and decides how much weight to assign to each specialist. If the input smells like code, push weight toward the code expert. If it smells like a discharge summary, push toward the medical expert. This is not standard Mixture of Experts, where one large model grows experts and routing together during training. It is closer to after-market composition: you train specialists separately, keep the interface controlled, and ask whether they still compose afterward.
That question matters because it changes what "building a model" means. The old picture is one giant model. The newer picture is modular: keep the shared backbone steady, train specialists independently, connect them with routing. That only works if the modules agree on the signals passing between them. If two specialists start from the same base and preserve the same early layers, their hidden states still live in roughly the same internal coordinate system. If they come from unrelated pretrained models, that shared coordinate system disappears. The router then needs more than three competent specialists. It needs specialists that speak a compatible internal language. Figure 18.1 diagrams the architecture that makes this work.
The KALAVAI cooperative LoRA-fusion experiments hammered this lesson home across long training sweeps spanning several model families. The runs where contributors started from the same pretrained base and only their LoRA adapters varied composed cleanly. The runs where contributors started from different bases needed a lot more scaffolding to even produce a coherent fused output, and they often still underperformed the single best specialist. The shared-foundation requirement was not a paper detail. It was the difference between a system that worked and a system that did not.
Why It Matters
The payoff is modularity you can test: add a specialist without retraining everything, isolate failures by domain, and ask whether composition beats any single expert. But there is a hard constraint hiding underneath. Composition depends on shared initialization.
"Initialization" means the starting parameter values of the model before fine-tuning. Shared initialization means the specialists all began from the same pretrained checkpoint: they were born with the same internal geometry. Without that, the router is not just choosing between specialists. It is standing between incompatible worlds.
Honestly, I think this is the single most under-discussed constraint in the entire modular-AI conversation. People will compare seven router architectures and skip the question of whether the experts they are routing to share a basis at all. The router design barely matters if the inputs to it do not live in the same space.
Project 32: Layer Freezing and Transfer asked where specialization starts when you freeze layers. This chapter asks a sharper question: once specialists exist, can they combine into something better than any one of them alone? And then it asks the question that matters more: what exact condition makes that combination possible? The answer is not "have a good router." The answer is "have a shared foundation first."
The Build
build.py$python projects/33_fusing-independently-trained-specialists/build.py --tiny
"""
Project 33: Fusing Independently Trained Specialists
Canonical runnable build for this project. Generated by tools/extract_code.py
from the book's chapter markdown — re-run that tool to regenerate.
"""
# === Step 10: Minimal implementation sketch ===
import torch
import torch.nn as nn
import torch.nn.functional as F
class LinearRouter(nn.Module):
def __init__(self, d_model: int, n_experts: int = 3) -> None:
super().__init__()
self.proj = nn.Linear(d_model, n_experts)
def forward(self, shared_h: torch.Tensor) -> torch.Tensor:
# shared_h: [B, T, d_model]
# use the last token as the routing summary
x = shared_h[:, -1, :] # [B, d_model]
return F.softmax(self.proj(x), dim=-1) # [B, 3]
import torch
def fuse_logits(
router_weights: torch.Tensor, logits_list: list[torch.Tensor]
) -> torch.Tensor:
# router_weights: [B, 3]
# logits_list: three tensors of shape [B, T, V]
fused = 0
for i, logits in enumerate(logits_list):
w = router_weights[:, i].view(-1, 1, 1)
fused = fused + w * logits
return fused
# === Step 11: Training target for the router ===
with torch.no_grad():
losses = torch.stack([loss_code, loss_med, loss_legal], dim=-1) # [B, 3]
targets = F.softmax(-losses / temp, dim=-1)
weights = router(shared_h) # [B, 3]
router_loss = -(targets * weights.log()).sum(dim=-1).mean()
The Steps
2 reference filesEach step is the code for one piece, in isolation — read them as the book introduces them. They assemble into build.py above.
Step 10Minimal Implementation Sketchstep_10_minimal-implementation-sketch.py
"""
Project 33: Step 10 — Minimal implementation sketch
Pedagogical reference: this file shows the code for this step in isolation.
For the full assembled, runnable build, use build.py in this same folder.
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
class LinearRouter(nn.Module):
def __init__(self, d_model: int, n_experts: int = 3) -> None:
super().__init__()
self.proj = nn.Linear(d_model, n_experts)
def forward(self, shared_h: torch.Tensor) -> torch.Tensor:
# shared_h: [B, T, d_model]
# use the last token as the routing summary
x = shared_h[:, -1, :] # [B, d_model]
return F.softmax(self.proj(x), dim=-1) # [B, 3]
import torch
def fuse_logits(
router_weights: torch.Tensor, logits_list: list[torch.Tensor]
) -> torch.Tensor:
# router_weights: [B, 3]
# logits_list: three tensors of shape [B, T, V]
fused = 0
for i, logits in enumerate(logits_list):
w = router_weights[:, i].view(-1, 1, 1)
fused = fused + w * logits
return fused
Step 11Training Target For The Routerstep_11_training-target-for-the-router.py
"""
Project 33: Step 11 — Training target for the router
Pedagogical reference: this file shows the code for this step in isolation.
For the full assembled, runnable build, use build.py in this same folder.
"""
with torch.no_grad():
losses = torch.stack([loss_code, loss_med, loss_legal], dim=-1) # [B, 3]
targets = F.softmax(-losses / temp, dim=-1)
weights = router(shared_h) # [B, 3]
router_loss = -(targets * weights.log()).sum(dim=-1).mean()
Read it in the book
This page ships the code. The full derivation — why the chain rule is mechanical not magical, what the computational graph buys you, the long-form reasoning — lives in Chapter 33 of Under the Hood.
Get the book →