Dense Teacher Distillation Experiment

Finished2026-06-21imported · line YPosttraining and Adaptation Structured Execution and Compilers GitHub ↗

The one idea you need

As a program runs, the set of value pairs that two hidden numbers could still be keeps shifting and tangling. The model must capture that whole set on one fixed-size sticky note that never grows, no matter how many possibilities pile up.

The question

If a model is shown the perfect answer about every remaining possibility at each step of a program, can its fixed-size memory learn to track it all exactly?

What we found

No. Even with a flawless teacher revealing the exact set of still-possible answers at every step, the fixed-size memory learned only a rough approximation. The best version placed 52% of its confidence on the correct final answer but just 22% on the exact full set of possibilities, and both collapse as programs get longer. More memory helps but never reaches exact.

Why it matters

If your task must track every remaining possibility exactly, not just a likely answer, piling on better training data won't get a compact memory there. The memory's shape is the limit, so build the bookkeeping structure in instead.

Getting the answer right52%confidence the best model put on the correct final answer for eight-step programs, a partial but useful signal

Getting the whole picture right22%confidence on the exact full set of remaining possibilities at that same length, far from exact

Longer programs, faster collapse22% → 5%full-picture confidence dropping from eight-step to twenty-four-step programs

Bigger memory helps but not enough9% → 19%full-picture confidence at eight-step programs when doubling the memory size

On this page

Results at a glance
Overview
Report
Experiment log
Figures
Data files
Reproduce
Related

Results at a glance 5

Right answer versus the full picture as programs get longer

How to read

Horizontal axis is program length in steps; vertical axis is how much of the model's confidence lands on the target, higher is better. One line tracks the correct final answer, the other the exact full set of still-possible pairs.

Takeaway → The answer line sits far above the full-picture line and both fall as programs lengthen, showing the model learned a useful shortcut, not exact tracking.

Data table

program length	projected query mass	exact pair-belief mass
4	80.5%	50.1%
8	52.1%	21.9%
12	36.5%	11.5%
16	27.9%	7.8%
24	20.5%	5.1%

Numbers from reports/dense_teacher_distillation_paper.md tables (matches analysis/threshold_summary.csv)

Technical framing

Best dense student: projected query mass vs exact belief mass (mod 31) — The best model (residual MLP, width 512) tracks queries far better than the exact belief, so it learned a compressed approximation, not exact execution.

Bigger memory lifts right-answer accuracy the most

How to read

Horizontal axis is program length; vertical axis is confidence on the correct final answer, higher is better. Three lines compare small memory, doubled memory, and doubled memory with a redesigned update rule.

Takeaway → Doubling the memory raises every point the most and the redesigned update adds a smaller steady edge, but all three still decay with program length.

Data table

program length	GRU, width 256	GRU, width 512	residual, width 512
4	69.9%	78.1%	80.5%
8	36.1%	49.2%	52.1%
12	23.7%	35.4%	36.5%
16	16.8%	26.3%	27.9%
24	11.7%	18.7%	20.5%

Numbers from reports/dense_teacher_distillation_paper.md table (matches analysis/threshold_summary.csv)

Technical framing

Projected query mass by program length across main variants (mod 31) — Doubling state width from 256 to 512 gives the biggest gain; the residual transition adds a smaller consistent edge over the GRU.

Exact full-picture accuracy collapses with length for every version

How to read

Horizontal axis is program length; vertical axis is confidence on the exact full set of still-possible pairs, higher is better. Same three versions: small memory, doubled memory, and doubled memory with a redesigned update rule.

Takeaway → Every line sinks toward the bottom as programs lengthen; wider memory helps a little but none stays high, so no version tracks exactly.

Data table

program length	GRU, width 256	GRU, width 512	residual, width 512
4	32.6%	45.7%	50.1%
8	9%	19.3%	21.9%
12	4%	9.9%	11.5%
16	2.5%	6.8%	7.8%
24	1.6%	4.3%	5.1%

Numbers from reports/dense_teacher_distillation_paper.md table (matches analysis/threshold_summary.csv)

Technical framing

Exact pair-belief mass by program length across main variants (mod 31) — Strict belief mass collapses with length for every variant: wider states help but no dense student becomes an exact belief executor.

Easier warm-up: comparing memory designs and sizes

How to read

Four model versions run along the horizontal axis; each has a pair of bars, one for landing on the right answer and one for reproducing the full set of possibilities, on shorter six-step programs. Taller is better.

Takeaway → The version that scores the two hidden numbers separately has clearly the shortest bars, which is why it was dropped from the main tests.

Data table

variant	projected query mass	exact pair-belief mass
GRU MLP d128	61.9%	37.8%
GRU MLP d256	71.3%	52.9%
residual MLP d256	71.6%	52.8%
GRU low-rank d256	50.3%	21.2%

Numbers from reports/dense_teacher_distillation_paper.md pilot table (matches analysis/threshold_summary.csv)

Technical framing

Pilot (mod 11): decoder and width ablation at train length 6 — The low-rank decoder is clearly worse at representing correlated pair beliefs, so it was dropped from the main sweep.

A fresh readout finds no hidden exact answer in the memory

How to read

Horizontal axis is program length; vertical axis is confidence on the exact full set of possibilities, higher is better. One line is the fully trained readout, the other a simple fresh readout attached to the same memory.

Takeaway → The two lines lie almost on top of each other, meaning no hidden exact picture is stored, so the memory itself, not the readout, is the bottleneck.

Data table

Program length	Trained belief decoder	Frozen linear probe
4	50.1%	44.5%
8	21.9%	20%
12	11.5%	10.6%
16	7.8%	7.7%
24	5.1%	5.1%

Numbers from experiments/dense_teacher_distillation/analysis/threshold_summary.csv (main_residual_mlp_d512)

Technical framing

Frozen linear probe vs trained decoder: belief mass at K = L (best variant, mod 31) — The frozen probe matches the trained decoder at every length: no hidden exact belief exists in the state, so the state itself is the bottleneck.

In the author’s words from the Report · “Abstract”

Full prefix-belief distillation greatly improves dense recurrent execution, and increasing state width from 256 to 512 substantially improves both query accuracy and strict belief accuracy. A residual transition cell is consistently better than a GRU at width 512. However, even the best dense student does not learn exact belief execution on modulus 31. At train-length 8, the best model assigns 52.1% mass to the correct query support but only 21.9% mass to the exact pair-belief support. At held-out length 24, those fall to 20.5% and 5.1%.

Overview

This experiment tests whether a fixed-width dense recurrent state can learn exact modular belief-state execution when the training signal is the full teacher belief distribution at every prefix step.

src/dense_teacher_distillation_experiment.py: training, probing, checkpointing, and evaluation harness.
src/analyze_dense_teacher_distillation.py: analysis and figure generation.
reports/dense_teacher_distillation_experiment_log.md: chronological experiment log.
reports/dense_teacher_distillation_paper.md: standalone written report.
reports/dense_teacher_distillation_paper.html: standalone HTML report.
runs/: JSON and CSV run outputs.
analysis/: generated summaries and figures.
checkpoint_manifest.csv: checkpoint paths and sizes.

Checkpoints are written outside the experiment directory under:

../../large_artifacts/dense_teacher_distillation/checkpoints/

Download this experiment directory for the normal research bundle. Download ../../large_artifacts/dense_teacher_distillation/ only when saved model weights are needed.

Report

Rendered from reports/dense_teacher_distillation_paper.md

Abstract

This experiment tests whether a fixed-width dense recurrent state can learn exact belief-state execution when supervision is not the limiting factor. Programs operate on two modular registers (A,B) using arithmetic updates and observation filters. A teacher computes the exact belief distribution over all (A,B) pairs after every prefix. A student recurrent executor receives the same program and initial relation, stores only a dense hidden vector, and is trained to decode the exact teacher belief at each recurrent step.

The result is clear but mixed. Full prefix-belief distillation greatly improves dense recurrent execution, and increasing state width from 256 to 512 substantially improves both query accuracy and strict belief accuracy. A residual transition cell is consistently better than a GRU at width 512. However, even the best dense student does not learn exact belief execution on modulus 31. At train-length 8, the best model assigns 52.1% mass to the correct query support but only 21.9% mass to the exact pair-belief support. At held-out length 24, those fall to 20.5% and 5.1%.

Task

Each example starts with a correlated support over two registers:

B = A + delta (mod p)

The program then applies a sequence of operations:

constant arithmetic updates to A or B
cross-register arithmetic updates such as A = A + B
observation filters such as A mod m = r or B mod m = r

The teacher maintains the exact belief distribution over all p^2 possible (A,B) pairs after each prefix. The query asks for one of four values:

A
B
A + B
A - B

All query values are computed modulo p.

Student Model

The recurrent student has three parts:

An instruction embedding for the current operation and argument.
A recurrent transition over a fixed-width dense hidden state.
A belief decoder that maps the hidden state to a distribution over all (A,B) pairs.

Two transition types were tested:

gru: a GRUCell followed by a residual MLP update.
residual: a gated residual MLP transition.

Two decoder families were tested in the pilot:

mlp: unconstrained MLP logits over all pairs.
low_rank: a mixture of rank-1 distributions over A and B.

The low-rank decoder was dropped from the main sweep because it was consistently worse.

Training

The student is trained with full prefix-belief distillation. At each recurrent step t, the decoder produces logits over all (A,B) pairs, and the training loss is cross-entropy against the teacher belief after prefix t.

The direct query head is left unsupervised in the main runs. It remains in the harness as a diagnostic, but the headline query metric is computed by projecting the decoded pair distribution onto the requested query.

Metrics

The primary metrics are:

decoder_query_target_mass: mass assigned to the correct query support after projecting the decoded pair distribution.
decoder_belief_target_mass: mass assigned to the exact teacher pair-belief support.
probe_belief_target_mass: mass assigned by a separately trained frozen-state probe to the exact pair-belief support.
query_target_mass: direct query-head mass, included only as a diagnostic because the query head is not trained.

The headline row for each program length uses the first evaluation step K such that K >= L, where L is the program length.

Pilot Results

The pilot used modulus 11, train lengths up to 6, and evaluation lengths 3, 6, 9, and 12.

Variant	L=3 Query	L=6 Query	L=9 Query	L=12 Query	L=3 Belief	L=6 Belief	L=9 Belief	L=12 Belief
GRU MLP d128	90.5%	61.9%	44.8%	35.2%	70.6%	37.8%	21.6%	15.0%
GRU MLP d256	94.7%	71.3%	53.1%	41.4%	81.9%	52.9%	32.7%	22.0%
Residual MLP d256	94.8%	71.6%	54.7%	42.6%	81.0%	52.8%	33.8%	22.5%
GRU low-rank d256	78.8%	50.3%	36.4%	28.5%	49.4%	21.2%	12.0%	8.9%

The pilot established three facts:

Increasing state width improved both projected query mass and strict belief mass.
The residual transition was slightly stronger than the GRU on longer lengths.
The low-rank decoder was a poor fit for correlated pair beliefs.

Main Results

The main sweep used modulus 31, train lengths up to 8, and evaluation lengths 4, 8, 12, 16, and 24.

Decoder-Projected Query Mass

Variant	L=4	L=8	L=12	L=16	L=24
GRU MLP d256	69.9%	36.1%	23.7%	16.8%	11.7%
GRU MLP d512	78.1%	49.2%	35.4%	26.3%	18.7%
Residual MLP d512	80.5%	52.1%	36.5%	27.9%	20.5%

Main decoder-query mass

Strict Decoded-Belief Mass

Variant	L=4	L=8	L=12	L=16	L=24
GRU MLP d256	32.6%	9.0%	4.0%	2.5%	1.6%
GRU MLP d512	45.7%	19.3%	9.9%	6.8%	4.3%
Residual MLP d512	50.1%	21.9%	11.5%	7.8%	5.1%

Main decoded-belief mass

Interpretation

Full prefix-belief supervision is not enough to make a fixed-width dense recurrent state behave like an exact symbolic belief executor at modulus 31. The student learns useful approximate state updates, and the approximation improves with more width, but exact pair-belief mass remains low.

The gap between query mass and belief mass is important. A model can assign useful probability to the correct projected query support while still failing to represent the full pair belief. The best main model reaches 52.1% projected query mass at train-length 8, but only 21.9% exact pair-belief mass. That means it learned a partially useful compressed computation, not exact latent belief execution.

The frozen probes do not reveal a hidden exact state. Probe belief mass is close to decoder belief mass on the main fixed-length evaluation rows. This suggests the primary bottleneck is the learned recurrent state and transition, not merely the trained belief decoder.

The residual transition is the best tested transition, but its advantage is modest. Its larger effect is not qualitative: it improves the curve without changing the conclusion.

Conclusion

Dense teacher distillation improves recurrent belief execution, and state capacity matters. The strongest tested model is the residual MLP decoder with 512 hidden dimensions. It is better than the GRU at the same width and much better than the 256-dimensional model.

The experiment does not support the claim that full teacher supervision alone makes a dense recurrent student learn exact modular belief-state execution. On the hard modulus-31 setting, the dense state remains an approximation whose accuracy decays quickly with program length.

The most direct next test is to add structure back into the state or transition: factorized register beliefs, explicit pair-state tables, sparse support tracking, or a differentiable update rule constrained to preserve and transform belief mass. The dense-state route is viable as an approximation, but the results show a clear exact-execution bottleneck.

Reproducibility

Code and lightweight artifacts are in:

experiments/dense_teacher_distillation/

Checkpoints are stored separately in:

large_artifacts/dense_teacher_distillation/checkpoints/

The normal research bundle is the experiment directory. The checkpoint directory is only needed when loading saved model weights.

Key files:

src/dense_teacher_distillation_experiment.py
src/analyze_dense_teacher_distillation.py
analysis/summary.md
analysis/threshold_summary.csv
analysis/final_metrics_query_mean.csv
checkpoint_manifest.csv

Experiment log 10

Show the running log (10 entries, 2026-06-21)

Objective

Test whether a recurrent executor with only a fixed-width dense hidden state can represent exact belief-state computation when every prefix state is supervised by an exact teacher distribution.

The task uses modular programs over two registers (A,B). Programs combine arithmetic updates with observation filters. The teacher computes the exact belief distribution over all (A,B) pairs after each prefix. The student receives the same symbolic program and initial relation, keeps a dense hidden vector, and is trained to decode the teacher belief at each recurrent step.

Primary Questions

Does full prefix-belief distillation make the dense recurrent executor solve the task, or does it still fail despite high-bandwidth supervision?
Does increasing dense state capacity improve exact belief decoding and decoder-projected query accuracy?
Does a residual transition cell improve over a GRU cell when the target is exact belief-state execution?
Does a low-rank decoder expose a decoder bottleneck by improving or degrading projected query accuracy relative to an MLP decoder?

Metrics

decoder_belief_target_mass: probability assigned by the student belief decoder to the exact target support.
decoder_query_target_mass: query probability obtained by projecting the decoded pair distribution and measuring mass on the exact query support.
probe_belief_target_mass: probability assigned by a separately trained frozen-state probe to the exact target support.
query_target_mass: direct query-head mass, treated as diagnostic when the query-head loss is disabled.

The headline metric is decoder_query_target_mass, with decoder_belief_target_mass as the stricter state-execution metric.

Artifact Layout

Code and lightweight outputs: experiments/dense_teacher_distillation/
Checkpoints: large_artifacts/dense_teacher_distillation/checkpoints/
Run outputs: experiments/dense_teacher_distillation/runs/<variant>/
Analysis outputs: experiments/dense_teacher_distillation/analysis/

Planned Sequence

Smoke test on modulus 7 with tiny capacity and short chains to validate the harness.
Pilot on modulus 11 across state sizes and transition/decoder variants.
Main bottleneck sweep on modulus 31 with longer held-out chains.
Generate tables, figures, checkpoint manifest, standalone report, and HTML report.

Variant Plan

smoke_gru_mlp_d32: minimal GRU with MLP belief decoder.
smoke_residual_mlp_d32: minimal residual transition with MLP belief decoder.
smoke_gru_lowrank_d32_r4: minimal GRU with low-rank pair decoder.
pilot_gru_mlp_d128: baseline dense state at moderate capacity.
pilot_gru_mlp_d256: capacity scaling check.
pilot_residual_mlp_d256: transition structure check.
pilot_gru_lowrank_d256_r16: decoder structure check.
main_gru_mlp_d256: main baseline.
main_gru_mlp_d512: main capacity scaling check.
main_residual_mlp_d512: main transition check.

2026-06-21 Setup

Created the standalone experiment directory:

experiments/dense_teacher_distillation/src/
experiments/dense_teacher_distillation/reports/
experiments/dense_teacher_distillation/runs/
experiments/dense_teacher_distillation/analysis/figures/
large_artifacts/dense_teacher_distillation/checkpoints/

Implemented the first harness changes:

Teacher-belief distillation is the default and only supervision mode.
Added decoder-projected query metrics from decoded pair distributions.
Added variant_name, transition, decoder_type, and decoder_rank controls.
Added GRU and residual transition choices.
Added MLP and low-rank belief decoder choices.
Checkpoints default to the external large_artifacts tree.

Next action: run smoke tests and fix any harness issues before starting the pilot sweep.

2026-06-21 Smoke Tests

Ran three two-step CUDA smoke tests at modulus 7:

Variant	Status	Notes
`smoke_gru_mlp_d32`	passed	GRU transition and MLP decoder path writes full metrics and checkpoint.
`smoke_residual_mlp_d32`	passed	Residual transition path writes full metrics and checkpoint.
`smoke_gru_lowrank_d32_r4`	passed	Low-rank decoder path writes full metrics and checkpoint.

Smoke artifacts:

Run directories: experiments/dense_teacher_distillation/runs/smoke_*
Checkpoints: large_artifacts/dense_teacher_distillation/checkpoints/smoke_*

The metric CSVs include variant metadata, decoder belief metrics, probe belief metrics, and decoder-projected query metrics. The external checkpoint layout is working.

Next action: run a modulus-11 pilot sweep to decide which architecture and capacity settings deserve a larger modulus-31 run.

2026-06-21 Pilot Sweep

Ran four modulus-11 pilot variants with train lengths up to 6 and evaluation lengths 3, 6, 9, and 12:

Variant	L=3 decoder-query	L=6 decoder-query	L=9 decoder-query	L=12 decoder-query	Decision
`pilot_gru_mlp_d128`	90.5%	61.9%	44.8%	35.2%	Keep as capacity baseline only.
`pilot_gru_mlp_d256`	94.7%	71.3%	53.1%	41.4%	Carry forward as main baseline.
`pilot_residual_mlp_d256`	94.8%	71.6%	54.7%	42.6%	Carry forward at larger width.
`pilot_gru_lowrank_d256_r16`	78.8%	50.3%	36.4%	28.5%	Drop from main sweep.

Strict decoded-belief mass showed the same ordering. The low-rank decoder also produced weaker frozen-probe results, so the constrained decoder appears to alter the learned state rather than merely reducing the final readout.

Analysis artifacts generated:

analysis/all_metrics_long.csv
analysis/final_metrics_query_mean.csv
analysis/threshold_summary.csv
analysis/summary.md
analysis/figures/

Main sweep decision:

Use MLP belief decoders.
Test state capacity at modulus 31.
Include a residual-transition run because it was slightly better on longer held-out chains.
Evaluate lengths beyond the training range so the key outcome is length transfer, not only short-chain fit.

2026-06-21 Main Sweep

Ran three modulus-31 main variants with train lengths up to 8 and evaluation lengths 4, 8, 12, 16, and 24:

Variant	L=4 decoder-query	L=8 decoder-query	L=12 decoder-query	L=16 decoder-query	L=24 decoder-query
`main_gru_mlp_d256`	69.9%	36.1%	23.7%	16.8%	11.7%
`main_gru_mlp_d512`	78.1%	49.2%	35.4%	26.3%	18.7%
`main_residual_mlp_d512`	80.5%	52.1%	36.5%	27.9%	20.5%

Strict decoded-belief mass at K >= L:

Variant	L=4 belief	L=8 belief	L=12 belief	L=16 belief	L=24 belief
`main_gru_mlp_d256`	32.6%	9.0%	4.0%	2.5%	1.6%
`main_gru_mlp_d512`	45.7%	19.3%	9.9%	6.8%	4.3%
`main_residual_mlp_d512`	50.1%	21.9%	11.5%	7.8%	5.1%

Main interpretation:

Width matters. Moving from d256 to d512 substantially improved both decoder-projected query mass and strict decoded-belief mass.
The residual transition was consistently, but modestly, better than the GRU at the same d512 width.
Full teacher-belief supervision did not make the dense state exact on modulus 31. The best run reached strong query support mass for short programs but only 21.9% strict belief mass at train-length 8 and 5.1% at held-out length 24.
Frozen probes were close to the trained decoder on fixed eval rows, so the bottleneck is mostly the learned dense state and transition, not only the final decoder.

Next action: write the standalone report and HTML artifact, then verify the directory layout and references.