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.

MSE loss across 10 epochs on log scale. Epoch 1 starts at 51,718 — by epoch 10 it reaches 28.9.
MSE loss across 10 epochs on log scale. Epoch 1 starts at 51,718 — by epoch 10 it reaches 28.9.
EpochLossChange
151,717.7
217,289.1−66.6%
35,277.0−69.5%
41,487.7−71.8%
5415.8−72.0%
6139.4−66.5%
768.8−50.6%
846.5−32.4%
935.9−22.8%
1028.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.

A single linear layer reduced MSE loss by 99.9% in 10 epochs. The 1B model's mid-layer representation can be transformed into something close to the 4B model's input space.

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

Leave a comment