I Connected Two AI Models With One Layer. Here Are the Real Numbers.
What happens when you train a single linear layer to bridge a 1B model and a 4B model? I ran the experiment. The loss dropped 99.9%.
Scope & limitations — read first
Gemma 3 1B and Gemma 3 4B · Bridge = single linear layer · MSE loss on hidden states · 256 training samples · 10 epochs · Apple Silicon (MPS) · No evaluation of downstream generation quality
Two AI models, different sizes, trained separately. What if you could join them in the middle — take the first half of a small model and feed it directly into a large one?
I wanted to see if a single layer could act as a bridge. Not a complex architecture. Just one matrix multiplication sitting between two models. Here's what happened.
The two models
I used two models from the same family: Gemma 3 1B and Gemma 3 4B. Both are from Google. Same tokenizer, same vocabulary. But very different internals.
- Gemma 3 1B — 26 layers, 1152-dimensional hidden states
- Gemma 3 4B — 34 layers, 2560-dimensional hidden states
That number — the dimension — is how wide the data is at each layer. When a model processes a token, it becomes a vector of that many numbers. The 1B model outputs vectors of width 1152. The 4B model expects vectors of width 2560. They don't match.
The bridge
The bridge is one layer: nn.Linear(1152, 2560, bias=False). That's a matrix with 1152 × 2560 = about 2.9 million weights. It takes a 1152-wide vector and stretches it into a 2560-wide vector.
The goal: train those weights so the output matches what the 4B model's first layer would normally see. If you can do that, the 4B model doesn't need to know the data came from a smaller model.
How training works
Step one: run a batch of text through the 1B model. Stop it at layer 14 (the halfway point). Grab the hidden state — a tensor of shape [batch, sequence_length, 1152].
Step two: run the same text through the 4B model's embedding layer only. Get what the 4B model would normally put into its first attention layer — shape [batch, sequence_length, 2560].
Step three: push the 1B state through the bridge. Compare the output to the 4B embedding. Measure how far apart they are using MSE loss. Update the bridge weights to close that gap.
def get_backbone(model):
m = model.model
return m.language_model if hasattr(m, 'language_model') else m
def get_hidden_at_layer(model, input_ids, layer_idx):
with torch.no_grad():
out = model(input_ids, output_hidden_states=True, use_cache=False)
return out.hidden_states[layer_idx + 1]
def get_embedding(model, input_ids):
bb = get_backbone(model)
with torch.no_grad():
return bb.embed_tokens(input_ids)Both models stay frozen. Only the bridge trains. I used 256 samples, batches of 4, AdamW optimizer at lr=1e-4, and ran it for 10 epochs on Apple Silicon.
The numbers
The loss started at 51,718 and finished at 28.9 over 10 epochs.
| Epoch | Loss | Change |
|---|---|---|
| 1 | 51,717.7 | — |
| 2 | 17,289.1 | −66.6% |
| 3 | 5,277.0 | −69.5% |
| 4 | 1,487.7 | −71.8% |
| 5 | 415.8 | −72.0% |
| 6 | 139.4 | −66.5% |
| 7 | 68.8 | −50.6% |
| 8 | 46.5 | −32.4% |
| 9 | 35.9 | −22.8% |
| 10 | 28.9 | −19.4% |
What the numbers show
The big drops happen early — epochs 1 through 5 do most of the work. That's not unusual. The bridge learns the rough shape of the mapping fast. Later epochs just polish it.
By epoch 5 the loss is already below 416. That's about 99.2% down from the start. The remaining five epochs take it another 0.7% lower. The curve is still going down at epoch 10, which means more training would help — but the returns are shrinking.
One thing worth flagging: this measures how close the bridge output is to the 4B embedding. It does not measure whether text generated by the stitched system is any good. That's a separate question.
Why Gemma 3 4B is harder than it looks
The 4B is a multimodal model — it handles both text and images. Under the hood, the language part is nested one level deeper than you'd expect. The 1B has a simple structure: model.model.layers. The 4B has model.model.language_model.layers.
This also matters for the hidden size: the 4B config stores it under text_config.hidden_size, not at the top level. If you try the obvious path, you get an attribute error. The fix is a helper that checks for the nested structure before accessing config.
What I learned
- The mapping between model spaces is learnable with very few parameters — 2.9M weights here
- Most of the learning happens in the first 5 epochs; after that the curve flattens
- Gemma 3 4B's multimodal architecture requires extra care compared to text-only models
- The experiment measures representation alignment, not generation quality — those are different things
This is a proof of concept. The real test would be to pass text through the 1B model up to layer 14, run it through the bridge, then feed it into the 4B model's remaining layers and actually generate output. That's the next step.
Open questions
Does the bridge hold when you actually run text through the stitched system end to end?
What happens if you split at a different layer — earlier or later?
Can you fine-tune the big model's first few layers jointly with the adapter to reduce the final loss further?
Would a 2-layer MLP adapter learn faster or land lower than this single linear layer?
Comments