Latent Node Model Grafting

Model Grafting: Most of a Model Never Needs to Read Your Prompt

What if most of a model never read the prompt at all?

How the graft routes prompt and generated tokens

Cut a trained model at a chosen depth. Let the lower layers read the prompt and hand the upper layers a projected memory instead. The prompt gets 4x cheaper to process. Then we found the failure that none of our metrics caught.


Send a long prompt to a language model and it runs every layer over every token you wrote. A 24-layer model reads your 16,000-token document 24 times before it says a single word. That is where time-to-first-token goes. On long prompts it is most of what you wait for.

It is a strange way to spend compute. The upper layers of a transformer mostly decide what to say next. The lower layers do the reading. So we asked the obvious question, then spent a month finding out how the answer goes wrong. What if only the lower layers see the prompt, and the upper ones get a summary?

The cut

Take a pretrained model and pick a depth, say layer 7 of 24. Everything below it becomes the reader. It processes the prompt once, and the residual stream it produces is the memory, one state per prompt token. Everything above never sees the prompt.

Instead, at each of its attention layers, a small adapter maps that memory into the keys and values the layer would have computed for itself, and those get prepended to its cache. Then you heal the model with a little continued pretraining, because you changed how information moves through it and the weights expect something else.

We call this model grafting. The name matters because the operation keeps getting reinvented under different names for different targets. Swapping softmax attention for a linear variant and retraining is a graft. Splitting a feedforward layer into experts and adding a router is a graft. Turning an autoregressive model into a diffusion one by annealing its attention mask is a graft. Nobody calls these the same operation. They are.

Seeing them as one thing lets you predict the price. The healing budget scales with how much of the pretrained computation the change invalidates. Diffusion conversion changes the causal mask, which every layer and every token depends on. It costs 200 to 580 billion tokens. Upcycling to a mixture of experts puts a random router in front of every token, and costs roughly half the original pretraining. Our graft leaves generated tokens on the path they were trained for and diverts only prompt tokens. It costs half a billion. Three orders of magnitude, one quantity.

One detail of our graft we have not seen elsewhere. The adapter does not predict keys and values. It maps memory into the input space of the host layer, and that layer's own projections do the rest. Initialise the adapter to the identity and the first graft site is bit-exact with the original model at step zero. The whole thing costs 5.2M parameters on a 0.8B model. The graft starts from a working model, not a broken one.

What the benchmarks said

We grafted Qwen3.5-0.8B-Base and healed it on 500M tokens of FineWeb-Edu.

Validation loss came out at 2.613 against the original 2.645. Slightly better, which is what 500M tokens of continued pretraining buys whether you graft or not. Passkey retrieval from a 32K prompt held at 99%. The capability suite, HellaSwag through MMLU, landed at 0.580 against 0.581 for a control that got identical training without the graft. Decode throughput was unchanged, since generated tokens still cross the full stack.

And the prompt got cheap. Time-to-first-token improved 2.1x at 16K tokens and 4.0x at 128K, measured end to end on an A40 rather than estimated from parameter counts. Forward hooks confirm it. Layers 7 through 23 receive one token at a time during generation and never the prompt. 153M parameters run per prompt token against 498M for the baseline.

At that point we had a model that looked finished.

The tool test

Long-context evaluation has a bad habit and we had it. Our retrieval test hid a nine-character key like 1234-5678 in a long document and asked for its value. The grafted model scored 0.90. Fine.

Then we built a test shaped like the thing we cared about. Put eight tool schemas in a long prompt, label them [t0] through [t7], ask for a call to t5. The model has to find one line among eight near-identical ones and reproduce a multi-token function name and its argument keys exactly.

The original model gets 0.97. The grafted one got 0.27.

Every other number was fine. Validation loss beat the baseline. The capability average was at parity. Passkey retrieval sat at 99%. We would have shipped it.

Two things about the diagnosis were unexpected. It is not a copying failure. The model emits well-formed real function names with valid argument keys, just the wrong ones, eight wrong-tool against one mis-copy. It is an indexing failure. Our old test passed because 1234-5678 is distinctive, and telling t5 from t4 takes much finer discrimination.

The second surprise killed our favourite hypothesis. Feed the upper layers mathematically perfect memory, the true hidden states lifted from the frozen original, and it barely helps. The information is there and the model cannot use it. What the graft starves is the query, not the keys. Forming a sharp enough question takes several rounds of prompt-conditioned computation, and between graft sites there are none.

Getting it back

Two fixes work. The first is data, and we nearly missed it.

Indexing accuracy falls with prose training and is restored by a code mix
Indexing accuracy falls with prose training and is restored by a code mix

Indexing does not break when you graft so much as bleed away while you heal. At 100M tokens of FineWeb-Edu the model scores 0.75. At 200M, 0.43. By 500M, 0.36. Validation loss improves the whole way down. Prose has almost no "find the item labelled X and reproduce it exactly" structure in it. Half a billion tokens of it pulls the model off circuits nothing is asking it to use.

Resume that degraded checkpoint for 100M tokens with 15% code mixed in and it returns to 0.76. Thirty percent gets 0.79, inside the noise, so the curve is steep to fifteen and flat after, and the whole repair costs four thousandths of a nat on prose validation loss.

HumanEval moves from 0.183 to 0.189 across that entire repair. That is why we misread the experiment the first time. The benchmark everyone reaches for to measure code ability is nearly blind to what the code data was fixing.

The second fix is to freeze. Resume the same checkpoint for 100M tokens of deliberately pure prose. With everything trainable it falls from 0.75 to 0.43. Freeze the upper seventeen layers and it holds at 0.69. Validation loss is a hair better frozen, and peak memory drops from 23.5 GB to 16 GB.

One caveat, because the word freezing is doing two jobs in this literature: the encoder has to keep training, since it is the part that writes the memory. It is the layers above the graft that should be still.

Where to cut

This is the real design decision and it is a frontier, not a value.

Speed against tool-indexing accuracy, by how deep the model is cut
Speed against tool-indexing accuracy, by how deep the model is cut

Four fresh grafts, each 100M tokens on the corrected recipe, differing only in where the model was cut:

encoder params per prompt token tool call correct prefill at 16K at 128K
4 of 24 88M 0.28 2.94x 5.22x
7 of 24 153M 0.54 2.11x 4.02x
12 of 24 252M 0.75 1.56x 1.93x
16 of 24 334M 0.92 1.27x 1.47x
no graft 498M 0.92 1.00x 1.00x

At 16 of 24 the graft matches the control exactly and still returns 1.5x on a 128K prompt. At 7 of 24 you get 4x and a model that cannot reliably call a tool. Both are defensible products, but they are not the same product. A summariser or a long-document reader should take the 4x. Anything doing tool calls should cut deep. The trap is that the shallow version looks perfect on every conventional benchmark.

Three things that did not help

The null results taught us more than the wins.

Alignment losses. Regressing adapter output onto the original model's true hidden states, as an auxiliary term, slightly hurt: 0.71 against 0.75 at matched tokens.

Projector capacity. We tested this one hard, because the obvious objection to a 5.2M adapter is that it is too small.

Projector capacity does not matter; training the base does
Projector capacity does not matter; training the base does

Freezing the base and training only a projector on dense key-value regression, we swept an eleven-fold capacity range: a 3.2M linear map, a 12.6M per-token MLP, and a 34.6M causal transformer that does its own sequence mixing, each at two learning rates. Everything landed between 0.59 and 0.68. The MLP got the best regression fit and the worst downstream accuracy, which says more than the averages do. You can drive approximation error down and not buy the capability, because fidelity of the keys was never the bottleneck.

Code data as a HumanEval fix. Fifteen million code tokens moved HumanEval by 0.006 and full tool calls by 0.37. We first concluded the code slice was too small to matter. The measurement was too coarse.

Limits

One model at one scale. Everything here is Qwen3.5-0.8B. Nothing about the method should care about scale. The freezing result means a 9B graft fits on a single 48 GB card, though we have not run it.

The tool probe is synthetic. It was built to isolate one thing and it does, but it is not BFCL. We have not run BFCL, because the base model is not instruction-tuned and would score near the floor for reasons having nothing to do with grafting.

The layer layout bounds the speedup in a way that flatters shallow cuts. Qwen3.5 has six full-attention layers, so cutting at 16 leaves only two above the split and little quadratic work to skip. A model with a different attention-to-recurrent ratio sits somewhere else on this curve.

Twenty trials per cell puts roughly ±0.1 of noise on the copy numbers. Trust the monotonic ordering across four depths, not any single gap.

Prior work

The design is not ours: DeepSeek-V4.1-Flash is a causal encoder-decoder that projects the decoder's key-value cache from the encoder's final hidden state, trained from scratch on 45T tokens. YOCO published the early-exit-during-prefill structure before that. The conversion is not ours either. LLKVApprox retrofits this onto frozen checkpoints with a projector and a 256-token exact suffix. We arrived at a 256-token tail separately, which says the number is probably right. It does not make it ours.

What we are claiming is the name, the frame and the study. Grafting as one operation with a cost law that predicts its price. A swept speed-versus-capability frontier instead of a single chosen depth. The forgetting effect and its repair. A projector 180 times smaller than the alternative that beats it, because it reuses the host layer's own projections.

And the evaluation result, which is the most portable thing here. A model can improve on validation loss, hold its benchmark average, retrieve from 32K, and have quietly lost the ability to tell two similar labels apart. If you are grafting anything, measure that before you ship.

Next

A final run on the recipe this study arrived at. Freeze the upper layers, keep the encoder training, 15% code in the mix, a 256-token tail, and the split chosen by workload rather than by default. Then the 9B graft that the freezing result makes affordable, and an instruction-tuned model so we can put a real tool-calling benchmark against the probe that started this.

The source, trained checkpoints, datasets and run logs behind this study are available to sponsors.

Become a sponsor