Building a model's inputs from shared parts

The question and why it matters

Every language model starts the same way. Text is cut into tokens, each token is replaced by an integer id, and the id selects one row of a learned table. That row, a list of a few hundred numbers called the token's embedding, is what the first layer of the network reads. The table has one row per token and nothing else: the row for a token is the same in every sentence, and two rows share nothing unless training happens to make them similar.

This study replaces the table with a small constructor. The constructor keeps one shared dictionary of K = 128 learned vectors, called atoms. For every position in the input it computes a weight for each atom from the current token and the tokens before it, keeps the k = 4 largest weights and sets the rest to zero, and adds up the four weighted atoms. The result is the vector the network reads. A small shared module can then adjust the weights a few times (J = 2 by default) before the vector is handed on. Everything after the front end is an ordinary transformer and stays the same in every condition.

The question, as the research plan states it: does a sparse, shared-dictionary constraint on a context-conditioned input constructor change accuracy on held-out combinations relative to an equally expensive dense constructor, and does the value of shared iteration depend on that constraint?

The comparison is scored on held-out combinations. Neural networks are good at examples like the ones they trained on and often poor at new combinations of parts they have seen separately. The standard demonstration is SCAN, a toy command language: a model that has seen jump only on its own and has seen walk twice and run twice will often fail on jump twice. The ability to handle such combinations is called compositional generalization. A front end that must build every input from a few shared parts might help with it, because the same parts serve many tokens and so may carry meaning that transfers to new combinations.

Parts of this idea have been tried. A 2026 paper builds embeddings from a dictionary with weights that depend on context. A 2024 paper adds a learned code at the input and reaches 99 % on the SCAN split described above. What has not been measured is the comparison this study runs: sparse against an equally large dense constructor, with and without shared refinement, scored on held-out combinations at matched parameter counts and compute. A gain for the sparse version would point to a cheap change at the input of any model. Equal scores would show that the context, and not the sparsity, does the work. The plan puts the chance that the sparse constructor beats the dense one by more than seed noise at 30 to 40 %.

The models are small on purpose: 10.7 million parameters, 6 layers, 384 numbers wide. A run takes minutes on one graphics processor, so the study can afford 270 runs: three random seeds per setting and a control for each comparison.

The basics, built for this problem

Tokens in, next-token predictions out

Each example is a pair: a source (a program, or a SCAN command) and a target (the program's result, or the action sequence). Both are sequences of tokens from a small vocabulary: 19 symbols for most program tasks, 25 for the rebinding task, 23 for SCAN. Four ids are reserved: padding, a begin mark, a separator and an end mark. The model sees one row of ids: begin, source, separator, target, end.

The model is a decoder: at every position it produces a score for each vocabulary entry, predicting the next token from the tokens to its left and never from the ones to its right. That restriction is called causal. During training the model sees the whole row at once and is graded at every target position. At test time it gets only the begin mark, the source and the separator, and writes the target itself, one token at a time, always taking the highest-scoring token (greedy decoding), until it writes the end mark.

What training means here

The model has about 10.7 million adjustable numbers, its parameters. Training repeats one step 12,000 times (15,000 for SCAN): take a batch of 128 examples, compute the cross-entropy (the average of minus the log probability the model gave to each correct target token), work out how each parameter should move to lower it, and move every parameter a little. How far it moves is set by the learning rate. Every model starts from random parameters; nothing is pretrained.

Exact match

A generated target counts as right only if every token matches and the end mark comes in the right place. One wrong digit makes the whole example wrong. The score of a run on a set of examples is the fraction it gets entirely right, its exact match. Per-token accuracy is logged too, but the hypotheses are read on exact match.

What "generalize" means here

Every task has at least two evaluation sets, both disjoint from training (no evaluation source appears among the training sources):

• the in-distribution or iid split: fresh examples drawn the same way as the training examples;
• one or more held-out splits: examples built around a combination that never occurs in training, such as one operation directly after another, an operation with a particular argument, or a longer program than any seen.

The two can be far apart. In the check run described in section 3, the plain model scored 1.000 on the iid split of the depth task and 0.000 on the held-out depths. The iid split is the control. A front end that raises held-out accuracy and leaves iid accuracy alone has changed how the model generalizes. A front end that raises both may only have given the model more capacity or an easier optimization.

Seeds and seed noise

Each setting is trained three times, with seeds 0, 1 and 2. The seed changes the random starting parameters and the order and sampling of training examples. It does not change which combinations are held out: that design comes from a separate fixed seed, so the three runs are replicates of one design. The spread between seeds, the seed noise, is the yardstick for every comparison in section 6.

The data

Programs on a 7 by 7 grid

Four of the six tasks are generated programs. The state is a pair (x, y) of integers modulo 7, a point on a 7 by 7 grid that wraps around at the edges. A program gives a start state and then a list of operations; the generator works out the answer exactly, so the targets are always right.

The target is a trace: the state after each operation. For the combination task the target is the state after every second operation (and after the last), so each output step has to compose a pair of operations in one go. An earlier form asked for the final state only; the plain model learned it too slowly (12 % exact after 4,000 steps), and a state after every single operation made the held-out pairs trivial, since no pair ever had to be composed.

Each program task has 100,000 training programs, 2,000 iid examples and 2,000 examples in each held-out split. The generator balances how often each operation appears in training.

What each task holds out

The rebinding task is the one the hypothesis speaks to most directly. An episode starts with three demonstrations of a made-up word, for example S 5 5 fa 5 = 5 3 ;, which shows what fa does to a state, and then a short program that uses fa among ordinary operations. The demonstrations are drawn until they rule out every other operation with the same kind of argument. In training each alias stands for any operation except one; the held-out split binds each alias to that one operation. The model must build the meaning of the alias from the context, which is what a context-conditioned constructor is for: the sparse version has to express the new meaning as a mix of atoms it already has.

The two SCAN tasks use the official splits of the public benchmark. A command such as run thrice after jump around left maps to a sequence of actions: I_TURN_LEFT I_JUMP four times, then I_RUN three times. The iid split is 1,000 commands taken out of the official training file, which leaves 13,670 training commands for add-jump and 15,990 for length.

The low-data condition, and why it exists

Before any cluster run, each program task was trained once with the plain table model for 2,500 steps on a desktop computer's graphics processor, to check that the baseline learns it and that the held-out split stays hard.

Two readings follow. The depth task looks like a floor for this kind of backbone: it is kept because the original notes asked for it, but no pass line depends on it. Rebinding at 100,000 episodes leaves little room for any front end to differ (0.965 held-out), while at 20,000 episodes the held-out score falls to 0.347. So the study adds a low-data condition: rebinding and role at 20,000 training examples instead of 100,000, for all nine front ends. These check runs are single runs of one condition. They set the form of the tasks and are not results of the study.

The model

The backbone, the same everywhere

The backbone is a decoder-only transformer of 6 blocks, 384 numbers wide, with 6 attention heads and an MLP of width 1,536 with GELU activation in each block. Normalization comes before each sub-layer (pre-norm), positions are encoded by rotating the attention vectors (rotary positions), attention is causal, and the output head has its own weights, separate from any input table. The blocks hold about 10.6 million parameters. Between the front end and the first block sits a normalization with no learned parameters, the same for every condition, so a front end that produces larger or smaller numbers cannot gain an advantage in optimization from scale alone.

The front end, step by step

The contextual front ends work in three stages. An identity table gives each token a short vector u of 128 numbers. A context block, one causal attention block 128 wide, turns the sequence of u into context vectors c, so that c at a position depends on that token and the ones before it. A composer, a two-layer MLP with a hidden width of 256, reads u and c side by side and produces either 128 atom weights (the sparse versions) or 384 numbers directly (the dense versions).

The constraint sits in the sparse code. Of 128 weights, all but the four largest in size are set to zero; the gradient flows through the four that are kept. The vector is then the weighted sum of four atoms.

The constraint has a geometric form. With K = 128 atoms and vectors of 384 numbers, every vector the sparse constructor can make lies in the 128-dimensional space spanned by the dictionary, and with k = 4 it lies in one of the 4-dimensional subspaces spanned by four atoms. Such a set is a union of subspaces. The normalization after the front end is a fixed map and does not lift the bound. The sweep over K (32, 128 and 256) moves the dictionary across the 384-dimension mark, and the sweep over k (2, 4 and 8) loosens or tightens the union.

Refinement

The recurrent versions add a refiner, a small MLP shared across passes. On each pass it reads the current code (or, in the dense control, the current vector) together with u and c, and adds a correction. In the sparse version the four largest weights are kept again after every pass, so a pass can swap one atom for another. The refiner starts as the identity: its last layer is zero at initialization, so refinement begins by changing nothing and learns what to change.

The nine front ends

The comparisons the hypotheses rest on are close in cost. The dense control C has 0.15 % more parameters and 0.15 % more compute per token than D. E exceeds F by 1.3 % in parameters and 2.4 % in compute at J = 2 (4.6 % compute at J = 4). The extra block of A+ costs 16.6 % more, so it is a larger control rather than an equal one. Q is a simplified stand-in for the published method, and its token vector receives gradient through two paths, once directly and once through the matched code.

Every front end must be causal, like the backbone: the vector at a position may depend on that token and earlier ones only. A self-test checks this for all nine before any training on the cluster, by changing later tokens and confirming that vectors and scores at earlier positions stay the same.

The training schema

Objective and optimizer

The loss is the cross-entropy on the target tokens and the end mark. Two front ends add a small extra term: the soft-threshold variant adds an L1 penalty (10⁻³ times the mean absolute weight) that pushes weights to zero, and Q adds the usual pull between each token vector and its nearest code. The optimizer is AdamW (momentum terms 0.9 and 0.95, weight decay 0.1), gradients are clipped at norm 1.0, and the batch is 128 examples. The learning rate climbs linearly for 500 steps to 0.001 and then follows half a cosine down to 0.0001.

Every 2,000 steps the run scores a fixed random subset of 2,000 examples from each evaluation split. At the end it scores every split in full, saves 24 decoded examples per split so failures can be read by eye, and, for the sparse versions, records how often each atom was used.

The grid of runs

A grid is a table with one row per run: task, front end, size, seed, steps and any extra settings. The study has two.

How a run is packed onto a node

The cluster's GPU nodes each have 128 CPU cores, 512 GB of memory and four NVIDIA A100 graphics processors with 40 GB each. One job takes a whole node and starts four slots, one per graphics processor. Each slot walks down the grid table and claims the first run that is neither finished nor claimed, by creating a lock folder for it; creating a folder either succeeds or fails in one operation, so two slots can never claim the same run. A slot trains its run, then claims the next, until no rows are left. Runs are short (a little over four minutes on average in grid1), so a node finishes several dozen per hour.

Chunks and resuming

The grids go to the scheduler as a chain of 4-hour jobs, called chunks; each chunk starts when the one before it ends. Short jobs start much sooner: on the day of submission a 4-hour whole-node job was estimated to start in 1.6 hours and a 12-hour one in 18.7 hours. When a chunk reaches its time limit, the scheduler signals every slot, and each run in progress saves a checkpoint (parameters, optimizer state, position in the shuffled data, random number states and elapsed time) and stops. The next chunk's slots find those runs locked by a job that has ended, take them over and continue from the checkpoint with the same data order and random state it had when it stopped. Runs also save a checkpoint every 2,000 steps. A run is finished when it has written its final metrics file, and finished runs are skipped, so submitting the same chain again carries on where it stopped.

Hardware, stage by stage

Moving between machines changes the arithmetic and the speed; the batch, steps and learning rate stay the same. On the A100s the forward and backward passes run in bfloat16, a 16-bit format with the range of 32-bit floats, while the desktop runs used full 32-bit arithmetic. The evaluation uses the same precision as training, so the four atoms chosen at test time are the ones the model trained with. The cluster also has nodes with AMD MI300A accelerators; the same code would run there from a second container image built on ROCm in place of CUDA, and each run records the device it ran on. grid1 and grid1b ran only on the A100 nodes.

Evaluation and the pass lines

Each number in the comparisons is the exact match on the held-out split, from the final model, scored on the whole split, averaged over the three seeds. One condition beats another when the difference between their means is larger than twice the pooled seed standard deviation.

The pass lines, quoted from the research plan:

• H1 (sparse reuse): D beats C on at least two of {combo, role, role-20k, rebind-20k} with iid within 2 points.
• H2 (learning the dictionary matters): D beats D₀ wherever H1 holds.
• H3 (iteration × constraint): the gain from J = 1 → 4 is larger for F than for E on the cells that carry the J sweep, progs-combo and progs-rebind-20k (grid1b); a monotone curve in J for F and flat for E is the clean version.
• H4 (context matters): D beats B on rebind (B cannot do rebinding at the input at all; this is a sanity line, not a result).
• Reference: Q must reach > 90 % on scan-addjump (the published result for the full method is 99 %); if it does not, our proxy is weak and the published code is run before any claim.

H1 asks whether the sparse constructor generalizes better than its dense twin on the tasks with room to move, while matching it in distribution. H2 asks whether the gain needs a learned dictionary or whether random atoms would do. H3 asks whether refinement helps more when it works on a sparse code. H4 checks that a front end with context can do what a context-free one cannot. The reference line checks that the simplified code-at-the-input baseline is strong enough to stand in for the published method.

Null splits and checks

• On the iid splits and on rebind_seen no condition should differ from the plain table by more than seed noise. A difference there points to a bug or a change in capacity.
• The causality self-test passes on the cluster before any run.
• A sparse condition that uses fewer than a quarter of its atoms, or whose atom-usage perplexity is below 8, has collapsed onto a few atoms; that is reported and the K sweep consulted before its scores are read.
• Ten failures per comparison are read by hand from the saved examples.
• Only the final scores count. Runs also record the best score seen during training, but that picks a step using the held-out split itself, so it is never quoted.

What a null result would mean

Each way the hypotheses can fail has its own reading, within the parameter and compute budgets tested:

If neither D nor F wins anywhere, the study ends with that finding. The next phase, the same comparison at about 110 million parameters on natural text, runs only if one of them wins somewhere.

The status page and the replay

The status page is refreshed from the cluster about every 10 minutes while one of us is logged in. When the figures are more than 30 minutes old the page says so and shows when they were taken.

The tiles and sections

• Finished: runs with a final metrics file, out of the grid's total, with a bar.
• Chain: which chunk is running or waiting.
• Time left: running time the unfinished runs need at the current pace, and whether the queued chunks cover it.
• Cost: core-hours so far and projected at the end (128 core-hours per node-hour).
• Chain section: the timing job and each chunk with its state, run time, node and how many runs were finished when it ended.
• Running: one row per graphics processor with the run, step of steps, time left, last loss and tokens per second. A run silent for 20 minutes is marked quiet.
• Failed and paused: runs that stopped with an error, and runs saved at a chunk limit that the next chunk will resume.
• Results: tasks down, front ends across. Each cell is the held-out exact match, mean over seeds, with ± the standard deviation, or a count such as 2/3 while seeds are still running; the best in each row is bold. Switches show the iid scores and the sweeps over K, k, J and the soft threshold. Run names read as task, front end and seed: progs-rebind-20k__sparse_ctx__s2.
• Dictionary use: for the sparse front ends, how many of the atoms are used at all, how many are active per token, and the usage perplexity (the effective number of atoms in use; 128 would mean all used equally).
• Cost and disk: totals so far, the projection and its basis, and disk space under the run folders.

The replay

The replay plays the training back as it happened. Play, pause and the step buttons move a scrubber through time; a clock shows the time of day on the cluster and a second readout the core-hours spent up to that moment. Below them, one chart per measure (training loss, held-out exact match, iid exact match) draws one line per front end, and a marker moves along the step axis. An event list beside the charts names what happened at each moment: a chunk starting, runs finishing, a chunk reaching its limit, a run resuming. Each line holds at most 500 points, taken evenly from the run's log.

The held-out lines can rise late, or not at all, after the loss has flattened; that gap is the difference between fitting the training set and generalizing. The pass lines are read on the last points of the C and D lines, and of E and F, on the tasks they name.

Glossary

A100

An NVIDIA graphics processor with 40 GB of memory; each GPU node of the cluster has four.

AdamW

The optimizer used: it scales each parameter's step by running averages of its gradient and squared gradient, and shrinks weights slightly each step (weight decay).

atom

One of the K learned vectors in the dictionary; every constructed input is a weighted sum of a few atoms.

batch

The 128 examples used for one training step.

bfloat16

A 16-bit number format with the same range as 32-bit floats and less precision; used for the arithmetic on the A100s.

causal

Depending only on the current and earlier positions, never later ones.

checkpoint

A saved copy of everything a run needs to continue: parameters, optimizer state, data position, random number states.

chunk

One 4-hour job in the chain that carries a grid; the next chunk resumes where it stopped.

code

The 128 atom weights for one position, of which k are nonzero.

composer

The small MLP that turns identity and context vectors into atom weights or a dense vector.

compositional generalization

Handling new combinations of parts that were each seen in training.

core-hour, node-hour

One CPU core for one hour; a whole node of 128 cores for one hour is a node-hour, 128 core-hours.

cross-entropy

The training loss: the average of minus the log probability given to each correct token.

dictionary

The shared set of K atoms, 128 by default, each 384 numbers long.

embedding

The vector of numbers that stands for a token at the input of the network.

exact match

The fraction of examples whose generated target is entirely right.

held-out split

Evaluation examples built around a combination that never occurs in training.

iid split

Evaluation examples drawn the same way as the training examples (independent and identically distributed), none of them seen in training.

J

The number of refinement passes; 2 by default, 1 and 4 in the sweep.

K, k

The number of atoms in the dictionary (128), and the number active at each position (4).

learning rate

How far each training step moves the parameters; here up to 0.001.

MLP

A multilayer perceptron: linear maps with a nonlinearity between them.

node

One computer in the cluster; the GPU nodes have 128 CPU cores, 512 GB of memory and four A100s.

parameter

One adjustable number in the model.

pass line

A condition written into the plan before the runs that a result must meet to count.

refiner

The small shared MLP that adjusts the code or vector on each of the J passes.

rotary positions

A way of telling attention where tokens are by rotating its query and key vectors by an angle that depends on position.

SCAN

A public benchmark of simple navigation commands and their action sequences, with splits that test compositional generalization.

seed

The number that fixes a run's random choices; each setting runs with seeds 0, 1 and 2.

slot

One of the four workers in a job, each bound to one graphics processor and running one training run at a time.

soft threshold

An alternative to keeping the k largest weights: every weight shrinks toward zero by a learned amount, and small ones become zero.

token

One symbol of the input or output: a digit, an operation name, a word of SCAN.

trace

A target listing the state after each operation (or after every second one).

transformer block

One layer of the backbone: causal attention across positions followed by an MLP at each position, each with a normalization and a skip connection.

union of subspaces

The set of vectors reachable when only k atoms may be active: one k-dimensional flat piece for each choice of k atoms.

usage perplexity

The effective number of atoms in use, from how evenly they are used; low values mean a few atoms do most of the work.