A Tool Inside the Forward Pass Cuts Text-Calling Costs in Narrow AI Tasks

Moving the call inside the model

In the preprint A Tool in the Forward Pass: Frozen Tape Modules as Layers of a Small Language Model, dated 2026-10-31, Michael Granitzer and colleagues examine a specific source of computational overhead: the tokens a language model generates to call a tool. Their alternative places a small neural executor inside the model’s forward pass. The reported benefit concerns bounded computations with single-symbol results, not a general replacement for external tools.

The interface has three parts. A learned writer reads the model’s hidden states and selects symbols from the prompt to form a discrete tape. A separately trained, frozen module repeatedly reduces nearby tape cells until a result remains. A learned reader adds that result back into the model’s hidden state, allowing the remaining layers to produce an answer.

The researchers test arithmetic modulo ten and stories in which people exchange objects. For arithmetic, the tape contains an expression. For object tracking, it contains the exchange participants followed by the person being asked about. The module traces that person’s final object backwards through the exchanges to its original owner.

Accuracy with evidence of tool use

With supervised input extraction, the authors report correct answers on all 750 items of BIG-Bench Hard’s tracking shuffled objects benchmark using gemma-3-1b-it. Training used synthetic stories, not those benchmark items.

Importantly, the study checks whether the model actually uses the module. Replacing the learned executor with exact code preserves accuracy on the evaluated inputs. Removing the reader’s contribution drops performance to chance. These checks matter because an earlier version achieved substantial accuracy while writing incorrect tapes: trainable adapters were answering around the tool rather than relying on it.

The benchmark also needs a qualification. Although its items were excluded from training, findings from object tracking informed later decisions about supervision, training and gating. The authors therefore describe it as a transfer case study, not an untouched test of the final design.

What the timing comparison establishes

The closest comparison uses the same host model, trained to generate the tool’s inputs as text, with exact code performing the exchanges. At seven people, this text-tool baseline achieves 99% accuracy against the tape layer’s 100%.

At batch size one, the authors measure 112.4 milliseconds of GPU kernel time per item for the text tool and 11.0 milliseconds for the layer: approximately a tenfold difference. Standard-decoder wall-clock measurements are 328 and 31 milliseconds, respectively. The text tool generates 17 tokens per item, including its answer step; the layer reads the answer in the same forward pass that processes the prompt.

This is an implementation-specific result, not a universal speedup. Measurements used one RTX 4090, small timing samples and batch sizes one and sixteen. Batching narrows the kernel-time advantage to roughly threefold per item on object tracking. Speculative decoding was not measured, and longer tapes or additional reduction rounds would increase the layer’s cost.

Learning the input remains difficult

The executor is only part of the problem. The writer must identify which names, numbers or operators belong in the computation and discard the rest.

The authors also train writers without labelled tapes, rewarding sampled tapes that produce correct answers. This still requires hand-written guidance: a copying prior supplies symbol meanings, including a fixed name-to-index mapping. Reward then teaches which tokens to select rather than discovering the symbols’ meanings unaided.

With the host frozen and no adapters, five of six object-tracking runs succeed in the 1B model and two of four in the 4B model. These counts span several training configurations. One failed 1B run exploits the module with a misleading exchange prefix that works on training stories but fails on benchmark items. Validation accuracy does not catch that failure.

Knowing when to stay silent

Separately trained arithmetic and tracking layers can share one frozen model without reducing their task accuracy in the composition experiments. Preventing unwanted activation is harder. Tape checks and a learned prompt gate eventually keep both layers silent on 3,800 tested general-text and look-alike prompts, preserving the base model’s measured MMLU accuracy.

That restraint has a cost: valid prompts with unfamiliar wording can be rejected. The strongest rewritten-prompt results come from a different, ungated configuration. As the authors emphasize, no single configuration reliably demonstrates the main cost result, answer-only learning and conservative composition across new tasks.

The contribution is therefore an interface result, not evidence that learned execution beats code. Exact code serves equally well on the evaluated inputs of both studied tasks. The preprint shows that avoiding textual calls can reduce measured computation costs; reliably extracting the right inputs and deciding when to act remain the central challenges.