cs.AIOct 31, 2026
SaveA Tool in the Forward Pass: Frozen Tape Modules as Layers of a Small Language Model
Organizations: University of Passau, Passau, Germany · IT:U Interdisciplinary Transformation University Austria, Linz, Austria
Abstract
Calling a tool through generated text adds decoding steps, even when the computation itself is short. We study an alternative: a small neural module that runs inside a language model's forward pass. A learned writer turns hidden states into a symbolic tape, a short sequence of discrete symbols that encodes the computation. The module, trained separately and then frozen, reduces the tape to a result, and a learned reader returns that result to the model. We evaluate arithmetic and object tracking in gemma-3-1b-it. With a writer trained on labelled tapes, the layer answers all 750 items of BIG-Bench Hard's tracking shuffled objects correctly. At seven people it needs about 1/10 of the GPU kernel time of a text-tool baseline that uses the same host model and exact code, with accuracy 1.00 against 0.99. The writer can also be learned from answers alone, guided by a hand-written copying prior, but less reliably: 5 of 6 runs succeed in the 1B host and 2 of 4 in the 4B host. Two independently trained layers keep their task accuracy when attached to one model, and a learned gate keeps them silent on the tested general-text and look-alike prompts. The results support a narrow use case: a bounded computation that returns one symbol can avoid the decoding cost of a textual tool call. Extracting the right input reliably, and deciding when to act, remain the main open problems.
Figures & tables
Figure 1. One computation, three interfaces. (a) The model generates a program for an interpreter. (b) It generates only the inputs to a fixed tool. (c) A learned writer constructs those inputs directly from hidden states. A frozen neural module computes the result, and a reader returns it to the model. The tape traces Bob’s final object back through the swaps to its original owner. (d) The layer’s components. (e) Measured GPU kernel time at seven people, relative to the layer, with each method’s accuracy (§ 5.2 ).
| module, trained from its final tape only | trained on | at the longest trained | at 4 | 16-place carry chain | formed |
|---|---|---|---|---|---|
| addition, discrete tape (straight-through) | 1–4 digits | 0.08–1.00 | 0.00–1.00 | 0.00–1.00 | 3 of 4 |
| + a continuous state per cell | 1–4 digits | 1.00 | 0.98–1.00 | 0.99–1.00 | 4 of 4 |
| trained soft, used hard | 1–4 digits | 0.02–0.30 | 0.00 | – | 0 of 4 |
| multiplication (either tape) | 1–4 digits | 0.00 | 0.00 | – | 0 of 8 |
| operator order: selector + column, length curriculum | 1–6 operators | 1.00 | 0.97–1.00 | – | 4 of 4 |
| no curriculum | 1–6 operators | 0.19–1.00 | 0.12–1.00 | – | 2 of 4 |
Table 1. Modules trained alone, from their final tape only (four seeds each, used hard). Accuracy at the longest trained length and at four times it (sixteen digits, twenty-four operators; ∗ fourteen swaps, 2.3 times). “Formed”: seeds exact at the longest trained length. Hand-written: the tape layout, the splice and the number of rounds.
| how the writer is trained | tape labels | 4–6 ops | unseen templates | 12 ops | tapes exact | swap | reader off |
|---|---|---|---|---|---|---|---|
| adapters only, no tape | – | 0.35–0.38 | 0.28–0.34 | 0.10–0.12 | – | – | – |
| answers only, gradient through the soft module | none | 0.37 | 0.29–0.31 | 0.10–0.12 | 0.00 | 0.28–0.36 | 0.12 |
| answers only, straight-through | none | 0.34–0.41 | 0.29–0.33 | 0.09 | 0.00 | 0.32–0.34 | 0.11–0.16 |
| answers only, reward (module a black box) | none | 0.36–0.39 | 0.30–0.31 | 0.07–0.11 | 0.00 | 0.32–0.34 | 0.14–0.17 |
| answers only, ReinMax ‡ | none | 0.24–0.30 | 0.17–0.18 | 0.12 | 0.00 | 0.20–0.27 | 0.20–0.25 |
| answers only, I-MLE ‡ | none | 0.24–0.25 | 0.16–0.20 | 0.08–0.10 | 0.00 | 0.20–0.21 | 0.18–0.21 |
Table 2. The frozen module as a layer of gemma-3-1b-it on the development task, by how its writer is trained (two seeds; three for the first row, the tape-labels row and the last label-free row; used hard; trained on one to six operators). “Swap”: accuracy with exact code in place of the learned module; “reader off”: with the result not added back. Unseen templates: two sentence templates never trained on; † a set of eight, six of them new, not comparable with the rows above; ‡ adapters below the writer frozen. (code) : hand-written.
| accuracy | wall-clock (ms) | GPU kernels (ms) | floor | ||||
|---|---|---|---|---|---|---|---|
| per question | 4–6 ops | unseen templates | B1 | B16 | B1 | B16 | (ms) |
| text tool: adapters copy the expression, exact evaluator (code) | 0.99–1.00 | 0.93–0.98 | 235 | 255 | 103 | 138 | 32 |
| tape layer (frozen module after layer 8) | 1.00 | 0.55–0.59 | 20 | 25 | 8.9 | 16.0 | 2 |
| gemma-3-4b-it, chain of thought | 0.94 | – | 3091 | 8010 | – | – | – |
Table 3. Cost at near-equal accuracy on the development task (gemma-3-1b-it unless named; RTX 4090, bf16; four to six operators). Wall-clock: the standard decoder with a dynamic cache. GPU kernels: summed kernel time from the profiler. B1, B16: batch size one and sixteen. Floor: the model’s weights read once per forward pass at 1 TB/s.
| accuracy, people | per item at 7 people | ||||||||
|---|---|---|---|---|---|---|---|---|---|
| 3 | 5 | 7 | tokens | wall B1 (ms) | kernels B1 (ms) | wall B16 (ms) | floor (ms) | layer | |
| tape layer: frozen swap module after layer 8, writer taught which tokens to mark | 1.00 | 1.00 | 1.00 | 0 | 31 | 11.0 | 5.1 | 2 | 1 |
| text tool: adapters write the names, code runs the swaps | 1.00 | 0.98 | 0.99 | 17 | 328 | 112.4 | 23.5 | 34 | 10 |
| gemma-3-1b-it, chain of thought | 0.85 | 0.57 | 0.34 | 485 | 8099 | 2842.8 | 597.3 | 970 | 257 |
| gemma-3-1b-it, program-of-thought | 0.05 | 0.07 | 0.05 | 439 | 7372 | 2124.9 | 838.5 | 878 | 192 |
| gemma-3-4b-it, chain of thought | 1.00 | 0.96 | 0.93 | 534 | 11832 | 6864.0 | 927.1 | 4596 | 621 |
Table 4. Cost per item on the held-out task (tracking shuffled objects, BIG-Bench Hard, 750 items; RTX 4090, bf16). Accuracy on all items (chain of thought: Table 5 ); cost at seven people. Tokens: generated per item (the tool’s includes the answer step). Wall B1: Hugging Face generate with a dynamic cache, one item at a time, mean of eight. Kernels B1: summed GPU kernel time from the profiler, mean of two. Wall B16: per item in a batch of sixteen. Floor: forward passes weight bytes at 1 TB/s. layer: kernel time at batch 1 relative to the layer. Program of thought: the model’s program is run in a sandbox, with a call appended when it defines but never calls a function (PAL’s convention); the last printed line is matched to an option. (code) in the tool: the name-to-index map and the swaps.
| BBH, people (= swaps) | synthetic | tapes exact | reader off | tokens | |||
|---|---|---|---|---|---|---|---|
| model | 3 | 5 | 7 | 7 people, 10 swaps | BBH 7 | BBH 7 | generated |
| gemma-3-1b-it, answers directly | 0.32 | 0.21 | 0.13 | – | – | – | 4 |
| gemma-3-1b-it, chain of thought | 0.85 | 0.57 | 0.34 | – | – | – | 492 |
| gemma-3-4b-it, answers directly | 0.31 | 0.12 | 0.08 | – | – | – | 5 |
| gemma-3-4b-it, chain of thought | 1.00 | 0.96 | 0.93 | – | – | – | 536 |
| 1B + adapters on our stories, no tape | 0.98–1.00 | 0.93–0.95 | 0.78–0.83 | 0.46–0.54 | – | – | 0 |
Table 5. The held-out task: tracking shuffled objects (BIG-Bench Hard, 750 items), never trained on. gemma rows: zero-shot, greedy. Middle rows: gemma-3-1b-it with adapters trained on our own stories (three seeds for the first and third rows, two for the others), with and without the frozen swap module as a layer; the answer is read over the option letters in the same forward pass. Bottom rows: gemma-3-1b-it frozen, no adapters, only the writer and reader trained — with taught marks (two seeds), and from answers alone after a hand-written name prior (two variants of the prior, runs split by whether their benchmark tapes are exact). Chance at 3 / 5 / 7 people: 0.33 / 0.20 / 0.14.
| rewritten, people | original, people | tape right | fires | |||||
|---|---|---|---|---|---|---|---|---|
| 3 | 5 | 7 | 3 | 5 | 7 | rewritten, 7 | rewritten | |
| layer, adapters, marks taught | 0.43–0.81 | 0.40–0.86 | 0.27–0.85 | 1.00 | 1.00 | 1.00 | 0.21–0.82 | 1.00 |
| layer, no adapters, marks taught | 0.94–0.96 | 0.94 | 0.87–0.88 | 1.00 | 1.00 | 1.00 | 0.87–0.88 | 1.00 |
| layer, no adapters, from answers alone | 0.94–0.95 | 0.91 | 0.90–0.94 | 1.00 | 1.00 | 1.00 | 0.90–0.94 | 1.00 |
| with the prompt gate of Table 8 | 0.67 | 0.45 | 0.62–0.63 | 1.00 | 1.00 | 0.97–1.00 | 0.76–0.78 | 0.39–0.53 |
| gate trained with rewrites of our stories | 0.84–0.86 | 0.83–0.84 | 0.82–0.83 | 1.00 | 1.00 | 0.97–1.00 | 0.76–0.78 | 0.73–0.88 |
Table 6. Open wording. 234 judged rewrites of tracking shuffled objects (the first 100 items per split; the hand-written parser reads none), and the same items in their original wording. Nothing is trained on them. Layers read the answer in one forward pass (seeds or runs as in Table 5 ; ranges over them). “Tape right”: the written tape, run through exact code, gives the answer. “Fires”: the share of rewrites on which the layer acts (otherwise the base model answers). Chance at 3 / 5 / 7 people: 0.33 / 0.20 / 0.14.
| alone | together | change | tapes exact | fires: other layer; arithmetic / swap writer | |
| arithmetic, 4–6 operators | 1.00 | 1.00 | +0.000 | 1.00 | 0.00 |
| unseen templates | 0.39–0.73 | 0.39–0.73 | +0.000 | 0.34–0.69 | 0.01 |
| 12 operators | 0.96–1.00 | 0.96–1.00 | +0.000 | 0.95–0.99 | 0.00 |
| state tracking, 7 people, 10 swaps | 1.00 | 1.00 | +0.000 | 1.00 | 0.00 |
| BBH, 3 people | 1.00 | 1.00 | +0.000 | 1.00 | 0.00 |
| BBH, 5 people | 1.00 | 1.00 | +0.000 | 1.00 | 0.00 |
Table 7. Two frozen modules as two layers of one frozen gemma-3-1b-it , each layer trained alone, no adapters (two seeds, used hard). “Together”: both layers attached. “Fires”: the share of prompts on which a writer marks a token — in the task rows the other task’s writer, in the bottom rows the arithmetic / the swap writer. Hand-written: the packing, the gate (an empty tape adds nothing), the tape languages.
| MMLU (base 0.382) | fires: general text | fires: BBH logic | fires: GSM8K | BBH tracking, 7 | |
|---|---|---|---|---|---|
| writer trained on general text; marks nothing adds nothing | 0.368–0.370 | 7.2–9.0 % | 86.8–92.4 % | 54.0–80.8 % | 1.00 |
| + question negatives, well-formed tape | 0.376–0.386 | 7.2–9.0 % | 36.0–38.0 % | 28.2–37.2 % | 0.99–1.00 |
| + the tape must compute (an operator, a swap) | 0.382 | 0.4 % | 0.0–4.8 % | 0.2–0.6 % | 0.99–1.00 |
| + a learned gate on the whole prompt | 0.382 | 0.0 % | 0.0 % | 0.0 % | 0.99–1.00 |
Table 8. Keeping two layers silent on general text (both layers attached, two seeds each). “Fires”: the share of prompts on which the more active writer adds its result — general text: the Pile and BoolQ; BBH logic: logical deduction with five objects (no exchanges); GSM8K: word problems (test split). None of these is trained on, except that the gate’s negatives include GSM8K’s training split. Hand-written: the well-formedness check and the requirement that the tape computes; learned: the question negatives and the gate.
Appendix figures & tables1 asset
Supplementary material from the paper’s appendix.
Appendix
| label (measurement) | prediction | outcome |
|---|---|---|
| e2e1b-a (operator order) | at 6 operators on 3 of 4; at 24 on 2 of 4 | confirmed (4 of 4; 0.97–1.00) |
| e2e1b-b, c (curriculum; straight-through) | without curriculum lower; straight-through | refuted on the median (forms on 2 of 4); confirmed |
| e2e1b-d, e (no-ops) | twice the rounds change ; cost | confirmed |
| e2e2-a, b (layer) | adapters alone ; tape labels , swap within 0.02 | confirmed (0.35–0.38; 1.00, swap identical) |
| e2e2-c (answers only) | on one seed | refuted (0.37, no exact tape) |
| e2e2-d, e (warm start; reader off) | warm within 0.05 of labels; reader off | confirmed, d only through the checkpoint kept while supervised |
Table 9. Registered predictions and their outcomes.
More in this topic
LLM Tool Use
Large language model tool use studies how models select, call, and coordinate external tools to extend their reasoning and take actions.
See all 307 papers