Prefix-Tuning for Arbitrary Output Sequences on Pretrained Transformers
Peter Cho-Ho Lam1, Ziyi Wang1, Zirui Zhou1
1Huawei Technologies Canada
NeurIPS 2026
Prefix-tuning adapts a pretrained language model by prepending a few learnable vectors to the keys and values of every attention layer, leaving all weights frozen. How expressive is this small set of parameters? We show that, under mild conditions, a prefix can be constructed in closed form that makes a pretrained autoregressive Transformer generate any output sequence the model can already produce, for any input of the same length that ends with the same token. The construction, attention matching, needs no training: it records the attention outputs of a run that produces the target, and solves one small linear system per attention head so that the prefix reproduces them. This repository contains the implementation for GPT-2 (multi-head attention) and Qwen3-4B (grouped-query attention with QK-Norm), the datasets, and the scripts that reproduce the results of the paper.
👉 Try our demo: pick a sentence, and Qwen3-4B with the prefix will say it, whatever you type.
Contents
File structure
.
├── attention_matching/ # the method, as a small Python package
│ ├── solver.py # the linear system S W = R and the prefix keys
│ ├── adapters.py # GPT-2 and Qwen3: queries/keys/values, prefix injection
│ └── prefix.py # build_prefix(), Prefix (save/load), load_model()
├── data/
│ ├── gpt2/{and,but,comma,dot}.txt # datasets of Table 1 (GPT-2 tokenizer)
│ └── qwen3/{and,but,comma,dot}.txt # the same datasets for the Qwen3 tokenizer
├── docs/ # project page source, built and published by GitHub Actions
├── examples/
│ └── quickstart.py # build one prefix and apply it to unseen inputs
├── scripts/
│ ├── evaluate.py # success rates under the protocol of Section 5
│ └── make_table.py # tables from the results of evaluate.py
├── verify/
│ ├── check_matching.py # hidden states of the prefixed run vs. the source run
│ └── check_projection_rank.py # every prefix key is reachable from an input vector
└── results/ # per-cell results of the paper's experiments
Environment setup
Python 3.10 or later and a recent PyTorch are required. The patched attention depends on the internals of Hugging Face Transformers, so the version is pinned.
conda create -n attention-matching python=3.12 -y
conda activate attention-matching
pip install -r requirements.txt
Download
The datasets are part of this repository: data/gpt2/ for GPT-2 and data/qwen3/ for Qwen3-4B, built from the same sentences for the respective tokenizer. Each of the four datasets contains five groups of 50 sentences with exactly 5, 10, 20, 50 and 100 tokens, and all its sentences end with the same token ( and, but, , or .).
The models are downloaded from the Hugging Face Hub on first use; to fetch them in advance:
huggingface-cli download gpt2
huggingface-cli download Qwen/Qwen3-4B
Usage
Python. Build a prefix from a source input x', then generate from any input of the same length that ends with the same token:
from attention_matching import adapter_for, build_prefix, greedy, load_model
model, tok = load_model("Qwen/Qwen3-4B") # or "gpt2"
adapter = adapter_for(model) # patches the attention layers
source = tok("The old man sat and")["input_ids"]
prefix = build_prefix(adapter, source, k=10) # one linear system per layer and KV head
print(tok.decode(prefix.target_ids)) # y: the model's own continuation of x'
x = tok("End the long night and")["input_ids"] # n = p, and x ends like x'
with adapter.use(prefix, repeat=1): # repeat: the factor r of Section 5
print(tok.decode(greedy(model, x, prefix.k)[0])) # prints y
prefix.save("prefix.safetensors") # Prefix.load(...) reads it back
Reproducing the tables. scripts/evaluate.py evaluates every ordered pair of sentences within each length group, for output lengths k ∈ {1, 5, 10}. We ran all experiments on a single GPU with 24 GB of memory, with model weights in float32 and the linear systems solved in float64; GPT-2 also runs on a CPU, which is enough for the example.
# Table 1: GPT-2, attn_lim = 20, r = 90
python scripts/evaluate.py --model gpt2 --repeats 90 --out results/gpt2.json
python scripts/make_table.py results/gpt2.json
# Section 6: Qwen3-4B, r = 1 and r = 10
python scripts/evaluate.py --model Qwen/Qwen3-4B --repeats 1 10 --out results/qwen3.json
python scripts/make_table.py results/qwen3.json --summary
--datasets, --in-lens and --out-lens select a subset of the grid. Results are appended to the output file, so an interrupted run resumes where it stopped.
Checks.
python verify/check_matching.py --model gpt2 # hidden states match the source run
python verify/check_projection_rank.py --model Qwen/Qwen3-4B
Example
python examples/quickstart.py builds a prefix from the first 10-token sentence of data/gpt2/and.txt and applies it to the next five:
This is the beginning of the README. Read the rest at the original. Read the full README ↗
