diff --git a/experiments/mess3_token_guess_cycle_2/NOTES.md b/experiments/mess3_token_guess_cycle_2/NOTES.md new file mode 100644 index 00000000..6f792656 --- /dev/null +++ b/experiments/mess3_token_guess_cycle_2/NOTES.md @@ -0,0 +1,78 @@ +# MESS3 token-guess cycle 2 + +Cycle 2 is a planned re-run of the token-guess comparison ahead of presenting it. +It does not yet contain any training conditions. + +What is here now is the measurement work that has to precede the re-run: + +- `REVIEW.md` — an audit of `mess3_token_guess_cycle_1` and + `mess_3_kelly_cycle_1/2/3`, which are one experiment split over four + directories. +- `PLAN.md` — a proposed design for the re-run, with the open questions that + still need deciding. +- `metric_references.py` and `references/` — the floors and ceilings both + headline metrics are measured against. Cycle 1 reported bare numbers, and both + metrics turn out to have task-imposed ranges narrow enough that the bare + numbers mislead. +- `operating_point.py` and `task_parameters/` — how the metric's range and + precision vary with the MESS3 transition and emission parameters, so the + process can be chosen to make the metric sensitive rather than degenerate. +- `statistics.py` and `audit/` — seed-level aggregation with intervals, paired + comparison, multiplicity correction, and power planning, plus a re-analysis of + the results already committed. + +## What the references say + +| reference | belief-probe R² | greedy accuracy | +|---|---:|---:| +| randomly initialised transformer | 0.8733 | 0.3412 | +| affine probe on the raw observations | 0.9668 | — | +| exact Bayesian filter, one observation | — | 0.6732 | +| exact Bayesian filter, saturated | — | 0.6883 | +| supervised next-token replication | 0.9989 | 0.6859 | + +Belief-probe R² therefore moves through 0.032 and greedy accuracy through 0.015. +Cycle 1's reward-only and max-entropy arms land below the untrained network on +the first metric, and at the repeat-the-previous-token rule on the second. + +## What the audit says + +Re-reading the three-seed Kelly cycles with a t-based interval on the mean and +pairing across shared seeds, none of the 34 published pairwise orderings survives +a Holm correction. The reported `±` values were population standard deviations of +three samples, roughly a third of the width of a confidence interval. + +## What the operating-point analysis says + +Cycle 1's parameters are close to the worst choice in the symmetric MESS3 family +for metric sensitivity. Slowing the chain to a self-transition of 0.995 at +`alpha=0.70` roughly triples the belief-probe range and multiplies the accuracy +range by ten, at the cost of a probe interval that widens from ±0.0006 to +±0.0046 because the rollout decorrelates far more slowly. + +It also shows that belief-probe R² declines with continued *supervised* training +while cross-entropy stays at the Bayes floor. The decline cycle 1 saw over 20M +PPO steps is therefore not a property of reinforcement learning, and optimiser +and learning rate move the headline metric more than most of the arms do. + +Two operational consequences: + +- `CausalTransformerEncoder` bands attention at every layer, so `context_len=64` + over three layers reaches 192 observations while the belief needs about 32. + Dropping to `context_len=16` costs nothing measurable and makes a learner step + 3.7x faster, a cached rollout step 3.4x faster, and a whole PPO run 1.6x + faster. +- A 30,000-step probe of the slower chain has a measurement error of ±0.0025, + the same size as the smallest seed-to-seed spreads. Raise the probe to about + 500,000 test steps, where it falls to ±0.0005. + +Regenerate all three with: + +```bash +uv run rl-harness experiments.mess3_token_guess_cycle_2.references.experiment +uv run rl-harness experiments.mess3_token_guess_cycle_2.task_parameters.experiment +uv run rl-harness experiments.mess3_token_guess_cycle_2.audit.experiment +``` + +Pass `--no-upload-artifacts` when `B2_*` variables are configured as secrets, or +the bucket name is written into tracked results. diff --git a/experiments/mess3_token_guess_cycle_2/PLAN.md b/experiments/mess3_token_guess_cycle_2/PLAN.md new file mode 100644 index 00000000..cb46fca8 --- /dev/null +++ b/experiments/mess3_token_guess_cycle_2/PLAN.md @@ -0,0 +1,328 @@ +# Proposed design for token-guess cycle 2 + +Draft for discussion. `REVIEW.md` establishes the problems; this is one way to +address them. Everything here is negotiable, and the open questions at the end +are the parts I would not decide alone. + +## Principle + +Cycle 1 ran a large number of arms, measured each once, and read a ranking off +the resulting table. Cycle 2 should run fewer arms, measure each well enough to +put an interval on it, and pre-commit to which comparisons it is trying to +resolve. The compute cost of doing this is trivial — the numbers are at the end — +so the binding constraint is design discipline, not GPU budget. + +## 1. Consolidate into one study + +One directory, one `shared.py`, one task class, one probe. Today the same +experiment is spread over four directories with two copies of the task, three +copies of the probe, and three IQN implementations (`mess3_token_guess_cycle_1. +iqn_value.iqn`, `mess3_reward_state_cycle_1.iqn`, and the library's +`IQNValueMixin`). + +The library's `IQNValueMixin` should be the only one. Anything else that is a +reusable RL concept goes to `rl-harness`; Kelly wagering stays here. + +## 1b. Move the operating point, and pay for it in probe size + +Cycle 1 ran at `alpha=0.85` with a self-transition of 0.9, which is close to the +worst choice in this family for metric sensitivity: a fast-mixing chain lets the +belief be approximated by an exponentially weighted average of recent one-hot +observations, which is exactly what an affine probe computes. Slowing the chain +breaks that approximation. + +| point | α | p | R² floor | R² range | accuracy range | probe ±95% at 30k steps | +|---|---:|---:|---:|---:|---:|---:| +| cycle 1 | 0.85 | 0.900 | 0.9671 | 0.033 | 0.015 | 0.0006 | +| candidate A | 0.70 | 0.990 | 0.8967 | 0.103 | 0.134 | 0.0037 | +| candidate B | 0.60 | 0.995 | 0.8782 | 0.122 | 0.136 | 0.0061 | +| candidate C | 0.70 | 0.995 | 0.9061 | 0.094 | 0.145 | 0.0046 | + +Candidate C is my recommendation: about three times the belief-probe range and +ten times the accuracy range, while keeping the task clearly learnable at a +Bayes-optimal accuracy of 0.68. + +The trade is precision, and it is worth being concrete about why. + +The probe number is an estimate, computed from a finite rollout. Two rollouts of +the same checkpoint give two slightly different answers, and the spread between +them is the metric's measurement error. That error depends on how many +*independent* samples the rollout holds, which is not the same as how many steps +it has. On a slow chain the hidden state sits still for hundreds of steps, so +consecutive samples are near-duplicates and a long rollout can carry very little +new information. + +Measured by direct replication — running the identical probe 60 times on fresh +rollouts and taking the standard deviation of the answers: + +| operating point | 30k steps | 120k steps | 500k steps | +|---|---:|---:|---:| +| cycle 1 | 0.0004 | 0.0002 | 0.0001 | +| candidate C | 0.0025 | 0.0012 | 0.0005 | + +At 30,000 steps the slow chain's measurement error is 0.0025, which is the same +size as the smallest seed-to-seed spreads observed in the Kelly cycles. Probe +noise would then be indistinguishable from seed variance, and adding seeds would +not fix it. + +**Raise the probe to around 500,000 test steps**, where the measurement error +falls to 0.0005 and is safely below anything the design needs to resolve. Probe +rollouts are inference only and cost seconds, so this is the cheapest fix in the +plan and the easiest one to omit by accident. + +Two notes on method. The i.i.d. bootstrap used in `REVIEW.md` understates this +error, because resampling individual steps pretends they are independent; a +moving-block bootstrap predicted ±0.0046 at 30k on candidate C and replication +measured ±0.0047. And the simulation behind these numbers does not include the +512-step episode resets the real environment applies, which partially +decorrelate the rollout, so treat 0.0025 as an upper bound and confirm it once +against the real probe by running it repeatedly on a single fixed checkpoint +with different probe seeds. + +Two consequences follow. Episode length should rise from 512 so that the reset +back to a uniform belief is a smaller fraction of each trajectory. And at +γ = 0.99 the effective horizon of about 100 steps is now shorter than the state's +persistence, which is worth stating when γ is discussed. + +## 2. Report metrics against their range, always + +Never present a bare R². Every table gets the four reference rows from +`references.experiment` alongside the trained arms: + +| row | belief R² | source | +|---|---:|---| +| affine probe on raw observations | 0.9668 | analytic, no network | +| randomly initialised transformer | 0.8733 | same probe, same architecture | +| supervised next-token model | 0.9989 | `mess3_supervised` | +| exact Bayesian filter (accuracy) | 0.6883 | analytic | + +R² is a bad scale here because the interesting region is 0.967–0.999. I would +report three things instead, and let the audience pick: + +- **Probe MSE**, which has real dynamic range (0.0012 at R² = 0.99 against 0.03 + at R² = 0.85) and reads naturally on a log axis. +- **Incremental R²**: fit the probe on the raw-observation window alone, then on + the window plus the network's activations, and report the increment. This + answers "what does the network represent that its own inputs did not already + linearly provide", which is the question the study is actually asking. +- **Position in the usable range**, as `references.experiment` already computes, + for the one summary slide. + +## 3. Treat training duration as an axis, not a hyperparameter + +Log-spaced checkpoints for every arm, following `rl-harness/docs/ +checkpoint_strategy.md`, with the probe run at each. + +Duration should not be tuned, because belief-probe R² is not monotonic in it and +the decline is now known to be optimiser drift rather than anything about the +objective. Picking the step budget after seeing the results would be choosing the +ranking. So: pre-register the budget, pre-register whether the primary statistic +is the final checkpoint or the mean over the last k, and report the whole curve +either way. I prefer the mean over the last k, which removes the ±0.0035 +within-run fluctuation for free. + +The step budget must be identical across every arm. 2.5M is defensible; so is +5M. What is not defensible is the current mixture of 828K, 2.5M, 3M and 20M. + +## 4. Seeds: ten, from a documented spawn + +Ten seeds per condition, paired across arms, drawn by spawning from one master +`SeedSequence` and recorded in the manifest. + +The planning table below is from `statistics.seeds_for_power`, paired, 80% power, +95% confidence: + +| belief-R² gap | sd = 0.002 | sd = 0.005 | sd = 0.010 | sd = 0.020 | +|---:|---:|---:|---:|---:| +| 0.002 | 10 | 52 | 199 | 787 | +| 0.005 | 4 | 10 | 34 | 128 | +| 0.010 | 3 | 5 | 10 | 34 | +| 0.020 | 3 | 3 | 5 | 10 | +| 0.100 | 3 | 3 | 3 | 3 | + +Ten seeds resolves any gap of 0.01 or more at the spreads actually observed +(0.001–0.011). It will not resolve the 0.002–0.005 gaps between the middle arms, +and that is the point: those should be reported as ties rather than chased. Going +to the ~150 seeds that would settle them is possible on this compute budget but I +do not think it buys a better talk. + +Note the planning numbers are computed from effects observed at n=3 and are +therefore optimistic — the observed gap in a small sample is biased upward. Ten +is already padded for that; I would not go below eight. + +## 5. Sweep the optimiser and the coefficients, in two stages + +Only four things are worth sweeping, and the first one matters more than the +arms do. See `task_parameters/` for the evidence. + +| priority | swept | values | +|---|---|---| +| 1 | optimiser | AdamW, Muon | +| 1 | learning rate | 1e-4, 3e-4, 1e-3 | +| 2 | predictive aux weight λ | 0.03, 0.1, 0.3, 1.0 | +| 2 | max-entropy reward coefficient α | 0.01, 0.05, 0.2, 0.5 | +| 2 | IQN loss coefficient | 0.125, 0.5, 2.0 | +| 2 | Kelly direct-loss weight | 0.25, 1.0, 4.0 | + +**Optimiser and learning rate move the headline metric more than the arms do.** +Training the study architecture on next-token prediction and probing it as it +goes, belief-probe R² *falls* with continued training while cross-entropy sits at +the Bayes floor — 0.9567 to 0.9318 over 6,000 AdamW steps at the cycle-1 +parameters, and 0.9487 to 0.9367 at the slower candidate. That is supervised +training, so the decline over 20M PPO steps in cycle 1 is not an artefact of +reinforcement learning. `mess3_supervised` independently found SGD at 0.9979 +against Muon at 0.9843 with identical next-token accuracy, a 0.014 swing from the +optimiser alone that is larger than most of the between-arm gaps cycle 2 is +trying to resolve. An arm comparison run at one fixed optimiser and learning rate +is partly measuring optimiser drift. + +Sweeping the learning rate per arm is separately necessary because arms differ in +parameter count and in how much auxiliary gradient enters the shared trunk, so an +arm can lose by being mis-tuned rather than by being a worse idea. + +**Do not sweep:** context length (see below — shrink it once and fix it), GAE λ +(irrelevant at γ = 0, low leverage at γ = 0.99), batch size, minibatch size, epoch +count, clip parameter, `d_model`, or `n_layers`. Fix these and record them. γ is a +scientific factor, not a nuisance hyperparameter; keep it in the design rather +than tuning it away. + +## 5b. Shrink the context length once, then leave it alone + +`CausalTransformerEncoder` applies its causal band at *every layer*, so the +receptive field is `n_layers × context_len`. Confirmed by perturbing inputs at +increasing distance and watching when the output stops changing: three layers at +`context_len=64` reach exactly 192 observations. The belief needs about 32. + +The cost is paid twice, because the learner recomputes over +`lookback + max_seq_len = 192 + 32 = 224` positions to produce 32 useful +embeddings, and the cached rollout path attends over the same band. + +| context_len | receptive field | learner step | rollout step | belief R² available | +|---:|---:|---:|---:|---:| +| 64 | 192 | 305 ms | 15.9 ms | 1.00000 | +| 32 | 96 | 146 ms (2.1x) | 6.9 ms (2.3x) | 1.00000 | +| 16 | 48 | 82 ms (3.7x) | 4.7 ms (3.4x) | 1.00000 | +| 12 | 36 | 70 ms (4.3x) | 4.6 ms (3.5x) | 0.99997 | +| 8 | 24 | 57 ms (5.4x) | 3.9 ms (4.0x) | 0.99966 | + +Belief R² is at candidate C; candidate B needs more, reaching 0.99983 only at a +receptive field of 48. + +**Use `context_len=16`.** It costs nothing measurable in available belief R² at +any candidate operating point, leaves a 50% margin over what the belief needs, +and keeps the choice valid if α or p move later. `context_len=12` is defensible +if candidate C is locked in, but the extra saving is small. + +`context_len` affects compute only, not capacity — positions are encoded with +RoPE, so the parameter count does not change. + +End-to-end this is worth **1.6x**, measured with a real PPO loop at matched +environment steps. It is less than the 3.4–3.7x component figures because +environment stepping is pure NumPy and unaffected. The split differs on a 4090 +with 16 env runners, so measure it once there before re-budgeting. + +**Two stages, with disjoint seeds.** Stage 1 sweeps on seeds 0–2 and selects one +configuration per arm. Stage 2 re-runs the selected configurations on ten +*fresh* seeds, and only stage 2 enters the results table. Selecting and reporting +on the same seeds is what makes a swept comparison optimistic, and it is easy to +avoid here. + +Selection criterion for stage 1 must be pre-registered. I suggest belief-probe +MSE, since accuracy is saturated. + +## 6. Pre-register the comparison family + +Ranking eight arms is 28 tests. Declare the primary comparisons in advance and +correct within that family; report everything else as exploratory. + +Proposed primaries, all on belief-probe MSE at γ = 0.99: + +1. reward-only PPO against the raw-observation floor — does task reward degrade + the representation? +2. reward-only PPO against an untrained network — is the degradation relative to + initialisation? +3. predictive aux against reward-only — does the auxiliary objective recover it? +4. IQN against reward-only — does distributional value do the same? +5. Kelly against IQN — are they doing the same thing or different things? +6. best combined arm against the best single arm — do they compose? + +Six tests, Holm-corrected. Everything else is exploratory and labelled as such. + +## 7. Arms + +Two axes crossed, with max-entropy as an orthogonal third: + +| axis | levels | +|---|---| +| critic | scalar mean, IQN | +| representation pressure | none, predictive aux, Kelly wager | +| entropy in reward | off, on | +| γ | 0, 0.99 | + +The full cross is 24. I would run the 2 × 3 × 2 = 12 critic × pressure × γ cells +as the main grid, and treat entropy as a separate 2 × 2 (critic × γ) add-on at +its swept coefficient, giving 16 conditions. The γ = 1.0 and differential +average-reward objectives from cycle 1 are worth keeping as a third γ level if +the budget is there, since the average-reward arm was the one that restored sane +critic diagnostics. + +## 8. Compute + +Measured from the committed manifests: a 2.5M-step arm costs 6–14 minutes on an +RTX 4090, median about 6.5. Adding eight checkpoint probes puts it near 12. + +At `context_len=16` a run costs about 1.6x less than the measured 6–14 minutes, +so budget roughly 7 minutes per run including checkpoint probes. + +| stage | runs | GPU-hours | cost at $0.35/hr | +|---|---:|---:|---:| +| stage 1 sweep (7 arms × 6 optimiser/lr × 3 coefficient × 3 seeds) | 378 | 47 | ~$17 | +| stage 2 confirmation (16 conditions × 10 seeds) | 160 | 20 | ~$7 | +| references, operating point, audit | — | <1 | — | +| **total** | **538** | **67** | **~$24** | + +Roughly nine wall-clock hours across eight Vast boxes. The entire committed +history of these four studies is 19.4 GPU-hours, so this is about three times all +prior work on the question, for the price of lunch. + +`devops.vast.provision up -n N --run "..."` gives one command per box, so the +sweep matrix needs a small dispatcher that maps (arm, coefficient, seed) to +`--run` strings. That is the only new infrastructure required; there is no +multi-seed or sweep support in the harness today. + +## Open questions + +1. **Is the degradation result the headline?** The strongest thing in the data is + that reward-only PPO ends up below both the raw observations and its own + initialisation. That reframes the study from "which objective induces belief + geometry" to "task reward destroys it and these objectives protect it". It is + a better talk, but it is a different talk. Which one do you want to give? + +2. **Is belief-probe R² still the right dependent variable?** It was chosen to + match the supervised replication, but on this task it is compressed into the + top 3% of its range and it measures linear decodability rather than use. An + intervention experiment — ablate the probe-identified directions and see + whether actions change — would answer the question colleagues will actually + ask. That is more work than a re-run. Worth it before the presentation, or + after? + +3. **How much of the Kelly programme survives?** Cycles 1–3 ran fourteen Kelly + variants and none of the orderings among them is supported. Cycle 2 could + carry one Kelly arm (conditional decoupled, the best performer) rather than + the family. Is the wager-calibration result something you want to present in + its own right, or is Kelly here purely as a representation-shaping device? + +4. **Do you want γ as a swept axis or a fixed choice?** γ is currently confounded + with the study directory, and the γ = 0 versus γ = 0.99 contrast is one of the + larger apparent effects. Making it a proper factor doubles the grid. My + inclination is yes, because a clean γ effect is presentable and the current + evidence for it is cross-study. + +5. **Do we keep the 20M-step observation?** The R² decay from 0.990 to 0.968 over + 20M steps is arguably the most interesting single curve in the repository, and + it exists for one arm at one seed. Three seeds of a long run for two arms + would cost about six GPU-hours and would turn an anecdote into a result. + +6. **What is the presentation format?** If there is a slide budget, that should + drive how many arms survive. Sixteen conditions with intervals is a dense + table; six conditions with intervals and a reference band is a slide. diff --git a/experiments/mess3_token_guess_cycle_2/REVIEW.md b/experiments/mess3_token_guess_cycle_2/REVIEW.md new file mode 100644 index 00000000..6c31b65d --- /dev/null +++ b/experiments/mess3_token_guess_cycle_2/REVIEW.md @@ -0,0 +1,242 @@ +# Review of the MESS3 token-guess studies + +Scope: `mess3_token_guess_cycle_1` and `mess_3_kelly_cycle_1/2/3`. These are four +directories but one experiment. All of them train a 3-layer, 96-wide transformer +with PPO on passive MESS3 (`alpha=0.85`, `delay=1`, episode length 512), reward +it 1 for guessing the next emitted token, and then fit a held-out rank-2 affine +probe from its residual stream to the exact Bayesian belief. `NextTokenGuessTask` +and `RawNextTokenTask` are the same task written twice. + +The findings below are ordered by how much they would change what you tell your +colleagues. Every number is reproducible from the repository: + +```bash +uv run rl-harness experiments.mess3_token_guess_cycle_2.references.experiment +uv run rl-harness experiments.mess3_token_guess_cycle_2.audit.experiment +``` + +--- + +## 1. The headline metric is mostly measuring something trivial + +Belief-probe R² was reported as a bare number from 0.85 to 0.99, and read as +"how much of the Bayesian belief did this objective induce". It cannot be read +that way, because the metric has a floor that nobody measured. + +An affine probe fitted directly to the **one-hot encoded raw observations**, with +no network and no training anywhere in the pipeline, scores **R² = 0.9668**. It +gets there on four observations and is saturated by six. + +| probe input | R² | +|---|---:| +| last 1 observation | 0.8043 | +| last 2 observations | 0.9302 | +| last 4 observations | 0.9648 | +| last 8 observations | 0.9669 | +| last 64 observations | 0.9668 | + +A **randomly initialised** copy of the study's own transformer, probed by exactly +the same code at exactly the same site, scores **R² = 0.8733**. + +The supervised next-token replication in `mess3_supervised` reaches **0.99888**. + +So the range this metric can move through is 0.9668 to 0.9989 — a span of +**0.032** — and the published scores land like this: + +| condition | reported R² | position in the usable range | +|---|---:|---:| +| `comparison/reward_only` | 0.8552 | −348% | +| `comparison/max_entropy` | 0.8558 | −346% | +| *randomly initialised transformer* | *0.8733* | *−293%* | +| `comparison/predictive_loss` | 0.9319 | −109% | +| *affine probe on raw observations* | *0.9668* | *0%* | +| `iqn_value` | 0.9760 | +29% | +| `kelly_cycle_3/conditional_decoupled_kelly_iqn` | 0.9824 | +49% | +| `kelly_cycle_2/correctness_iqn` | 0.9857 | +59% | +| *supervised next-token replication* | *0.9989* | *100%* | + +Two consequences. + +**Reward-only PPO and max-entropy PPO are below the untrained network.** After +2.5M steps their residual streams are *less* linearly informative about the +belief than they were at initialisation. That is a real and interesting result — +task-reward training is destroying linearly-available structure — but it is the +opposite of the claim "reward-only PPO reaches R² = 0.855 of the belief". + +**All eight Kelly cycle-2 wager arms sit below the no-network floor.** The +cycle-2 note reads their 0.94–0.96 as respectable. Against the floor they are +negative. + +This is safe because the task is passive: `NextTokenGuessTask.resolve_action` +returns the same transition matrix for every action, so the distribution of +beliefs the probe sees is identical for every arm, for the untrained network, and +for the raw-observation baseline. The comparison is exact, not approximate. + +## 2. Greedy token accuracy is saturated and cannot discriminate + +An exact Bayesian filter allowed only the previous observation — the trivial +"repeat the last token" rule — scores **0.6732**. Given eight or more +observations it saturates at **0.6883**. Total usable range: **0.0151**. + +Reported accuracies span 0.6733 to 0.6860. So the best arms are already within +0.002 of Bayes-optimal, `reward_only` at 0.6733 is at the repeat-last-token +rule, and the whole comparison lives inside 1.5 percentage points. Meanwhile the +probe's own sampling error at 30,000 test steps is ±0.005 at 95%, and that +assumes independent samples, which consecutive steps of a slowly-mixing chain +are not. + +Reporting this metric to four decimal places invites exactly the question you do +not want from the audience. + +## 3. None of the published orderings is statistically supported + +Cycles 2 and 3 ran three seeds and reported `mean ± std`. That `±` is a +population standard deviation (`ddof=0`), not an interval; for three samples it +sits 18% below the sample standard deviation and about 3x below a 95% confidence +interval on the mean. + +Recomputed properly, and paired across the seeds each pair of arms shares: + +- **Kelly cycle 3**: 0 of 6 pairwise orderings survive a Holm correction. +- **Kelly cycle 2**: 0 of 28 pairwise orderings survive a Holm correction. + +The headline gap is the clearest case. Cycle 3 reported PPO `0.8742 ± 0.0432` +against IQN `0.9688 ± 0.0090`, which reads as decisive. Paired on the three +shared seeds, the difference is +0.095 with a 95% interval of **[−0.017, ++0.206]** and p = 0.067. PPO's own interval, [0.743, 1.006], extends past the +maximum value R² can take. + +Some claims in the notes are far past what three seeds can carry. Cycle 3 states +that "conditional Kelly with a scalar critic narrowly exceeds plain IQN in belief +R²". That gap is 0.0029 against a paired spread of 0.013; resolving it needs +**165 seeds**. + +## 4. Final-checkpoint scores are drawn from a moving target + +Every arm reports the checkpoint at its step budget, with `num_to_keep=1`. The +one arm that kept a checkpoint curve, `iqn_value_20m`, shows why that is unsafe: + +| steps | R² | greedy accuracy | +|---:|---:|---:| +| 0.83M | 0.9902 | 68.12% | +| 2.48M | 0.9772 | 68.08% | +| 3.31M | 0.9673 | 68.34% | +| 10.76M | 0.9718 | 68.77% | +| 20.03M | 0.9686 | 68.47% | + +Belief R² **peaks at the first retained checkpoint and declines**, while accuracy +slowly rises. They are not merely uncorrelated, they move in opposite directions. +After 3M steps R² fluctuates with a standard deviation of 0.0035 and a range of +0.013 — comparable to most of the between-arm gaps being reported. + +So "IQN scores 0.976 at 2.5M" is one draw from a decaying, noisy trajectory. A +different step budget reorders the table. And the budget is not even constant +across the family: 828K, 2.5M, 3M and 20M all appear. + +## 5. Every intervention was tested at exactly one strength + +| intervention | values tried | +|---|---| +| predictive auxiliary loss λ | 0.1 | +| max-entropy reward coefficient α | 0.05 | +| IQN loss coefficient | 0.5 | +| Kelly direct-loss weight | 1.0 | +| learning rate | 3e-4 everywhere (4.2e-4 in the reward-state battery) | + +No conclusion of the form "X does not help" is available from one point. The +max-entropy arm is the clearest example: at α = 0.05 the entropy bonus is at most +`0.05 · ln 3 ≈ 0.055` against a reward of 0 or 1, so the arm was never given +enough signal to do anything, and "max-entropy training produced essentially the +same belief representation as reward-only" is a statement about that one +coefficient. + +The arms are also not matched on anything except environment steps. IQN arms set +`vf_loss_coeff=0` and add a quantile loss; Kelly and predictive arms add heads +and extra gradient into the shared trunk. Parameter count and total gradient +magnitude differ between arms while the learning rate is held fixed, so an arm +can win by being better tuned rather than by being a better idea. + +## 6. γ is confounded with the study + +Discount factor and condition set were varied together across directories. + +| γ, λ | where | +|---|---| +| 0.99, 0.95 | token-guess `comparison`, `iqn_value`, Kelly cycle 3 | +| 1.0, 0.95 | Kelly cycle 1, `iqn_gamma_1_3m` | +| 0.0, 0.0 | Kelly cycle 2 | +| 0.0, 0.95 | `iqn_return_objectives/gamma_zero` | + +Cycle 2 concludes that "moving from the prior γ = 0.99 reward-only result +(R² = 0.8552) to γ = 0 gives R² = 0.9834 without an auxiliary objective". That +compares n=1 to n=3, across two studies, two probe seed streams, and two +code paths. It is a plausible hypothesis, not a measured effect. + +The λ change is harmless in itself — GAE weights terms by `(γλ)^k`, so at γ = 0 +the value of λ cannot matter — but it should be made explicit rather than left as +an apparent second difference. + +## 7. Smaller items worth fixing + +**Probe protocol is not shared.** The RL probe uses a rank-2 constrained affine +map on 30,000 test steps at the post-final-LayerNorm site, from a single +train/test split. The supervised reference uses a full-rank probe on 472,390 +positions, sweeps five layer sites and reports the best. Those two numbers should +not be placed in the same table until they are measured the same way. Kelly cycle +2 also uses different probe seed streams (200/201/202) from the token-guess cycle +(100/101/102), so cross-study numbers carry an extra source of variation. + +**Probe R² and MSE are reported, but no uncertainty.** Bootstrapping the probe's +test set puts its own sampling noise at ±0.0005, which is small — but it needed +to be measured rather than assumed. + +**Pairing does not currently buy anything.** All arms already run seeds 42/43/44, +so paired analysis is free. On the existing data the paired spread is frequently +*larger* than the unpaired spread, meaning seed effects are arm-specific +optimisation noise rather than shared data noise. Three seeds cannot settle this; +keep the paired design because it costs nothing, but do not budget for the +variance reduction. + +**Actual step counts differ between arms.** 2,515,821 against 2,512,179 within +Kelly cycle 1, because the stop condition triggers on batch boundaries. Small, +but it is free to record and equalise. + +**Run-to-run non-determinism is negligible.** The two reproduction arms re-ran +the same seed and code and differed by 0.00016 in R². Seed variance is roughly +200x larger. Determinism is not the problem here; sample size is. + +**Naming and layout.** `mess3_token_guess_cycle_1` and `mess_3_kelly_cycle_*` +differ in spelling, and one experiment is spread over four directories with three +copies of the probe and two copies of the task. `mess_3_kelly_cycle_3/shared.py` +imports IQN from `mess3_token_guess_cycle_1.iqn_value.iqn` while the reward-state +battery imports it from `mess3_reward_state_cycle_1.iqn`, and the library now has +its own `IQNValueMixin`. Three IQN implementations are in play. + +**Committed manifests leak the B2 bucket name.** Any run made with `B2_*` +configured writes the bucket and endpoint into tracked `results/`. Pass +`--no-upload-artifacts`, or scrub those fields, when running from an environment +where those are secrets. + +--- + +## What this means for the presentation + +The existing work supports a small number of claims very well: + +- Training reward-only PPO on this task **degrades** the linear decodability of + the Bayesian belief relative to both the raw observations and to the network's + own initialisation. That is a genuinely surprising, presentable result, and it + is large enough to survive the sample size. +- Adding a predictive or distributional objective **recovers and then exceeds** + the raw-observation floor. `correctness_iqn` at γ = 0 and + `conditional_decoupled_kelly_iqn` at γ = 0.99 are the only arms whose intervals + clear the floor. +- Greedy token accuracy is saturated near Bayes-optimal for every arm, so + representation quality and task performance are dissociable on this task. This + is the same dissociation the reward-state battery found, arrived at + independently. + +It does not support any fine-grained ranking among the middle arms, and the +current notes lead with several of those rankings. + +`PLAN.md` proposes what cycle 2 should do about it. diff --git a/experiments/mess3_token_guess_cycle_2/__init__.py b/experiments/mess3_token_guess_cycle_2/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/experiments/mess3_token_guess_cycle_2/audit/__init__.py b/experiments/mess3_token_guess_cycle_2/audit/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/experiments/mess3_token_guess_cycle_2/audit/experiment.py b/experiments/mess3_token_guess_cycle_2/audit/experiment.py new file mode 100644 index 00000000..091ebc61 --- /dev/null +++ b/experiments/mess3_token_guess_cycle_2/audit/experiment.py @@ -0,0 +1,298 @@ +"""Re-analyse the committed multi-seed token-guess results. + +This reads the results already in the repository and re-reports them against the +task's own floors and ceilings, with intervals that account for the three seeds +they were measured on. It trains nothing. Its purpose is to establish which of +the published orderings the existing evidence actually supports, so that cycle 2 +can be scoped to the questions that remain open. +""" + +from __future__ import annotations + +import glob +import json +from collections import defaultdict +from pathlib import Path +from typing import Any + +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt # noqa: E402 + +from experiments.mess3_token_guess_cycle_2.statistics import ( + compare, + holm_adjust, + summarise, +) +from harness.artifacts import RunArtifacts +from harness.context import RunContext + +STUDIES = ("mess_3_kelly_cycle_2", "mess_3_kelly_cycle_3") +METRICS = ("r_squared", "token_accuracy_greedy") +FALLBACK_REFERENCES = { + "belief_r2_floor": 0.9668, + "belief_r2_ceiling": 0.99888, + "accuracy_floor": 0.6732, + "accuracy_ceiling": 0.6883, + "untrained_module_r2": 0.8733, +} + + +def _references(experiment_dir: Path) -> dict[str, float]: + """Load the most recent reference run, or fall back to recorded values.""" + + candidates = sorted( + (experiment_dir.parent / "references" / "results").glob("*/references.json") + ) + if not candidates: + return dict(FALLBACK_REFERENCES) + data = json.loads(candidates[-1].read_text()) + untrained = data.get("untrained_module") or {} + return { + "belief_r2_floor": float(data["belief_r2_floor"]), + "belief_r2_ceiling": float(data["supervised_ceiling"]), + "accuracy_floor": float(data["accuracy_floor_repeat_previous_token"]), + "accuracy_ceiling": float(data["accuracy_ceiling_bayes"]), + "untrained_module_r2": float( + untrained.get("r_squared", FALLBACK_REFERENCES["untrained_module_r2"]) + ), + } + + +def collect_study(repository_root: Path, study: str) -> dict[str, dict[int, dict[str, float]]]: + """Read every ``condition_summary.json`` a study committed, keyed by seed.""" + + per_arm: dict[str, dict[int, dict[str, float]]] = defaultdict(dict) + pattern = str(repository_root / "experiments" / study / "*" / "results" / "*" / "condition_summary.json") + for path in sorted(glob.glob(pattern)): + summary = json.loads(Path(path).read_text()) + probe = summary.get("probe") or {} + seed = summary.get("seed") + if seed is None or "r_squared" not in probe: + continue + arm = Path(path).parents[2].name + per_arm[arm][int(seed)] = { + metric: float(probe[metric]) for metric in METRICS if metric in probe + } + return dict(per_arm) + + +def analyse_study( + per_arm: dict[str, dict[int, dict[str, float]]], + references: dict[str, float], +) -> dict[str, Any]: + """Summarise each arm and test every pairwise ordering within the study.""" + + floor = references["belief_r2_floor"] + ceiling = references["belief_r2_ceiling"] + conditions: dict[str, Any] = {} + r2_by_arm: dict[str, dict[int, float]] = {} + for arm, by_seed in sorted(per_arm.items()): + values = [by_seed[seed]["r_squared"] for seed in sorted(by_seed)] + if len(values) < 2: + continue + estimate = summarise(values) + r2_by_arm[arm] = {seed: by_seed[seed]["r_squared"] for seed in by_seed} + conditions[arm] = { + "seeds": sorted(by_seed), + "r_squared_mean": estimate.mean, + "r_squared_sample_sd": estimate.sample_sd, + "r_squared_ci": [estimate.ci_low, estimate.ci_high], + "fraction_of_usable_range": (estimate.mean - floor) / (ceiling - floor), + "exceeds_no_network_floor": bool(estimate.ci_low > floor), + } + + comparisons: dict[str, Any] = {} + arms = sorted(r2_by_arm) + raw_p: dict[str, float] = {} + for index, left in enumerate(arms): + for right in arms[index + 1:]: + name = f"{left} vs {right}" + result = compare(r2_by_arm[left], r2_by_arm[right]) + raw_p[name] = result.p_value + comparisons[name] = { + "difference": result.difference, + "difference_ci": [result.ci_low, result.ci_high], + "paired_sample_sd": result.sample_sd, + "n_shared_seeds": result.n, + "p_value": result.p_value, + "seeds_for_80_percent_power": result.seeds_for_power, + } + for name, adjusted in holm_adjust(raw_p).items(): + comparisons[name]["p_value_holm"] = adjusted + comparisons[name]["resolved"] = bool( + adjusted < 0.05 and comparisons[name]["difference_ci"][0] > 0.0 + ) or bool(adjusted < 0.05 and comparisons[name]["difference_ci"][1] < 0.0) + return {"conditions": conditions, "comparisons": comparisons} + + +def plot_against_references( + analysis: dict[str, Any], + references: dict[str, float], + *, + untrained_r2: float, + path: Path, +) -> None: + """Place every arm's interval inside the range the metric can move through.""" + + floor = references["belief_r2_floor"] + ceiling = references["belief_r2_ceiling"] + rows: list[tuple[str, float, float, float]] = [] + for study, result in analysis.items(): + label = study.replace("mess_3_", "").replace("_", " ") + for arm, values in sorted( + result["conditions"].items(), + key=lambda item: item[1]["r_squared_mean"], + ): + low, high = values["r_squared_ci"] + rows.append( + (f"{arm.replace('_', ' ')}\n({label})", values["r_squared_mean"], low, high) + ) + + figure, axis = plt.subplots(figsize=(9.5, 0.42 * len(rows) + 2.6)) + positions = range(len(rows)) + axis.axvspan(floor, ceiling, color="tab:green", alpha=0.10) + axis.axvline( + floor, + color="tab:green", + linestyle="--", + linewidth=1.4, + label=f"affine probe on raw observations ({floor:.4f})", + ) + axis.axvline( + untrained_r2, + color="tab:orange", + linestyle=":", + linewidth=1.4, + label=f"randomly initialised transformer ({untrained_r2:.4f})", + ) + axis.axvline( + ceiling, + color="tab:blue", + linestyle="-.", + linewidth=1.4, + label=f"supervised next-token model ({ceiling:.4f})", + ) + for position, (_, mean, low, high) in zip(positions, rows): + cleared = low > floor + axis.plot( + [low, high], + [position, position], + color="tab:green" if cleared else "tab:red", + linewidth=2.0, + alpha=0.75, + ) + axis.plot( + mean, + position, + "o", + color="tab:green" if cleared else "tab:red", + markersize=6, + ) + axis.set_yticks(list(positions), [row[0] for row in rows], fontsize=8) + axis.set_xlim(0.72, 1.02) + axis.set_xlabel("held-out belief-probe R² (95% CI over three seeds)") + axis.set_title( + "Committed token-guess results against the range the metric can move through" + ) + axis.grid(axis="x", alpha=0.2) + axis.legend(loc="lower left", fontsize=8, framealpha=0.9) + figure.tight_layout() + figure.savefig(path, dpi=200) + plt.close(figure) + + +def _findings(analysis: dict[str, Any], references: dict[str, float]) -> str: + floor = references["belief_r2_floor"] + ceiling = references["belief_r2_ceiling"] + lines = [ + "# Audit of the committed multi-seed token-guess results", + "", + "Belief-probe R² is reported against the range it can move through: 0% is " + f"an affine probe on the raw observations ({floor:.4f}) and 100% is the " + f"supervised next-token replication ({ceiling:.4f}).", + "", + ] + for study, result in analysis.items(): + lines.extend( + [ + f"## `{study}`", + "", + "| condition | R² | 95% CI | usable range | above floor |", + "|---|---:|---|---:|---|", + ] + ) + ordered = sorted( + result["conditions"].items(), + key=lambda item: -item[1]["r_squared_mean"], + ) + for arm, values in ordered: + low, high = values["r_squared_ci"] + above = "yes" if values["exceeds_no_network_floor"] else "no" + lines.append( + f"| `{arm}` | {values['r_squared_mean']:.4f} | " + f"[{low:.4f}, {high:.4f}] | " + f"{values['fraction_of_usable_range']:+.0%} | {above} |" + ) + resolved = [ + name + for name, values in result["comparisons"].items() + if values.get("resolved") + ] + total = len(result["comparisons"]) + lines.extend( + [ + "", + f"Of {total} pairwise orderings, {len(resolved)} survive a " + "Holm correction across the family.", + "", + "| comparison | difference | 95% CI | Holm p | seeds for 80% power |", + "|---|---:|---|---:|---:|", + ] + ) + for name, values in sorted( + result["comparisons"].items(), key=lambda item: item[1]["p_value"] + ): + low, high = values["difference_ci"] + needed = values["seeds_for_80_percent_power"] + lines.append( + f"| {name} | {values['difference']:+.4f} | " + f"[{low:+.4f}, {high:+.4f}] | " + f"{values['p_value_holm']:.3f} | " + f"{needed if needed < 10_000 else '>10000'} |" + ) + lines.append("") + return "\n".join(lines) + + +def run(context: RunContext) -> dict[str, Any]: + outputs = RunArtifacts.from_context(context) + outputs.prepare() + repository_root = Path(__file__).parents[3] + references = _references(Path(__file__).parents[1]) + analysis = { + study: analyse_study( + collect_study(repository_root, study), + references, + ) + for study in STUDIES + } + analysis = {study: result for study, result in analysis.items() if result["conditions"]} + if not analysis: + raise RuntimeError("no committed multi-seed results were found to audit") + figure_path = context.results_dir / "results_against_references.png" + plot_against_references( + analysis, + references, + untrained_r2=references["untrained_module_r2"], + path=figure_path, + ) + result = { + "references": references, + "studies": analysis, + "figure": str(figure_path), + } + outputs.write_json("audit.json", result) + (context.results_dir / "findings.md").write_text(_findings(analysis, references)) + return result diff --git a/experiments/mess3_token_guess_cycle_2/audit/results/20260725T234843Z-46848db2/audit.json b/experiments/mess3_token_guess_cycle_2/audit/results/20260725T234843Z-46848db2/audit.json new file mode 100644 index 00000000..e69f3147 --- /dev/null +++ b/experiments/mess3_token_guess_cycle_2/audit/results/20260725T234843Z-46848db2/audit.json @@ -0,0 +1,646 @@ +{ + "figure": "/workspace/experiments/mess3_token_guess_cycle_2/audit/results/20260725T234843Z-46848db2/results_against_references.png", + "references": { + "accuracy_ceiling": 0.6883, + "accuracy_floor": 0.6732, + "belief_r2_ceiling": 0.99888, + "belief_r2_floor": 0.9668, + "untrained_module_r2": 0.8733 + }, + "studies": { + "mess_3_kelly_cycle_2": { + "comparisons": { + "conditional_decoupled_kelly_iqn vs conditional_decoupled_kelly_mean": { + "difference": -0.006775144328574179, + "difference_ci": [ + -0.046992707495339764, + 0.03344241883819141 + ], + "n_shared_seeds": 3, + "p_value": 0.54388406860352, + "p_value_holm": 1.0, + "paired_sample_sd": 0.016189747845512227, + "resolved": false, + "seeds_for_80_percent_power": 47 + }, + "conditional_decoupled_kelly_iqn vs correctness_iqn": { + "difference": -0.03657977905470411, + "difference_ci": [ + -0.05634829673622678, + -0.016811261373181435 + ], + "n_shared_seeds": 3, + "p_value": 0.01541211009914722, + "p_value_holm": 0.27741798178464994, + "paired_sample_sd": 0.007957899269438538, + "resolved": false, + "seeds_for_80_percent_power": 3 + }, + "conditional_decoupled_kelly_iqn vs correctness_mean": { + "difference": -0.0343501413875947, + "difference_ci": [ + -0.05709634836414988, + -0.011603934411039518 + ], + "n_shared_seeds": 3, + "p_value": 0.022876187524714233, + "p_value_holm": 0.36601900039542773, + "paired_sample_sd": 0.009156580518448046, + "resolved": false, + "seeds_for_80_percent_power": 3 + }, + "conditional_decoupled_kelly_iqn vs coupled_kelly_iqn": { + "difference": 0.002415095533641387, + "difference_ci": [ + -0.017357691076054588, + 0.022187882143337362 + ], + "n_shared_seeds": 3, + "p_value": 0.651663786941959, + "p_value_holm": 1.0, + "paired_sample_sd": 0.007959617744285205, + "resolved": false, + "seeds_for_80_percent_power": 88 + }, + "conditional_decoupled_kelly_iqn vs coupled_kelly_mean": { + "difference": 0.01156453197739645, + "difference_ci": [ + -0.019415723606662535, + 0.04254478756145543 + ], + "n_shared_seeds": 3, + "p_value": 0.24947735945521132, + "p_value_holm": 1.0, + "paired_sample_sd": 0.012471231138884894, + "resolved": false, + "seeds_for_80_percent_power": 12 + }, + "conditional_decoupled_kelly_iqn vs decoupled_kelly_iqn": { + "difference": -0.00816893560137418, + "difference_ci": [ + -0.03588395840427449, + 0.019546087201526133 + ], + "n_shared_seeds": 3, + "p_value": 0.332373029722586, + "p_value_holm": 1.0, + "paired_sample_sd": 0.011156798059867716, + "resolved": false, + "seeds_for_80_percent_power": 17 + }, + "conditional_decoupled_kelly_iqn vs decoupled_kelly_mean": { + "difference": -0.003844048137894592, + "difference_ci": [ + -0.035160135613054276, + 0.02747203933726509 + ], + "n_shared_seeds": 3, + "p_value": 0.6501427336852326, + "p_value_holm": 1.0, + "paired_sample_sd": 0.012606421667780358, + "resolved": false, + "seeds_for_80_percent_power": 87 + }, + "conditional_decoupled_kelly_mean vs correctness_iqn": { + "difference": -0.02980463472612993, + "difference_ci": [ + -0.052096990111295635, + -0.007512279340964222 + ], + "n_shared_seeds": 3, + "p_value": 0.028914213440128043, + "p_value_holm": 0.4337132016019206, + "paired_sample_sd": 0.008973880666808305, + "resolved": false, + "seeds_for_80_percent_power": 3 + }, + "conditional_decoupled_kelly_mean vs correctness_mean": { + "difference": -0.027574997059020518, + "difference_ci": [ + -0.051758831760500626, + -0.00339116235754041 + ], + "n_shared_seeds": 3, + "p_value": 0.039125535965817734, + "p_value_holm": 0.5477575035214483, + "paired_sample_sd": 0.00973530355708022, + "resolved": false, + "seeds_for_80_percent_power": 4 + }, + "conditional_decoupled_kelly_mean vs coupled_kelly_iqn": { + "difference": 0.009190239862215566, + "difference_ci": [ + -0.016638458798726753, + 0.035018938523157885 + ], + "n_shared_seeds": 3, + "p_value": 0.26544296634561537, + "p_value_holm": 1.0, + "paired_sample_sd": 0.010397450406541008, + "resolved": false, + "seeds_for_80_percent_power": 13 + }, + "conditional_decoupled_kelly_mean vs coupled_kelly_mean": { + "difference": 0.018339676305970627, + "difference_ci": [ + 0.009096585841383842, + 0.02758276677055741 + ], + "n_shared_seeds": 3, + "p_value": 0.013444687050300122, + "p_value_holm": 0.263059411408765, + "paired_sample_sd": 0.0037208446298551134, + "resolved": false, + "seeds_for_80_percent_power": 4 + }, + "conditional_decoupled_kelly_mean vs decoupled_kelly_iqn": { + "difference": -0.0013937912728000008, + "difference_ci": [ + -0.014143649281822178, + 0.011356066736222175 + ], + "n_shared_seeds": 3, + "p_value": 0.6844041690055843, + "p_value_holm": 1.0, + "paired_sample_sd": 0.005132508535542739, + "resolved": false, + "seeds_for_80_percent_power": 109 + }, + "conditional_decoupled_kelly_mean vs decoupled_kelly_mean": { + "difference": 0.0029310961906795865, + "difference_ci": [ + -0.009622355191312848, + 0.01548454757267202 + ], + "n_shared_seeds": 3, + "p_value": 0.4208743533896648, + "p_value_holm": 1.0, + "paired_sample_sd": 0.005053444228398771, + "resolved": false, + "seeds_for_80_percent_power": 26 + }, + "correctness_iqn vs correctness_mean": { + "difference": 0.002229637667109411, + "difference_ci": [ + -0.0037297437911992045, + 0.008189019125418027 + ], + "n_shared_seeds": 3, + "p_value": 0.24873075583485868, + "p_value_holm": 1.0, + "paired_sample_sd": 0.0023989738693309475, + "resolved": false, + "seeds_for_80_percent_power": 12 + }, + "correctness_iqn vs coupled_kelly_iqn": { + "difference": 0.038994874588345496, + "difference_ci": [ + 0.022091693040576402, + 0.05589805613611459 + ], + "n_shared_seeds": 3, + "p_value": 0.009997654403249677, + "p_value_holm": 0.2199483968714929, + "paired_sample_sd": 0.006804446254253378, + "resolved": false, + "seeds_for_80_percent_power": 4 + }, + "correctness_iqn vs coupled_kelly_mean": { + "difference": 0.04814431103210056, + "difference_ci": [ + 0.034619652401099085, + 0.061668969663102036 + ], + "n_shared_seeds": 3, + "p_value": 0.004235687772692276, + "p_value_holm": 0.11436356986269144, + "paired_sample_sd": 0.00544440775848613, + "resolved": false, + "seeds_for_80_percent_power": 3 + }, + "correctness_iqn vs decoupled_kelly_iqn": { + "difference": 0.02841084345332993, + "difference_ci": [ + 0.018670338832495144, + 0.03815134807416472 + ], + "n_shared_seeds": 3, + "p_value": 0.0062894082495708005, + "p_value_holm": 0.1635246144888408, + "paired_sample_sd": 0.003921080773727154, + "resolved": false, + "seeds_for_80_percent_power": 3 + }, + "correctness_iqn vs decoupled_kelly_mean": { + "difference": 0.032735730916809515, + "difference_ci": [ + 0.021071896248949833, + 0.0443995655846692 + ], + "n_shared_seeds": 3, + "p_value": 0.006787762058179193, + "p_value_holm": 0.16969405145447983, + "paired_sample_sd": 0.004695325308531833, + "resolved": false, + "seeds_for_80_percent_power": 4 + }, + "correctness_mean vs coupled_kelly_iqn": { + "difference": 0.03676523692123609, + "difference_ci": [ + 0.013907050010196122, + 0.05962342383227605 + ], + "n_shared_seeds": 3, + "p_value": 0.02024825417309691, + "p_value_holm": 0.34422032094264743, + "paired_sample_sd": 0.009201658508269263, + "resolved": false, + "seeds_for_80_percent_power": 3 + }, + "correctness_mean vs coupled_kelly_mean": { + "difference": 0.045914673364991145, + "difference_ci": [ + 0.029292290665425925, + 0.06253705606455637 + ], + "n_shared_seeds": 3, + "p_value": 0.007005330955178829, + "p_value_holm": 0.16969405145447983, + "paired_sample_sd": 0.006691409506380806, + "resolved": false, + "seeds_for_80_percent_power": 3 + }, + "correctness_mean vs decoupled_kelly_iqn": { + "difference": 0.02618120578622052, + "difference_ci": [ + 0.013132866420991107, + 0.03922954515144993 + ], + "n_shared_seeds": 3, + "p_value": 0.01315297057043825, + "p_value_holm": 0.263059411408765, + "paired_sample_sd": 0.005252663450785712, + "resolved": false, + "seeds_for_80_percent_power": 3 + }, + "correctness_mean vs decoupled_kelly_mean": { + "difference": 0.030506093249700106, + "difference_ci": [ + 0.018615555078888603, + 0.04239663142051161 + ], + "n_shared_seeds": 3, + "p_value": 0.008106828749776702, + "p_value_holm": 0.18645706124486416, + "paired_sample_sd": 0.004786585749480611, + "resolved": false, + "seeds_for_80_percent_power": 3 + }, + "coupled_kelly_iqn vs coupled_kelly_mean": { + "difference": 0.009149436443755063, + "difference_ci": [ + -0.008841111720722017, + 0.027139984608232144 + ], + "n_shared_seeds": 3, + "p_value": 0.16013621152392507, + "p_value_holm": 1.0, + "paired_sample_sd": 0.007242170222439436, + "resolved": false, + "seeds_for_80_percent_power": 8 + }, + "coupled_kelly_iqn vs decoupled_kelly_iqn": { + "difference": -0.010584031135015567, + "difference_ci": [ + -0.02735921643873826, + 0.006191154168707127 + ], + "n_shared_seeds": 3, + "p_value": 0.11312782092798437, + "p_value_holm": 1.0, + "paired_sample_sd": 0.006752920832196076, + "resolved": false, + "seeds_for_80_percent_power": 6 + }, + "coupled_kelly_iqn vs decoupled_kelly_mean": { + "difference": -0.006259143671535979, + "difference_ci": [ + -0.029159237207066722, + 0.016640949863994767 + ], + "n_shared_seeds": 3, + "p_value": 0.3606159657395652, + "p_value_holm": 1.0, + "paired_sample_sd": 0.009218528194797735, + "resolved": false, + "seeds_for_80_percent_power": 20 + }, + "coupled_kelly_mean vs decoupled_kelly_iqn": { + "difference": -0.01973346757877063, + "difference_ci": [ + -0.023522459887497904, + -0.015944475270043355 + ], + "n_shared_seeds": 3, + "p_value": 0.001985514653420706, + "p_value_holm": 0.055594410295779766, + "paired_sample_sd": 0.0015252746620305286, + "resolved": false, + "seeds_for_80_percent_power": 5 + }, + "coupled_kelly_mean vs decoupled_kelly_mean": { + "difference": -0.01540858011529104, + "difference_ci": [ + -0.022898755032504922, + -0.007918405198077159 + ], + "n_shared_seeds": 3, + "p_value": 0.012524685671792615, + "p_value_holm": 0.2630183991076449, + "paired_sample_sd": 0.003015201162876065, + "resolved": false, + "seeds_for_80_percent_power": 3 + }, + "decoupled_kelly_iqn vs decoupled_kelly_mean": { + "difference": 0.0043248874634795875, + "difference_ci": [ + -0.0019633062598596335, + 0.010613081186818809 + ], + "n_shared_seeds": 3, + "p_value": 0.09773629300302608, + "p_value_holm": 1.0, + "paired_sample_sd": 0.0025313386184651347, + "resolved": false, + "seeds_for_80_percent_power": 5 + } + }, + "conditions": { + "conditional_decoupled_kelly_iqn": { + "exceeds_no_network_floor": false, + "fraction_of_usable_range": -0.5524138288833105, + "r_squared_ci": [ + 0.9288742088394529, + 0.9692829198993939 + ], + "r_squared_mean": 0.9490785643694234, + "r_squared_sample_sd": 0.008133347613701497, + "seeds": [ + 42, + 43, + 44 + ] + }, + "conditional_decoupled_kelly_mean": { + "exceeds_no_network_floor": false, + "fraction_of_usable_range": -0.34121855679558666, + "r_squared_ci": [ + 0.9330018921258707, + 0.9787055252701244 + ], + "r_squared_mean": 0.9558537086979976, + "r_squared_sample_sd": 0.009199094101761956, + "seeds": [ + 42, + 43, + 44 + ] + }, + "correctness_iqn": { + "exceeds_no_network_floor": true, + "fraction_of_usable_range": 0.5878535980089589, + "r_squared_ci": [ + 0.9841300313251463, + 0.9871866555231085 + ], + "r_squared_mean": 0.9856583434241274, + "r_squared_sample_sd": 0.0006152284117550854, + "seeds": [ + 42, + 43, + 44 + ] + }, + "correctness_mean": { + "exceeds_no_network_floor": true, + "fraction_of_usable_range": 0.5183511769644058, + "r_squared_ci": [ + 0.9789762448799507, + 0.9878811666340855 + ], + "r_squared_mean": 0.9834287057570181, + "r_squared_sample_sd": 0.0017923567022901537, + "seeds": [ + 42, + 43, + 44 + ] + }, + "coupled_kelly_iqn": { + "exceeds_no_network_floor": false, + "fraction_of_usable_range": -0.6276973554930759, + "r_squared_ci": [ + 0.9282569885039877, + 0.9650699491675765 + ], + "r_squared_mean": 0.9466634688357821, + "r_squared_sample_sd": 0.007409605451714317, + "seeds": [ + 42, + 43, + 44 + ] + }, + "coupled_kelly_mean": { + "exceeds_no_network_floor": false, + "fraction_of_usable_range": -0.9129042271812025, + "r_squared_ci": [ + 0.9231788371987796, + 0.9518492275852745 + ], + "r_squared_mean": 0.937514032392027, + "r_squared_sample_sd": 0.005770692633278706, + "seeds": [ + 42, + 43, + 44 + ] + }, + "decoupled_kelly_iqn": { + "exceeds_no_network_floor": false, + "fraction_of_usable_range": -0.29777119791778345, + "r_squared_ci": [ + 0.9466948105897789, + 0.9678001893518161 + ], + "r_squared_mean": 0.9572474999707975, + "r_squared_sample_sd": 0.004248029137476117, + "seeds": [ + 42, + 43, + 44 + ] + }, + "decoupled_kelly_mean": { + "exceeds_no_network_floor": false, + "fraction_of_usable_range": -0.4325868919165237, + "r_squared_ci": [ + 0.9412719883024149, + 0.964573236712221 + ], + "r_squared_mean": 0.9529226125073179, + "r_squared_sample_sd": 0.004690007381552904, + "seeds": [ + 42, + 43, + 44 + ] + } + } + }, + "mess_3_kelly_cycle_3": { + "comparisons": { + "conditional_decoupled_kelly_iqn vs conditional_decoupled_kelly_mean": { + "difference": 0.010632126578209148, + "difference_ci": [ + -0.0009456873161711038, + 0.0222099404725894 + ], + "n_shared_seeds": 3, + "p_value": 0.05848986190288163, + "p_value_holm": 0.35093917141728975, + "paired_sample_sd": 0.004660697287278206, + "resolved": false, + "seeds_for_80_percent_power": 4 + }, + "conditional_decoupled_kelly_iqn vs iqn": { + "difference": 0.013580476188700619, + "difference_ci": [ + -0.01114735073030133, + 0.03830830310770257 + ], + "n_shared_seeds": 3, + "p_value": 0.14193220803860646, + "p_value_holm": 0.35093917141728975, + "paired_sample_sd": 0.009954289893847574, + "resolved": false, + "seeds_for_80_percent_power": 7 + }, + "conditional_decoupled_kelly_iqn vs ppo": { + "difference": 0.10814250098337959, + "difference_ci": [ + -0.019795632818286987, + 0.23608063478504615 + ], + "n_shared_seeds": 3, + "p_value": 0.06798314215975723, + "p_value_holm": 0.35093917141728975, + "paired_sample_sd": 0.051502029535842854, + "resolved": false, + "seeds_for_80_percent_power": 5 + }, + "conditional_decoupled_kelly_mean vs iqn": { + "difference": 0.0029483496104914706, + "difference_ci": [ + -0.03041299605622082, + 0.03630969527720376 + ], + "n_shared_seeds": 3, + "p_value": 0.7403433780887467, + "p_value_holm": 0.7403433780887467, + "paired_sample_sd": 0.013429748885864219, + "resolved": false, + "seeds_for_80_percent_power": 165 + }, + "conditional_decoupled_kelly_mean vs ppo": { + "difference": 0.09751037440517045, + "difference_ci": [ + -0.030237345617045205, + 0.2252580944273861 + ], + "n_shared_seeds": 3, + "p_value": 0.0815332366987508, + "p_value_holm": 0.35093917141728975, + "paired_sample_sd": 0.05142537767449465, + "resolved": false, + "seeds_for_80_percent_power": 5 + }, + "iqn vs ppo": { + "difference": 0.09456202479467897, + "difference_ci": [ + -0.016734697690015884, + 0.2058587472793738 + ], + "n_shared_seeds": 3, + "p_value": 0.06735492093793352, + "p_value_holm": 0.35093917141728975, + "paired_sample_sd": 0.044802959980135275, + "resolved": false, + "seeds_for_80_percent_power": 5 + } + }, + "conditions": { + "conditional_decoupled_kelly_iqn": { + "exceeds_no_network_floor": true, + "fraction_of_usable_range": 0.48549322472516526, + "r_squared_ci": [ + 0.9788439340639512, + 0.9859053112344154 + ], + "r_squared_mean": 0.9823746226491833, + "r_squared_sample_sd": 0.0014212934204619135, + "seeds": [ + 42, + 43, + 44 + ] + }, + "conditional_decoupled_kelly_mean": { + "exceeds_no_network_floor": false, + "fraction_of_usable_range": 0.1540678326363491, + "r_squared_ci": [ + 0.9595826818879015, + 0.9839023102540466 + ], + "r_squared_mean": 0.9717424960709741, + "r_squared_sample_sd": 0.004894983931669666, + "seeds": [ + 42, + 43, + 44 + ] + }, + "iqn": { + "exceeds_no_network_floor": false, + "fraction_of_usable_range": 0.0621616727083165, + "r_squared_ci": [ + 0.941299174986669, + 0.9962891179342965 + ], + "r_squared_mean": 0.9687941464604828, + "r_squared_sample_sd": 0.01106821547926209, + "seeds": [ + 42, + 43, + 44 + ] + }, + "ppo": { + "exceeds_no_network_floor": false, + "fraction_of_usable_range": -2.885532367026067, + "r_squared_ci": [ + 0.7427668105896438, + 1.0056974327419637 + ], + "r_squared_mean": 0.8742321216658038, + "r_squared_sample_sd": 0.05292190946351729, + "seeds": [ + 42, + 43, + 44 + ] + } + } + } + } +} diff --git a/experiments/mess3_token_guess_cycle_2/audit/results/20260725T234843Z-46848db2/findings.md b/experiments/mess3_token_guess_cycle_2/audit/results/20260725T234843Z-46848db2/findings.md new file mode 100644 index 00000000..88874326 --- /dev/null +++ b/experiments/mess3_token_guess_cycle_2/audit/results/20260725T234843Z-46848db2/findings.md @@ -0,0 +1,69 @@ +# Audit of the committed multi-seed token-guess results + +Belief-probe R² is reported against the range it can move through: 0% is an affine probe on the raw observations (0.9668) and 100% is the supervised next-token replication (0.9989). + +## `mess_3_kelly_cycle_2` + +| condition | R² | 95% CI | usable range | above floor | +|---|---:|---|---:|---| +| `correctness_iqn` | 0.9857 | [0.9841, 0.9872] | +59% | yes | +| `correctness_mean` | 0.9834 | [0.9790, 0.9879] | +52% | yes | +| `decoupled_kelly_iqn` | 0.9572 | [0.9467, 0.9678] | -30% | no | +| `conditional_decoupled_kelly_mean` | 0.9559 | [0.9330, 0.9787] | -34% | no | +| `decoupled_kelly_mean` | 0.9529 | [0.9413, 0.9646] | -43% | no | +| `conditional_decoupled_kelly_iqn` | 0.9491 | [0.9289, 0.9693] | -55% | no | +| `coupled_kelly_iqn` | 0.9467 | [0.9283, 0.9651] | -63% | no | +| `coupled_kelly_mean` | 0.9375 | [0.9232, 0.9518] | -91% | no | + +Of 28 pairwise orderings, 0 survive a Holm correction across the family. + +| comparison | difference | 95% CI | Holm p | seeds for 80% power | +|---|---:|---|---:|---:| +| coupled_kelly_mean vs decoupled_kelly_iqn | -0.0197 | [-0.0235, -0.0159] | 0.056 | 5 | +| correctness_iqn vs coupled_kelly_mean | +0.0481 | [+0.0346, +0.0617] | 0.114 | 3 | +| correctness_iqn vs decoupled_kelly_iqn | +0.0284 | [+0.0187, +0.0382] | 0.164 | 3 | +| correctness_iqn vs decoupled_kelly_mean | +0.0327 | [+0.0211, +0.0444] | 0.170 | 4 | +| correctness_mean vs coupled_kelly_mean | +0.0459 | [+0.0293, +0.0625] | 0.170 | 3 | +| correctness_mean vs decoupled_kelly_mean | +0.0305 | [+0.0186, +0.0424] | 0.186 | 3 | +| correctness_iqn vs coupled_kelly_iqn | +0.0390 | [+0.0221, +0.0559] | 0.220 | 4 | +| coupled_kelly_mean vs decoupled_kelly_mean | -0.0154 | [-0.0229, -0.0079] | 0.263 | 3 | +| correctness_mean vs decoupled_kelly_iqn | +0.0262 | [+0.0131, +0.0392] | 0.263 | 3 | +| conditional_decoupled_kelly_mean vs coupled_kelly_mean | +0.0183 | [+0.0091, +0.0276] | 0.263 | 4 | +| conditional_decoupled_kelly_iqn vs correctness_iqn | -0.0366 | [-0.0563, -0.0168] | 0.277 | 3 | +| correctness_mean vs coupled_kelly_iqn | +0.0368 | [+0.0139, +0.0596] | 0.344 | 3 | +| conditional_decoupled_kelly_iqn vs correctness_mean | -0.0344 | [-0.0571, -0.0116] | 0.366 | 3 | +| conditional_decoupled_kelly_mean vs correctness_iqn | -0.0298 | [-0.0521, -0.0075] | 0.434 | 3 | +| conditional_decoupled_kelly_mean vs correctness_mean | -0.0276 | [-0.0518, -0.0034] | 0.548 | 4 | +| decoupled_kelly_iqn vs decoupled_kelly_mean | +0.0043 | [-0.0020, +0.0106] | 1.000 | 5 | +| coupled_kelly_iqn vs decoupled_kelly_iqn | -0.0106 | [-0.0274, +0.0062] | 1.000 | 6 | +| coupled_kelly_iqn vs coupled_kelly_mean | +0.0091 | [-0.0088, +0.0271] | 1.000 | 8 | +| correctness_iqn vs correctness_mean | +0.0022 | [-0.0037, +0.0082] | 1.000 | 12 | +| conditional_decoupled_kelly_iqn vs coupled_kelly_mean | +0.0116 | [-0.0194, +0.0425] | 1.000 | 12 | +| conditional_decoupled_kelly_mean vs coupled_kelly_iqn | +0.0092 | [-0.0166, +0.0350] | 1.000 | 13 | +| conditional_decoupled_kelly_iqn vs decoupled_kelly_iqn | -0.0082 | [-0.0359, +0.0195] | 1.000 | 17 | +| coupled_kelly_iqn vs decoupled_kelly_mean | -0.0063 | [-0.0292, +0.0166] | 1.000 | 20 | +| conditional_decoupled_kelly_mean vs decoupled_kelly_mean | +0.0029 | [-0.0096, +0.0155] | 1.000 | 26 | +| conditional_decoupled_kelly_iqn vs conditional_decoupled_kelly_mean | -0.0068 | [-0.0470, +0.0334] | 1.000 | 47 | +| conditional_decoupled_kelly_iqn vs decoupled_kelly_mean | -0.0038 | [-0.0352, +0.0275] | 1.000 | 87 | +| conditional_decoupled_kelly_iqn vs coupled_kelly_iqn | +0.0024 | [-0.0174, +0.0222] | 1.000 | 88 | +| conditional_decoupled_kelly_mean vs decoupled_kelly_iqn | -0.0014 | [-0.0141, +0.0114] | 1.000 | 109 | + +## `mess_3_kelly_cycle_3` + +| condition | R² | 95% CI | usable range | above floor | +|---|---:|---|---:|---| +| `conditional_decoupled_kelly_iqn` | 0.9824 | [0.9788, 0.9859] | +49% | yes | +| `conditional_decoupled_kelly_mean` | 0.9717 | [0.9596, 0.9839] | +15% | no | +| `iqn` | 0.9688 | [0.9413, 0.9963] | +6% | no | +| `ppo` | 0.8742 | [0.7428, 1.0057] | -289% | no | + +Of 6 pairwise orderings, 0 survive a Holm correction across the family. + +| comparison | difference | 95% CI | Holm p | seeds for 80% power | +|---|---:|---|---:|---:| +| conditional_decoupled_kelly_iqn vs conditional_decoupled_kelly_mean | +0.0106 | [-0.0009, +0.0222] | 0.351 | 4 | +| iqn vs ppo | +0.0946 | [-0.0167, +0.2059] | 0.351 | 5 | +| conditional_decoupled_kelly_iqn vs ppo | +0.1081 | [-0.0198, +0.2361] | 0.351 | 5 | +| conditional_decoupled_kelly_mean vs ppo | +0.0975 | [-0.0302, +0.2253] | 0.351 | 5 | +| conditional_decoupled_kelly_iqn vs iqn | +0.0136 | [-0.0111, +0.0383] | 0.351 | 7 | +| conditional_decoupled_kelly_mean vs iqn | +0.0029 | [-0.0304, +0.0363] | 0.740 | 165 | diff --git a/experiments/mess3_token_guess_cycle_2/audit/results/20260725T234843Z-46848db2/results_against_references.png b/experiments/mess3_token_guess_cycle_2/audit/results/20260725T234843Z-46848db2/results_against_references.png new file mode 100644 index 00000000..e9c4d217 Binary files /dev/null and b/experiments/mess3_token_guess_cycle_2/audit/results/20260725T234843Z-46848db2/results_against_references.png differ diff --git a/experiments/mess3_token_guess_cycle_2/audit/results/20260725T234843Z-46848db2/run_manifest.json b/experiments/mess3_token_guess_cycle_2/audit/results/20260725T234843Z-46848db2/run_manifest.json new file mode 100644 index 00000000..18a90378 --- /dev/null +++ b/experiments/mess3_token_guess_cycle_2/audit/results/20260725T234843Z-46848db2/run_manifest.json @@ -0,0 +1,75 @@ +{ + "command": [ + "/workspace/.venv/bin/rl-harness", + "experiments.mess3_token_guess_cycle_2.audit.experiment", + "--no-upload-artifacts" + ], + "dependency_lock": { + "file": "/workspace/uv.lock", + "sha256": "986725e46af5252423583fe947d30b2646f541a4863acb8c51d578a7ca1b8220" + }, + "ended_at": "2026-07-25T23:48:44.772373+00:00", + "error": null, + "experiment": { + "file": "/workspace/experiments/mess3_token_guess_cycle_2/audit/experiment.py", + "module": "experiments.mess3_token_guess_cycle_2.audit.experiment", + "source_sha256": "bd96a2324c233c02c6c1dfeedd9f77ceb5b7ff4fe61ef8bcfd99cc72cc0274f0" + }, + "framework_versions": { + "gymnasium": "1.2.2", + "numpy": "2.5.1", + "python": "3.14.6", + "ray": "2.56.0", + "rl-harness": "0.1.0", + "torch": "2.12.1" + }, + "git": { + "commit": "9015518e4fd44bbfef57c8b938bf4be828147be4", + "dirty": true, + "experiment_repository": { + "commit": "9015518e4fd44bbfef57c8b938bf4be828147be4", + "dirty": true, + "root": "/workspace" + }, + "library": { + "commit": "7c231655be83affb51c1d24d8cfd475b5d425eb5", + "dirty": false, + "package": "rl-harness", + "root": "/rl-harness-src", + "version": "0.1.0" + } + }, + "hardware": { + "cpu_count": 4, + "cuda_available": false, + "machine": "x86_64", + "mps_available": false, + "platform": "Linux-6.12.94+-x86_64-with-glibc2.39", + "profile": { + "learner_device": "cpu", + "name": "cpu", + "num_env_runners": null, + "num_envs_per_env_runner": 24, + "num_gpus_per_env_runner": 0.0, + "torch_threads": null + } + }, + "run_id": "20260725T234843Z-46848db2", + "runtime": { + "artifacts_dir": "/workspace/experiments/mess3_token_guess_cycle_2/audit/artifacts/20260725T234843Z-46848db2", + "overrides": { + "hardware_profile": "auto", + "resume_from": null, + "seed": 42, + "smoke": false, + "upload_artifacts": false + }, + "results_dir": "/workspace/experiments/mess3_token_guess_cycle_2/audit/results/20260725T234843Z-46848db2", + "resume_from": null, + "seed": 42, + "smoke": false + }, + "schema_version": 2, + "started_at": "2026-07-25T23:48:44.456151+00:00", + "status": "completed" +} diff --git a/experiments/mess3_token_guess_cycle_2/metric_references.py b/experiments/mess3_token_guess_cycle_2/metric_references.py new file mode 100644 index 00000000..d24a35ea --- /dev/null +++ b/experiments/mess3_token_guess_cycle_2/metric_references.py @@ -0,0 +1,342 @@ +"""Task-intrinsic reference points for the passive MESS3 token-guess study. + +Cycle 1 reported belief-probe R² and greedy token accuracy as bare numbers. Both +metrics have a task-imposed floor and ceiling that are large relative to the +differences between conditions, so a bare number cannot be read as a measure of +how much belief structure a training objective induced. + +This module derives those reference points from the environment definition and +from the same probe estimator the conditions use, so that condition scores can be +reported as a fraction of the range the metric can actually move through. + +Three references matter: + +``bayes_accuracy`` + Greedy token accuracy of the exact Bayesian filter, as a function of how many + past tokens it is allowed to see. One token reproduces the trivial + "repeat the previous token" rule; the sequence converges well inside the + 64-token context the agents receive. + +``raw_token_window_r2`` + Belief-probe R² for an affine probe read directly off the one-hot encoded + last ``k`` observations. No network and no training are involved, so this is + the score a policy earns for passing its own inputs through unchanged. + +``untrained_module_r2`` + The same probe applied to a randomly initialised copy of the study's + transformer. This controls for architecture, embedding width, and the final + LayerNorm, none of which the raw-token probe exercises. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +import numpy as np +import torch + +from analysis.probes import r2_score +from envs.hmm import HMMEnv +from envs.mess3.model import passive_model +from experiments.mess3_belief_geometry_2026_07.probe import ( + collect_probe_data, + make_transducer_target, +) +from experiments.mess3_token_guess_cycle_1.analysis import ( + PROBE_RANK, + fit_reduced_rank_affine, +) +from harness.seeding import named_seed_sequences, seed_sequence_to_int +from learners.models.transformer import TransformerModel + +ALPHA = 0.85 +WARMUP = 64 +CONTEXT_LENGTHS = (1, 2, 3, 4, 6, 8, 16, 64) +BOOTSTRAP_RESAMPLES = 400 + +_STREAM_KEYS = { + "reference_fit": (300,), + "reference_test": (301,), + "untrained_fit": (302,), + "untrained_test": (303,), + "bootstrap": (304,), +} + + +@dataclass(frozen=True, slots=True) +class TokenStream: + """Exact predictive beliefs aligned with the tokens they predict.""" + + beliefs: np.ndarray + tokens: np.ndarray + windows: np.ndarray + + +def simulate_stream(*, n_steps: int, seed: int, window: int = WARMUP) -> TokenStream: + """Simulate passive MESS3 and record the exact predictive belief per step. + + The belief is the distribution over the state that emits the token being + predicted, conditioned on every token revealed so far. This matches the + ``delay=1`` transducer target the conditions probe against, where the + filtering operator is ``diag(P(y|s)) @ T``. + """ + + if n_steps <= 0 or window <= 0: + raise ValueError("n_steps and window must be positive") + model = passive_model(alpha=ALPHA) + transition = np.asarray(model.transition_matrix, dtype=np.float64) + emission = np.asarray(model.emission_matrix, dtype=np.float64) + initial = np.asarray(model.initial_distribution, dtype=np.float64) + + rng = np.random.default_rng(seed) + state = int(rng.choice(len(initial), p=initial)) + belief = initial.copy() + history: list[int] = [0] * window + + beliefs = np.empty((n_steps, len(initial)), dtype=np.float64) + tokens = np.empty(n_steps, dtype=np.int64) + windows = np.empty((n_steps, window), dtype=np.int64) + + for step in range(n_steps + window): + token = int(rng.choice(emission.shape[1], p=emission[state])) + if step >= window: + index = step - window + beliefs[index] = belief + tokens[index] = token + windows[index] = history[-window:] + history.append(token) + posterior = belief * emission[:, token] + belief = (posterior / posterior.sum()) @ transition + state = int(rng.choice(transition.shape[1], p=transition[state])) + return TokenStream(beliefs=beliefs, tokens=tokens, windows=windows) + + +def _one_hot_window(windows: np.ndarray, k: int) -> np.ndarray: + return np.eye(3, dtype=np.float64)[windows[:, -k:]].reshape(len(windows), -1) + + +def _probe_r2( + fit_features: np.ndarray, + fit_targets: np.ndarray, + test_features: np.ndarray, + test_targets: np.ndarray, +) -> tuple[float, np.ndarray]: + weight, bias = fit_reduced_rank_affine( + fit_features, + fit_targets, + rank=PROBE_RANK, + ) + predicted = test_features @ weight + bias + return r2_score(predicted, test_targets), predicted + + +def raw_token_window_r2( + fit: TokenStream, + test: TokenStream, + *, + context_lengths: tuple[int, ...] = CONTEXT_LENGTHS, +) -> dict[int, float]: + """Score the study's probe against one-hot encoded raw observations.""" + + scores: dict[int, float] = {} + for k in context_lengths: + score, _ = _probe_r2( + _one_hot_window(fit.windows, k), + fit.beliefs, + _one_hot_window(test.windows, k), + test.beliefs, + ) + scores[k] = float(score) + return scores + + +def bootstrap_r2_interval( + fit: TokenStream, + test: TokenStream, + *, + context_length: int, + seed: int, + resamples: int = BOOTSTRAP_RESAMPLES, +) -> tuple[float, float]: + """Bound the probe's own sampling noise at the conditions' test-set size.""" + + _, predicted = _probe_r2( + _one_hot_window(fit.windows, context_length), + fit.beliefs, + _one_hot_window(test.windows, context_length), + test.beliefs, + ) + rng = np.random.default_rng(seed) + n = len(test.beliefs) + scores = [ + r2_score(predicted[index], test.beliefs[index]) + for index in (rng.integers(0, n, n) for _ in range(resamples)) + ] + low, high = np.percentile(scores, [2.5, 97.5]) + return float(low), float(high) + + +def bayes_accuracy_by_context( + stream: TokenStream, + *, + context_lengths: tuple[int, ...] = CONTEXT_LENGTHS, +) -> dict[int, float]: + """Greedy accuracy of an exact filter restricted to the last ``k`` tokens.""" + + model = passive_model(alpha=ALPHA) + transition = np.asarray(model.transition_matrix, dtype=np.float64) + emission = np.asarray(model.emission_matrix, dtype=np.float64) + initial = np.asarray(model.initial_distribution, dtype=np.float64) + + accuracies: dict[int, float] = {} + for k in context_lengths: + belief = np.repeat(initial[None, :], len(stream.windows), axis=0) + for offset in range(k): + belief = belief * emission[:, stream.windows[:, -k + offset]].T + belief /= belief.sum(axis=1, keepdims=True) + belief = belief @ transition + predicted = (belief @ emission).argmax(axis=1) + accuracies[k] = float((predicted == stream.tokens).mean()) + return accuracies + + +def _build_untrained_module(env_config: dict[str, Any], model_config: dict[str, Any]): + from ray.rllib.core.rl_module.rl_module import RLModuleSpec + + environment = HMMEnv(env_config) + try: + spec = RLModuleSpec( + module_class=TransformerModel, + model_config=dict(model_config), + observation_space=environment.observation_space, + action_space=environment.action_space, + ) + return spec.build() + finally: + environment.close() + + +def untrained_module_r2( + *, + env_config: dict[str, Any], + model_config: dict[str, Any], + seed: int, + fit_steps: int, + test_steps: int, + device: str = "cpu", +) -> dict[str, float]: + """Probe a randomly initialised copy of the study transformer. + + The module is never trained, so any score above the raw-token reference is + attributable to the architecture rather than to a learning objective. + """ + + streams = named_seed_sequences(seed, _STREAM_KEYS) + config = dict(env_config) + config["diagnostics"] = { + "state": True, + "belief": True, + "tokens": True, + "transitions": True, + } + torch.manual_seed(seed_sequence_to_int(streams["untrained_fit"], bits=64)) + module = _build_untrained_module(config, model_config) + + def make_environment(): + return HMMEnv(config) + + environment = make_environment() + try: + initial_belief, outcome_operator, initial_operator = make_transducer_target( + environment + ) + finally: + environment.close() + + common = { + "module": module, + "env_factory": make_environment, + "policy_mode": "greedy", + "device": device, + "warmup": WARMUP, + "initial_belief": initial_belief, + "action_outcome_operator": outcome_operator, + "initial_outcome_operator": initial_operator, + } + fit = collect_probe_data( + n_steps=fit_steps, seed=streams["untrained_fit"], **common + ) + test = collect_probe_data( + n_steps=test_steps, seed=streams["untrained_test"], **common + ) + score, _ = _probe_r2(fit.activations, fit.beliefs, test.activations, test.beliefs) + return { + "r_squared": float(score), + "token_accuracy_greedy": float(test.rewards.mean()), + "n_fit": int(len(fit.beliefs)), + "n_test": int(len(test.beliefs)), + } + + +def compute_references( + *, + seed: int, + fit_steps: int, + test_steps: int, + env_config: dict[str, Any] | None = None, + model_config: dict[str, Any] | None = None, + device: str = "cpu", +) -> dict[str, Any]: + """Compute every reference point the study needs to normalise its metrics.""" + + streams = named_seed_sequences(seed, _STREAM_KEYS) + fit = simulate_stream( + n_steps=fit_steps, + seed=seed_sequence_to_int(streams["reference_fit"]), + ) + test = simulate_stream( + n_steps=test_steps, + seed=seed_sequence_to_int(streams["reference_test"]), + ) + raw_r2 = raw_token_window_r2(fit, test) + saturated = max(raw_r2) + low, high = bootstrap_r2_interval( + fit, + test, + context_length=min(8, saturated), + seed=seed_sequence_to_int(streams["bootstrap"]), + ) + accuracy = bayes_accuracy_by_context(test) + references: dict[str, Any] = { + "alpha": ALPHA, + "seed": seed, + "n_fit": fit_steps, + "n_test": test_steps, + "probe_rank": PROBE_RANK, + "raw_token_window_r2": {str(k): v for k, v in raw_r2.items()}, + "belief_r2_floor": raw_r2[saturated], + "belief_r2_floor_context": saturated, + "belief_r2_probe_noise_95ci": [low, high], + "bayes_accuracy_by_context": {str(k): v for k, v in accuracy.items()}, + "accuracy_floor_repeat_previous_token": accuracy[1], + "accuracy_ceiling_bayes": accuracy[max(accuracy)], + } + if env_config is not None and model_config is not None: + references["untrained_module"] = untrained_module_r2( + env_config=env_config, + model_config=model_config, + seed=seed, + fit_steps=fit_steps, + test_steps=test_steps, + device=device, + ) + return references + + +def normalise(value: float, *, floor: float, ceiling: float) -> float: + """Express a score as the fraction of the floor-to-ceiling range it covers.""" + + if ceiling <= floor: + raise ValueError("ceiling must exceed floor") + return (value - floor) / (ceiling - floor) diff --git a/experiments/mess3_token_guess_cycle_2/operating_point.py b/experiments/mess3_token_guess_cycle_2/operating_point.py new file mode 100644 index 00000000..c343fac1 --- /dev/null +++ b/experiments/mess3_token_guess_cycle_2/operating_point.py @@ -0,0 +1,390 @@ +"""Choosing MESS3 transition and emission parameters for metric sensitivity. + +`REVIEW.md` found that at the parameters cycle 1 used (`alpha=0.85`, self-transition +0.9) an affine probe on the raw observations already scores R² = 0.967, leaving +the belief-probe metric only 0.03 to move through. The natural fix is to change +the process so that the belief is a less linear function of the observation +history. + +Two things have to be checked before committing to a new operating point, and +they pull in opposite directions. + +Lowering the floor means slowing the chain, because a fast-mixing chain lets the +belief be approximated by an exponentially weighted average of recent one-hot +observations, which is exactly what an affine probe computes. But slowing the +chain also raises the autocorrelation of any rollout drawn from it, and the probe +is estimated from a rollout. Past some point the metric gains range and loses +precision at the same rate, and nothing is won. + +This module measures both, plus the context length the belief actually requires, +so the operating point can be chosen rather than guessed. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import numpy as np + +from experiments.mess3_token_guess_cycle_1.analysis import PROBE_RANK +from experiments.mess3_token_guess_cycle_2.metric_references import ( + _probe_r2, + TokenStream, +) + +N_CHAINS = 4_000 +# `collect_probe_data` steps this many environments in parallel, so a probe +# rollout is this many trajectories rather than that many independent draws. +PROBE_CHAINS = 16 +BURN_IN = 600 +FLOOR_WINDOWS = (4, 8, 16, 32, 64) +CONTEXT_LENGTHS = (8, 16, 32, 64, 128) +BLOCK = 512 +RESAMPLES = 300 + + +@dataclass(frozen=True, slots=True) +class OperatingPoint: + """One symmetric MESS3 process, described by its two free parameters.""" + + alpha: float + self_transition: float + + @property + def transition_matrix(self) -> np.ndarray: + matrix = np.full((3, 3), (1.0 - self.self_transition) / 2.0) + np.fill_diagonal(matrix, self.self_transition) + return matrix + + @property + def emission_matrix(self) -> np.ndarray: + matrix = np.full((3, 3), (1.0 - self.alpha) / 2.0) + np.fill_diagonal(matrix, self.alpha) + return matrix + + @property + def state_correlation_time(self) -> float: + """1 / (1 - second eigenvalue), the timescale the hidden state persists.""" + + return 1.0 / (1.0 - (3.0 * self.self_transition - 1.0) / 2.0) + + +def simulate_parallel( + point: OperatingPoint, + *, + n_steps: int, + seed: int, + window: int = max(FLOOR_WINDOWS), + n_chains: int = N_CHAINS, +) -> TokenStream: + """Draw from ``n_chains`` chains stepped together, after burning each in. + + The chain count controls how correlated the returned samples are, which is a + property of the collection scheme rather than of the process. Use the default + for estimating a population quantity like the floor, and ``PROBE_CHAINS`` when + estimating how precise a real probe rollout would be. + """ + + transition = point.transition_matrix + emission = point.emission_matrix + stationary = np.full(3, 1.0 / 3.0) + rng = np.random.default_rng(seed) + cumulative_emission = emission.cumsum(1) + cumulative_transition = transition.cumsum(1) + + state = rng.choice(3, size=n_chains, p=stationary) + belief = np.repeat(stationary[None, :], n_chains, axis=0) + history = np.zeros((n_chains, window), dtype=np.int64) + + beliefs, tokens, windows = [], [], [] + per_chain = int(np.ceil(n_steps / n_chains)) + for step in range(BURN_IN + per_chain): + token = (rng.random(n_chains)[:, None] > cumulative_emission[state]).sum(1) + if step >= BURN_IN: + beliefs.append(belief.copy()) + tokens.append(token.copy()) + windows.append(history.copy()) + history = np.concatenate([history[:, 1:], token[:, None]], axis=1) + posterior = belief * emission[:, token].T + belief = (posterior / posterior.sum(1, keepdims=True)) @ transition + state = (rng.random(n_chains)[:, None] > cumulative_transition[state]).sum(1) + # Order by chain so that contiguous slices are contiguous in time, which the + # block bootstrap relies on. + return TokenStream( + beliefs=np.concatenate(beliefs).reshape(-1, n_chains, 3).transpose(1, 0, 2).reshape(-1, 3)[:n_steps], + tokens=np.concatenate(tokens).reshape(-1, n_chains).T.reshape(-1)[:n_steps], + windows=np.concatenate(windows).reshape(-1, n_chains, window).transpose(1, 0, 2).reshape(-1, window)[:n_steps], + ) + + +def _one_hot(windows: np.ndarray, k: int) -> np.ndarray: + return np.eye(3, dtype=np.float64)[windows[:, -k:]].reshape(len(windows), -1) + + +def belief_r2_floor(fit: TokenStream, test: TokenStream) -> float: + """The best affine readout of the raw observations, over window lengths.""" + + return max( + float(_probe_r2(_one_hot(fit.windows, k), fit.beliefs, _one_hot(test.windows, k), test.beliefs)[0]) + for k in FLOOR_WINDOWS + ) + + +def accuracy_bounds(point: OperatingPoint, test: TokenStream) -> tuple[float, float]: + """Bayes-optimal accuracy, and the repeat-the-previous-observation rule.""" + + emission = point.emission_matrix + ceiling = float((( test.beliefs @ emission).argmax(1) == test.tokens).mean()) + belief = np.repeat(np.full(3, 1.0 / 3.0)[None, :], len(test.windows), axis=0) + belief = belief * emission[:, test.windows[:, -1]].T + belief /= belief.sum(1, keepdims=True) + belief = belief @ point.transition_matrix + floor = float(((belief @ emission).argmax(1) == test.tokens).mean()) + return floor, ceiling + + +def context_requirement( + point: OperatingPoint, + test: TokenStream, + *, + context_lengths: tuple[int, ...] = CONTEXT_LENGTHS, +) -> dict[int, float]: + """Best belief R² available to a model limited to ``k`` observations. + + The exact posterior given the last ``k`` observations under a stationary + prior is what a ``k``-context model can compute at best, so this is an + architecture ceiling that no training objective can exceed. + """ + + transition = point.transition_matrix + emission = point.emission_matrix + scores: dict[int, float] = {} + for k in context_lengths: + if k > test.windows.shape[1]: + continue + belief = np.repeat(np.full(3, 1.0 / 3.0)[None, :], len(test.windows), axis=0) + for offset in range(k): + belief = belief * emission[:, test.windows[:, -k + offset]].T + belief /= belief.sum(1, keepdims=True) + belief = belief @ transition + residual = ((belief - test.beliefs) ** 2).sum() + total = ((test.beliefs - test.beliefs.mean(0)) ** 2).sum() + scores[k] = float(1.0 - residual / total) + return scores + + +def receptive_field(*, n_layers: int, context_len: int) -> int: + """Observations the model can actually reach back to. + + ``CausalTransformerEncoder`` applies a causal band of ``context_len`` at every + layer, so stacking them multiplies the reach. The study's ``context_len=64`` + over three layers is a 192-observation receptive field, not 64. + """ + + if n_layers <= 0 or context_len <= 0: + raise ValueError("n_layers and context_len must be positive") + return n_layers * context_len + + +def smallest_sufficient_context_len( + point: OperatingPoint, + test: TokenStream, + *, + n_layers: int = 3, + tolerance: float = 1e-3, + candidates: tuple[int, ...] = (4, 8, 12, 16, 20), +) -> int: + """Smallest ``context_len`` that still reaches the architecture ceiling. + + Compute in both the learner and the rollout path scales with the receptive + field, so anything beyond what the belief needs is paid for and discarded. + + Only candidates whose receptive field fits inside the observation window + recorded on ``test`` can be checked, so widen that window before widening + these candidates. + """ + + testable = [ + context_len + for context_len in sorted(candidates) + if receptive_field(n_layers=n_layers, context_len=context_len) + <= test.windows.shape[1] + ] + if not testable: + raise ValueError("no candidate fits inside the recorded observation window") + for context_len in testable: + field = receptive_field(n_layers=n_layers, context_len=context_len) + if context_requirement(point, test, context_lengths=(field,))[field] >= ( + 1.0 - tolerance + ): + return context_len + return testable[-1] + + +def integrated_autocorrelation(point: OperatingPoint, *, n_steps: int, seed: int) -> float: + """Sum the belief autocorrelation along a single chain. + + Rollouts are drawn as trajectories, so this is what divides the nominal probe + size to give the number of independent samples it really contains. + """ + + transition = point.transition_matrix + emission = point.emission_matrix + cumulative_emission = emission.cumsum(1) + cumulative_transition = transition.cumsum(1) + rng = np.random.default_rng(seed) + stationary = np.full(3, 1.0 / 3.0) + state = int(rng.choice(3, p=stationary)) + belief = stationary.copy() + series = np.empty(n_steps) + for step in range(n_steps + BURN_IN): + token = int((rng.random() > cumulative_emission[state]).sum()) + if step >= BURN_IN: + series[step - BURN_IN] = belief[0] + posterior = belief * emission[:, token] + belief = (posterior / posterior.sum()) @ transition + state = int((rng.random() > cumulative_transition[state]).sum()) + + centred = series - series.mean() + variance = centred.var() + tau = 1.0 + for lag in range(1, min(4_000, n_steps // 4)): + value = float((centred[:-lag] * centred[lag:]).mean() / variance) + if value <= 0.0: + break + tau += 2.0 * value + return tau + + +def block_bootstrap_interval( + fit: TokenStream, + test: TokenStream, + *, + window: int, + seed: int, + block: int = BLOCK, + resamples: int = RESAMPLES, +) -> tuple[float, float]: + """Resample contiguous blocks, which an i.i.d. bootstrap would understate.""" + + _, predicted = _probe_r2( + _one_hot(fit.windows, window), + fit.beliefs, + _one_hot(test.windows, window), + test.beliefs, + ) + rng = np.random.default_rng(seed) + n = len(test.beliefs) + n_blocks = max(1, n // block) + scores = [] + for _ in range(resamples): + starts = rng.integers(0, max(1, n - block), n_blocks) + index = (starts[:, None] + np.arange(block)[None, :]).ravel() + index = index[index < n] + residual = ((predicted[index] - test.beliefs[index]) ** 2).sum() + total = ( + (test.beliefs[index] - test.beliefs[index].mean(0)) ** 2 + ).sum() + scores.append(1.0 - residual / total) + low, high = np.percentile(scores, [2.5, 97.5]) + return float(low), float(high) + + +def replicate_probe_r2( + point: OperatingPoint, + *, + fit_steps: int, + test_steps: int, + replicates: int, + seed: int, + window: int = 32, +) -> np.ndarray: + """Re-run the whole probe on fresh rollouts and return every answer. + + This measures the metric's precision by direct replication rather than by + resampling one rollout, so it needs no assumption about the correlation + structure and can be used to check that a bootstrap is not lying. + """ + + scores = [] + for index in range(replicates): + offset = seed + 1_000 * index + fit = simulate_parallel( + point, n_steps=fit_steps, seed=offset, n_chains=PROBE_CHAINS + ) + test = simulate_parallel( + point, n_steps=test_steps, seed=offset + 500_000, n_chains=PROBE_CHAINS + ) + score, _ = _probe_r2( + _one_hot(fit.windows, window), + fit.beliefs, + _one_hot(test.windows, window), + test.beliefs, + ) + scores.append(float(score)) + return np.array(scores) + + +def probe_steps_for_spread( + measured_spread: float, + *, + measured_steps: int, + target_spread: float, +) -> int: + """Probe steps needed for a target spread, assuming the usual 1/sqrt(n). + + Verified against replication at 30k, 120k and 500k steps, where the observed + spread does fall as the square root of the step count. + """ + + if measured_spread <= 0.0 or target_spread <= 0.0: + raise ValueError("spreads must be positive") + if target_spread >= measured_spread: + return measured_steps + return int(np.ceil(measured_steps * (measured_spread / target_spread) ** 2)) + + +def evaluate( + point: OperatingPoint, + *, + fit_steps: int, + test_steps: int, + seed: int, +) -> dict[str, float | dict[int, float]]: + """Score one candidate operating point on range and on precision.""" + + fit = simulate_parallel(point, n_steps=fit_steps, seed=seed) + test = simulate_parallel(point, n_steps=test_steps, seed=seed + 1) + floor = belief_r2_floor(fit, test) + accuracy_floor, accuracy_ceiling = accuracy_bounds(point, test) + tau = integrated_autocorrelation(point, n_steps=min(120_000, 40 * test_steps), seed=seed + 2) + # Estimate the probe's precision from a rollout collected the way the study + # collects one, not from the near-independent sample used above. + probe_fit = simulate_parallel( + point, n_steps=fit_steps, seed=seed + 4, n_chains=PROBE_CHAINS + ) + probe_test = simulate_parallel( + point, n_steps=test_steps, seed=seed + 5, n_chains=PROBE_CHAINS + ) + low, high = block_bootstrap_interval(probe_fit, probe_test, window=32, seed=seed + 3) + smallest_context = smallest_sufficient_context_len(point, test) + return { + "alpha": point.alpha, + "self_transition": point.self_transition, + "state_correlation_time": point.state_correlation_time, + "belief_r2_floor": floor, + "belief_r2_range": 1.0 - floor, + "accuracy_floor": accuracy_floor, + "accuracy_ceiling": accuracy_ceiling, + "accuracy_range": accuracy_ceiling - accuracy_floor, + "context_requirement": context_requirement(point, test), + "smallest_sufficient_context_len": smallest_context, + "smallest_sufficient_receptive_field": receptive_field( + n_layers=3, context_len=smallest_context + ), + "integrated_autocorrelation": tau, + "effective_sample_size": test_steps / tau, + "probe_block_bootstrap_ci": [low, high], + "probe_ci_half_width": (high - low) / 2.0, + "probe_rank": PROBE_RANK, + } diff --git a/experiments/mess3_token_guess_cycle_2/references/__init__.py b/experiments/mess3_token_guess_cycle_2/references/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/experiments/mess3_token_guess_cycle_2/references/experiment.py b/experiments/mess3_token_guess_cycle_2/references/experiment.py new file mode 100644 index 00000000..05061b5b --- /dev/null +++ b/experiments/mess3_token_guess_cycle_2/references/experiment.py @@ -0,0 +1,202 @@ +"""Record the task-intrinsic floors and ceilings for the token-guess metrics.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt # noqa: E402 + +from experiments.mess3_token_guess_cycle_1.comparison.experiment import ( + BASE_MODEL_CONFIG, + ENV_CONFIG, +) +from experiments.mess3_token_guess_cycle_2.metric_references import ( + compute_references, + normalise, +) +from harness.artifacts import RunArtifacts +from harness.context import RunContext + +FIT_STEPS = 60_000 +TEST_STEPS = 30_000 +SMOKE_FIT_STEPS = 2_000 +SMOKE_TEST_STEPS = 1_000 + +# Best reported cycle-1 scores, used only to illustrate where the published +# numbers sit inside the range the metric can move through. +CYCLE_1_SCORES = { + "comparison/reward_only": 0.8552, + "comparison/predictive_loss": 0.9319, + "comparison/max_entropy": 0.8558, + "iqn_value": 0.9760, + "kelly_cycle_2/correctness_iqn": 0.9857, + "kelly_cycle_3/conditional_decoupled_kelly_iqn": 0.9824, +} +# Supervised next-token replication, final-LayerNorm probe, seed 0. See +# experiments/mess3_supervised/README.md. +SUPERVISED_CEILING = 0.99888 + + +def plot_landscape(references: dict[str, Any], *, path: Path) -> None: + """Show every published score against the references on one axis.""" + + floor = references["belief_r2_floor"] + untrained = references["untrained_module"]["r_squared"] + rows = [(name, value, "condition") for name, value in CYCLE_1_SCORES.items()] + rows.extend( + [ + ("randomly initialised transformer", untrained, "reference"), + ("affine probe on raw observations", floor, "reference"), + ("supervised next-token model", SUPERVISED_CEILING, "reference"), + ] + ) + rows.sort(key=lambda row: row[1]) + + figure, axis = plt.subplots(figsize=(9.5, 0.46 * len(rows) + 2.0)) + axis.axvspan(floor, SUPERVISED_CEILING, color="tab:green", alpha=0.10) + for position, (name, value, kind) in enumerate(rows): + reference = kind == "reference" + colour = "black" if reference else ("tab:green" if value > floor else "tab:red") + axis.barh( + position, + value - 0.80, + left=0.80, + height=0.55, + color=colour, + alpha=0.35 if reference else 0.85, + hatch="//" if reference else None, + ) + axis.text( + value + 0.002, + position, + f"{value:.4f}", + va="center", + fontsize=8, + ) + axis.set_yticks( + range(len(rows)), + [name.replace("_", " ") for name, _, _ in rows], + fontsize=8, + ) + axis.set_xlim(0.80, 1.03) + axis.set_xlabel("held-out belief-probe R²") + axis.set_title( + "The usable range of belief-probe R² is 0.967 to 0.999\n" + "(hatched bars are references; red bars fall below a probe on the raw " + "observations)", + fontsize=10, + ) + axis.grid(axis="x", alpha=0.2) + figure.tight_layout() + figure.savefig(path, dpi=200) + plt.close(figure) + + +def _findings(references: dict[str, Any]) -> str: + floor = references["belief_r2_floor"] + context = references["belief_r2_floor_context"] + low, high = references["belief_r2_probe_noise_95ci"] + lines = [ + "# Token-guess metric reference points", + "", + "## Belief-probe R²", + "", + f"An affine probe reading the one-hot encoded last {context} observations,", + "with no network and no training, already scores " + f"R² = {floor:.4f}.", + "The supervised next-token replication reaches " + f"{SUPERVISED_CEILING:.4f}.", + "Belief-probe R² therefore moves through a usable range of only " + f"{SUPERVISED_CEILING - floor:.4f}.", + "", + "| observations visible to the probe | R² |", + "|---:|---:|", + ] + for k, value in sorted( + references["raw_token_window_r2"].items(), key=lambda item: int(item[0]) + ): + lines.append(f"| {k} | {value:.4f} |") + untrained = references.get("untrained_module") + if untrained is not None: + lines.extend( + [ + "", + "A randomly initialised copy of the study transformer scores " + f"R² = {untrained['r_squared']:.4f} with greedy accuracy " + f"{untrained['token_accuracy_greedy']:.4f}.", + ] + ) + lines.extend( + [ + "", + "Bootstrap resampling of the probe's test set puts its own sampling " + f"noise at [{low:.4f}, {high:.4f}].", + "", + "## Where the cycle-1 scores sit", + "", + "| condition | reported R² | fraction of the floor-to-ceiling range |", + "|---|---:|---:|", + ] + ) + for condition, value in CYCLE_1_SCORES.items(): + fraction = normalise(value, floor=floor, ceiling=SUPERVISED_CEILING) + lines.append(f"| `{condition}` | {value:.4f} | {fraction:+.1%} |") + + accuracy = references["bayes_accuracy_by_context"] + lines.extend( + [ + "", + "## Greedy token accuracy", + "", + "| observations visible to an exact Bayesian filter | accuracy |", + "|---:|---:|", + ] + ) + for k, value in sorted(accuracy.items(), key=lambda item: int(item[0])): + lines.append(f"| {k} | {value:.4f} |") + lines.extend( + [ + "", + "One observation reproduces the trivial repeat-the-previous-token " + f"rule at {references['accuracy_floor_repeat_previous_token']:.4f}; " + "the filter saturates at " + f"{references['accuracy_ceiling_bayes']:.4f}. Greedy token accuracy " + "therefore moves through a usable range of only " + f"{references['accuracy_ceiling_bayes'] - references['accuracy_floor_repeat_previous_token']:.4f}.", + "", + ] + ) + return "\n".join(lines) + + +def run(context: RunContext) -> dict[str, Any]: + if context.seed is None: + raise ValueError("the reference computation requires a resolved seed") + outputs = RunArtifacts.from_context(context) + outputs.prepare() + references = compute_references( + seed=context.seed, + fit_steps=SMOKE_FIT_STEPS if context.smoke else FIT_STEPS, + test_steps=SMOKE_TEST_STEPS if context.smoke else TEST_STEPS, + env_config=ENV_CONFIG, + model_config=BASE_MODEL_CONFIG, + ) + references["supervised_ceiling"] = SUPERVISED_CEILING + references["cycle_1_normalised"] = { + condition: normalise( + value, + floor=references["belief_r2_floor"], + ceiling=SUPERVISED_CEILING, + ) + for condition, value in CYCLE_1_SCORES.items() + } + figure_path = context.results_dir / "belief_r2_landscape.png" + plot_landscape(references, path=figure_path) + references["figure"] = str(figure_path) + outputs.write_json("references.json", references) + (context.results_dir / "findings.md").write_text(_findings(references)) + return references diff --git a/experiments/mess3_token_guess_cycle_2/references/results/20260725T234936Z-4fc9f18d/belief_r2_landscape.png b/experiments/mess3_token_guess_cycle_2/references/results/20260725T234936Z-4fc9f18d/belief_r2_landscape.png new file mode 100644 index 00000000..d4c097d8 Binary files /dev/null and b/experiments/mess3_token_guess_cycle_2/references/results/20260725T234936Z-4fc9f18d/belief_r2_landscape.png differ diff --git a/experiments/mess3_token_guess_cycle_2/references/results/20260725T234936Z-4fc9f18d/findings.md b/experiments/mess3_token_guess_cycle_2/references/results/20260725T234936Z-4fc9f18d/findings.md new file mode 100644 index 00000000..cffc4761 --- /dev/null +++ b/experiments/mess3_token_guess_cycle_2/references/results/20260725T234936Z-4fc9f18d/findings.md @@ -0,0 +1,49 @@ +# Token-guess metric reference points + +## Belief-probe R² + +An affine probe reading the one-hot encoded last 64 observations, +with no network and no training, already scores R² = 0.9668. +The supervised next-token replication reaches 0.9989. +Belief-probe R² therefore moves through a usable range of only 0.0321. + +| observations visible to the probe | R² | +|---:|---:| +| 1 | 0.8043 | +| 2 | 0.9302 | +| 3 | 0.9581 | +| 4 | 0.9648 | +| 6 | 0.9667 | +| 8 | 0.9669 | +| 16 | 0.9669 | +| 64 | 0.9668 | + +A randomly initialised copy of the study transformer scores R² = 0.8733 with greedy accuracy 0.3412. + +Bootstrap resampling of the probe's test set puts its own sampling noise at [0.9664, 0.9673]. + +## Where the cycle-1 scores sit + +| condition | reported R² | fraction of the floor-to-ceiling range | +|---|---:|---:| +| `comparison/reward_only` | 0.8552 | -347.6% | +| `comparison/predictive_loss` | 0.9319 | -108.6% | +| `comparison/max_entropy` | 0.8558 | -345.7% | +| `iqn_value` | 0.9760 | +28.7% | +| `kelly_cycle_2/correctness_iqn` | 0.9857 | +58.9% | +| `kelly_cycle_3/conditional_decoupled_kelly_iqn` | 0.9824 | +48.7% | + +## Greedy token accuracy + +| observations visible to an exact Bayesian filter | accuracy | +|---:|---:| +| 1 | 0.6732 | +| 2 | 0.6732 | +| 3 | 0.6859 | +| 4 | 0.6860 | +| 6 | 0.6879 | +| 8 | 0.6884 | +| 16 | 0.6883 | +| 64 | 0.6883 | + +One observation reproduces the trivial repeat-the-previous-token rule at 0.6732; the filter saturates at 0.6883. Greedy token accuracy therefore moves through a usable range of only 0.0151. diff --git a/experiments/mess3_token_guess_cycle_2/references/results/20260725T234936Z-4fc9f18d/references.json b/experiments/mess3_token_guess_cycle_2/references/results/20260725T234936Z-4fc9f18d/references.json new file mode 100644 index 00000000..a03c16ae --- /dev/null +++ b/experiments/mess3_token_guess_cycle_2/references/results/20260725T234936Z-4fc9f18d/references.json @@ -0,0 +1,51 @@ +{ + "accuracy_ceiling_bayes": 0.6883, + "accuracy_floor_repeat_previous_token": 0.6732, + "alpha": 0.85, + "bayes_accuracy_by_context": { + "1": 0.6732, + "16": 0.6883, + "2": 0.6732, + "3": 0.6859, + "4": 0.686, + "6": 0.6879333333333333, + "64": 0.6883, + "8": 0.6883666666666667 + }, + "belief_r2_floor": 0.9667781106841088, + "belief_r2_floor_context": 64, + "belief_r2_probe_noise_95ci": [ + 0.966386807711707, + 0.9672839724278931 + ], + "cycle_1_normalised": { + "comparison/max_entropy": -3.4570585423198747, + "comparison/predictive_loss": -1.0864815569232977, + "comparison/reward_only": -3.4757490310352237, + "iqn_value": 0.287269363654747, + "kelly_cycle_2/correctness_iqn": 0.5894322645528668, + "kelly_cycle_3/conditional_decoupled_kelly_iqn": 0.48663457661845627 + }, + "figure": "/workspace/experiments/mess3_token_guess_cycle_2/references/results/20260725T234936Z-4fc9f18d/belief_r2_landscape.png", + "n_fit": 60000, + "n_test": 30000, + "probe_rank": 2, + "raw_token_window_r2": { + "1": 0.8043424580022087, + "16": 0.966855218046489, + "2": 0.9301666406047294, + "3": 0.9580682184473358, + "4": 0.9647961628967233, + "6": 0.9667392183948693, + "64": 0.9667781106841088, + "8": 0.9668577324292733 + }, + "seed": 42, + "supervised_ceiling": 0.99888, + "untrained_module": { + "n_fit": 60000, + "n_test": 30000, + "r_squared": 0.8733171203578343, + "token_accuracy_greedy": 0.34123333333333333 + } +} diff --git a/experiments/mess3_token_guess_cycle_2/references/results/20260725T234936Z-4fc9f18d/run_manifest.json b/experiments/mess3_token_guess_cycle_2/references/results/20260725T234936Z-4fc9f18d/run_manifest.json new file mode 100644 index 00000000..0ca95f37 --- /dev/null +++ b/experiments/mess3_token_guess_cycle_2/references/results/20260725T234936Z-4fc9f18d/run_manifest.json @@ -0,0 +1,75 @@ +{ + "command": [ + "/workspace/.venv/bin/rl-harness", + "experiments.mess3_token_guess_cycle_2.references.experiment", + "--no-upload-artifacts" + ], + "dependency_lock": { + "file": "/workspace/uv.lock", + "sha256": "986725e46af5252423583fe947d30b2646f541a4863acb8c51d578a7ca1b8220" + }, + "ended_at": "2026-07-25T23:49:59.308138+00:00", + "error": null, + "experiment": { + "file": "/workspace/experiments/mess3_token_guess_cycle_2/references/experiment.py", + "module": "experiments.mess3_token_guess_cycle_2.references.experiment", + "source_sha256": "606fa5cdba5b1aa2eefe3945514a955e43d0787844c83cbdd2e30906481c927e" + }, + "framework_versions": { + "gymnasium": "1.2.2", + "numpy": "2.5.1", + "python": "3.14.6", + "ray": "2.56.0", + "rl-harness": "0.1.0", + "torch": "2.12.1" + }, + "git": { + "commit": "9015518e4fd44bbfef57c8b938bf4be828147be4", + "dirty": true, + "experiment_repository": { + "commit": "9015518e4fd44bbfef57c8b938bf4be828147be4", + "dirty": true, + "root": "/workspace" + }, + "library": { + "commit": "7c231655be83affb51c1d24d8cfd475b5d425eb5", + "dirty": false, + "package": "rl-harness", + "root": "/rl-harness-src", + "version": "0.1.0" + } + }, + "hardware": { + "cpu_count": 4, + "cuda_available": false, + "machine": "x86_64", + "mps_available": false, + "platform": "Linux-6.12.94+-x86_64-with-glibc2.39", + "profile": { + "learner_device": "cpu", + "name": "cpu", + "num_env_runners": null, + "num_envs_per_env_runner": 24, + "num_gpus_per_env_runner": 0.0, + "torch_threads": null + } + }, + "run_id": "20260725T234936Z-4fc9f18d", + "runtime": { + "artifacts_dir": "/workspace/experiments/mess3_token_guess_cycle_2/references/artifacts/20260725T234936Z-4fc9f18d", + "overrides": { + "hardware_profile": "auto", + "resume_from": null, + "seed": 42, + "smoke": false, + "upload_artifacts": false + }, + "results_dir": "/workspace/experiments/mess3_token_guess_cycle_2/references/results/20260725T234936Z-4fc9f18d", + "resume_from": null, + "seed": 42, + "smoke": false + }, + "schema_version": 2, + "started_at": "2026-07-25T23:49:36.830424+00:00", + "status": "completed" +} diff --git a/experiments/mess3_token_guess_cycle_2/statistics.py b/experiments/mess3_token_guess_cycle_2/statistics.py new file mode 100644 index 00000000..cd84a2a8 --- /dev/null +++ b/experiments/mess3_token_guess_cycle_2/statistics.py @@ -0,0 +1,204 @@ +"""Seed-level aggregation, intervals, and power for the token-guess study. + +Cycle 1 and the Kelly cycles reported ``mean ± std`` over three seeds. Two +things go wrong with that. The ``±`` reads as an interval but is a population +standard deviation, which for three samples sits about 18% below the sample +standard deviation and roughly 3x below a 95% confidence interval on the mean. +And no comparison was checked for whether three seeds could resolve it, so +orderings were reported for gaps far smaller than the noise around them. + +This module supplies the aggregation the study should use instead: a sample +standard deviation, a t-based confidence interval on the mean, a paired +comparison between conditions that share seeds, and the seed count a given +comparison would need before it is worth reporting an ordering. +""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from dataclasses import dataclass + +import numpy as np +from scipy import stats + +CONFIDENCE = 0.95 +POWER = 0.80 + + +@dataclass(frozen=True, slots=True) +class Estimate: + """A condition's score summarised across seeds.""" + + mean: float + sample_sd: float + n: int + ci_low: float + ci_high: float + + @property + def half_width(self) -> float: + return (self.ci_high - self.ci_low) / 2.0 + + +@dataclass(frozen=True, slots=True) +class Comparison: + """A paired difference between two conditions evaluated on shared seeds.""" + + difference: float + sample_sd: float + n: int + ci_low: float + ci_high: float + p_value: float + seeds_for_power: int + + @property + def resolved(self) -> bool: + """Whether the interval excludes zero at the configured confidence.""" + + return self.ci_low > 0.0 or self.ci_high < 0.0 + + +def summarise(values: Sequence[float], *, confidence: float = CONFIDENCE) -> Estimate: + """Summarise one condition with a t-based interval on its mean.""" + + array = np.asarray(values, dtype=np.float64) + if array.ndim != 1 or len(array) < 2: + raise ValueError("summarising a condition needs at least two seeds") + n = len(array) + mean = float(array.mean()) + sample_sd = float(array.std(ddof=1)) + half = float(stats.t.ppf(0.5 + confidence / 2.0, n - 1)) * sample_sd / np.sqrt(n) + return Estimate( + mean=mean, + sample_sd=sample_sd, + n=n, + ci_low=mean - half, + ci_high=mean + half, + ) + + +MAX_SEEDS = 10_000 + + +def t_test_power( + n: int, + difference: float, + sample_sd: float, + *, + confidence: float = CONFIDENCE, + paired: bool = True, +) -> float: + """Power of a two-sided t-test at ``n`` seeds per condition.""" + + if n < 2 or sample_sd <= 0.0: + raise ValueError("power needs at least two seeds and positive spread") + degrees = n - 1 if paired else 2 * (n - 1) + scale = np.sqrt(n) if paired else np.sqrt(n / 2.0) + ncp = abs(difference) / sample_sd * scale + critical = float(stats.t.ppf(0.5 + confidence / 2.0, degrees)) + return float( + stats.nct.sf(critical, degrees, ncp) + stats.nct.cdf(-critical, degrees, ncp) + ) + + +def seeds_for_power( + difference: float, + sample_sd: float, + *, + power: float = POWER, + confidence: float = CONFIDENCE, + paired: bool = True, +) -> int: + """Seeds per condition needed to resolve a difference of this size. + + Solved against the noncentral t distribution rather than the usual normal + approximation. At the seed counts this study runs, the two disagree sharply: + a two-sided test on three seeds has only two degrees of freedom and a + critical value above four, so the normal approximation can suggest three + seeds suffice for a comparison that three seeds cannot in fact resolve. + """ + + if sample_sd < 0.0: + raise ValueError("sample_sd must be non-negative") + if difference == 0.0: + return MAX_SEEDS + if sample_sd == 0.0: + return 2 + for n in range(2, MAX_SEEDS): + if t_test_power( + n, + difference, + sample_sd, + confidence=confidence, + paired=paired, + ) >= power: + return n + return MAX_SEEDS + + +def compare( + left: Mapping[int, float], + right: Mapping[int, float], + *, + confidence: float = CONFIDENCE, + power: float = POWER, +) -> Comparison: + """Compare two conditions on the seeds they share. + + Pairing costs nothing when conditions are run on a common seed list and + removes any variation the seed induces in both conditions at once. It only + helps to the extent that such shared variation exists; the returned + ``sample_sd`` is of the per-seed differences, so it reports directly whether + it did. + """ + + shared = sorted(set(left) & set(right)) + if len(shared) < 2: + raise ValueError("a paired comparison needs at least two shared seeds") + differences = np.array( + [left[seed] - right[seed] for seed in shared], dtype=np.float64 + ) + n = len(differences) + mean = float(differences.mean()) + sample_sd = float(differences.std(ddof=1)) + half = float(stats.t.ppf(0.5 + confidence / 2.0, n - 1)) * sample_sd / np.sqrt(n) + statistic = stats.ttest_rel( + [left[seed] for seed in shared], + [right[seed] for seed in shared], + ) + return Comparison( + difference=mean, + sample_sd=sample_sd, + n=n, + ci_low=mean - half, + ci_high=mean + half, + p_value=float(statistic.pvalue), + seeds_for_power=seeds_for_power( + abs(mean), + sample_sd, + power=power, + confidence=confidence, + paired=True, + ), + ) + + +def holm_adjust(p_values: Mapping[str, float]) -> dict[str, float]: + """Holm-Bonferroni adjustment over a pre-registered comparison family. + + A study that reports every pairwise ordering across eight conditions is + running 28 tests, and at three seeds several will cross any uncorrected + threshold by chance. Declare the family up front and adjust within it. + """ + + if not p_values: + return {} + ordered = sorted(p_values.items(), key=lambda item: item[1]) + m = len(ordered) + adjusted: dict[str, float] = {} + running = 0.0 + for index, (name, value) in enumerate(ordered): + running = max(running, min(1.0, (m - index) * value)) + adjusted[name] = running + return adjusted diff --git a/experiments/mess3_token_guess_cycle_2/task_parameters/__init__.py b/experiments/mess3_token_guess_cycle_2/task_parameters/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/experiments/mess3_token_guess_cycle_2/task_parameters/experiment.py b/experiments/mess3_token_guess_cycle_2/task_parameters/experiment.py new file mode 100644 index 00000000..8aed5d99 --- /dev/null +++ b/experiments/mess3_token_guess_cycle_2/task_parameters/experiment.py @@ -0,0 +1,226 @@ +"""Choose the MESS3 operating point for cycle 2.""" + +from __future__ import annotations + +from typing import Any + +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt # noqa: E402 + +from experiments.mess3_token_guess_cycle_2.operating_point import ( + OperatingPoint, + accuracy_bounds, + belief_r2_floor, + evaluate, + simulate_parallel, +) +from harness.artifacts import RunArtifacts +from harness.context import RunContext + +FIT_STEPS = 60_000 +TEST_STEPS = 30_000 +SMOKE_FIT_STEPS = 8_000 +SMOKE_TEST_STEPS = 4_000 + +GRID_ALPHAS = (0.5, 0.6, 0.7, 0.85, 0.95) +GRID_TRANSITIONS = (0.9, 0.97, 0.99, 0.995) +SMOKE_ALPHAS = (0.7, 0.85) +SMOKE_TRANSITIONS = (0.9, 0.995) + +CANDIDATES = { + "cycle_1": OperatingPoint(alpha=0.85, self_transition=0.90), + "candidate_a": OperatingPoint(alpha=0.70, self_transition=0.99), + "candidate_b": OperatingPoint(alpha=0.60, self_transition=0.995), + "candidate_c": OperatingPoint(alpha=0.70, self_transition=0.995), +} + +# Measured by training the study architecture (3 layers, d_model 96, context 64) +# on next-token prediction to the Bayes cross-entropy, then probing it. Recorded +# rather than recomputed because it takes about twenty minutes per point on CPU. +SUPERVISED_PROBE_R2 = { + "cycle_1": {1_500: 0.9567, 3_000: 0.9532, 4_500: 0.9434, 6_000: 0.9318}, + "candidate_c": {1_500: 0.9487, 3_000: 0.9478, 4_500: 0.9439, 6_000: 0.9367}, +} + + +def _plot_grid(grid: dict[str, Any], *, path) -> None: + alphas = sorted({entry["alpha"] for entry in grid.values()}) + transitions = sorted({entry["self_transition"] for entry in grid.values()}) + figure, axes = plt.subplots(1, 2, figsize=(11.5, 4.4)) + for axis, key, title in ( + (axes[0], "belief_r2_range", "belief-probe R² range (1 − floor)"), + (axes[1], "accuracy_range", "greedy accuracy range (Bayes − repeat-last)"), + ): + for transition in transitions: + values = [ + grid[f"{alpha}_{transition}"][key] + for alpha in alphas + if f"{alpha}_{transition}" in grid + ] + axis.plot(alphas[: len(values)], values, "o-", label=f"p = {transition}") + axis.set_xlabel("emission concentration α") + axis.set_title(title, fontsize=10) + axis.grid(alpha=0.2) + axis.legend(fontsize=8) + axes[0].set_ylabel("usable range of the metric") + figure.suptitle( + "Both metrics are near-degenerate at the parameters cycle 1 used " + "(α = 0.85, p = 0.9)", + fontsize=11, + ) + figure.tight_layout() + figure.savefig(path, dpi=200) + plt.close(figure) + + +def _findings(result: dict[str, Any]) -> str: + lines = [ + "# Choosing the MESS3 operating point", + "", + "## Candidates", + "", + "`range` is what the belief-probe metric can move through; `ESS` is how " + "many independent samples a 30,000-step probe rollout actually contains, " + "given how slowly the chain mixes.", + "", + "| point | α | p | τ | R² floor | R² range | acc range | ESS | probe ±95% |", + "|---|---:|---:|---:|---:|---:|---:|---:|---:|", + ] + for name, entry in result["candidates"].items(): + lines.append( + f"| `{name}` | {entry['alpha']:.2f} | {entry['self_transition']:.3f} | " + f"{entry['state_correlation_time']:.0f} | " + f"{entry['belief_r2_floor']:.4f} | {entry['belief_r2_range']:.4f} | " + f"{entry['accuracy_range']:.4f} | " + f"{entry['effective_sample_size']:.0f} | " + f"{entry['probe_ci_half_width']:.4f} |" + ) + lines.extend( + [ + "", + "## Context length", + "", + "Best belief R² available to a model that can see only the last k " + "observations. No objective can beat this.", + "", + "| point | " + " | ".join( + f"k={k}" + for k in sorted( + next(iter(result["candidates"].values()))["context_requirement"] + ) + ) + " |", + "|---" * (1 + len(next(iter(result["candidates"].values()))["context_requirement"])) + "|", + ] + ) + for name, entry in result["candidates"].items(): + row = " | ".join( + f"{value:.4f}" for _, value in sorted(entry["context_requirement"].items()) + ) + lines.append(f"| `{name}` | {row} |") + lines.extend( + [ + "", + "`CausalTransformerEncoder` applies its causal band at every layer, so " + "the receptive field is `n_layers * context_len`. The study's " + "`context_len=64` over three layers reaches 192 observations, roughly " + "six times what the belief needs. Compute in both the learner and the " + "cached rollout path scales with that reach.", + "", + "| point | smallest sufficient context_len | receptive field |", + "|---|---:|---:|", + ] + + [ + f"| `{name}` | {entry['smallest_sufficient_context_len']} | " + f"{entry['smallest_sufficient_receptive_field']} |" + for name, entry in result["candidates"].items() + ] + + [ + "", + "Measured on this repository's model: dropping `context_len` from 64 to " + "16 makes a learner step 3.7x faster and a cached rollout step 3.4x " + "faster, for an end-to-end PPO speed-up of 1.6x once environment " + "stepping, which does not change, is included.", + "", + "## Supervised ceiling", + "", + "Training the study architecture on next-token prediction to the Bayes " + "cross-entropy, then probing it:", + "", + "| point | 1.5k steps | 3k | 4.5k | 6k |", + "|---|---:|---:|---:|---:|", + ] + ) + for name, curve in SUPERVISED_PROBE_R2.items(): + row = " | ".join(f"{curve[step]:.4f}" for step in sorted(curve)) + lines.append(f"| `{name}` | {row} |") + lines.extend( + [ + "", + "Belief-probe R² falls with continued training at both points while " + "cross-entropy stays at the Bayes floor. This is supervised training, " + "so the decline cycle 1 saw over 20M PPO steps is not caused by " + "reinforcement learning. It is optimiser-driven drift in a " + "representation the task no longer constrains, which makes learning " + "rate, optimiser, and training duration larger influences on the " + "headline metric than most of the differences between arms.", + "", + ] + ) + return "\n".join(lines) + + +def run(context: RunContext) -> dict[str, Any]: + if context.seed is None: + raise ValueError("choosing an operating point requires a resolved seed") + outputs = RunArtifacts.from_context(context) + outputs.prepare() + fit_steps = SMOKE_FIT_STEPS if context.smoke else FIT_STEPS + test_steps = SMOKE_TEST_STEPS if context.smoke else TEST_STEPS + alphas = SMOKE_ALPHAS if context.smoke else GRID_ALPHAS + transitions = SMOKE_TRANSITIONS if context.smoke else GRID_TRANSITIONS + + grid: dict[str, Any] = {} + for alpha in alphas: + for transition in transitions: + point = OperatingPoint(alpha=alpha, self_transition=transition) + fit = simulate_parallel(point, n_steps=fit_steps, seed=context.seed) + test = simulate_parallel(point, n_steps=test_steps, seed=context.seed + 1) + floor = belief_r2_floor(fit, test) + accuracy_floor, accuracy_ceiling = accuracy_bounds(point, test) + grid[f"{alpha}_{transition}"] = { + "alpha": alpha, + "self_transition": transition, + "belief_r2_floor": floor, + "belief_r2_range": 1.0 - floor, + "accuracy_floor": accuracy_floor, + "accuracy_ceiling": accuracy_ceiling, + "accuracy_range": accuracy_ceiling - accuracy_floor, + } + + candidates = { + name: evaluate( + point, + fit_steps=fit_steps, + test_steps=test_steps, + seed=context.seed, + ) + for name, point in CANDIDATES.items() + } + figure_path = context.results_dir / "operating_point_grid.png" + _plot_grid(grid, path=figure_path) + result = { + "seed": context.seed, + "smoke": context.smoke, + "grid": grid, + "candidates": candidates, + "supervised_probe_r2": { + name: {str(step): value for step, value in curve.items()} + for name, curve in SUPERVISED_PROBE_R2.items() + }, + "figure": str(figure_path), + } + outputs.write_json("task_parameters.json", result) + (context.results_dir / "findings.md").write_text(_findings(result)) + return result diff --git a/experiments/mess3_token_guess_cycle_2/task_parameters/results/20260727T065519Z-ef24bfec/findings.md b/experiments/mess3_token_guess_cycle_2/task_parameters/results/20260727T065519Z-ef24bfec/findings.md new file mode 100644 index 00000000..06436553 --- /dev/null +++ b/experiments/mess3_token_guess_cycle_2/task_parameters/results/20260727T065519Z-ef24bfec/findings.md @@ -0,0 +1,45 @@ +# Choosing the MESS3 operating point + +## Candidates + +`range` is what the belief-probe metric can move through; `ESS` is how many independent samples a 30,000-step probe rollout actually contains, given how slowly the chain mixes. + +| point | α | p | τ | R² floor | R² range | acc range | ESS | probe ±95% | +|---|---:|---:|---:|---:|---:|---:|---:|---:| +| `cycle_1` | 0.85 | 0.900 | 7 | 0.9671 | 0.0329 | 0.0151 | 2324 | 0.0006 | +| `candidate_a` | 0.70 | 0.990 | 67 | 0.8967 | 0.1033 | 0.1337 | 198 | 0.0037 | +| `candidate_b` | 0.60 | 0.995 | 133 | 0.8782 | 0.1218 | 0.1362 | 104 | 0.0061 | +| `candidate_c` | 0.70 | 0.995 | 133 | 0.9061 | 0.0939 | 0.1447 | 104 | 0.0046 | + +## Context length + +Best belief R² available to a model that can see only the last k observations. No objective can beat this. + +| point | k=8 | k=16 | k=32 | k=64 | +|---|---|---|---|---| +| `cycle_1` | 1.0000 | 1.0000 | 1.0000 | 1.0000 | +| `candidate_a` | 0.9418 | 0.9969 | 1.0000 | 1.0000 | +| `candidate_b` | 0.7623 | 0.9431 | 0.9965 | 1.0000 | +| `candidate_c` | 0.9232 | 0.9951 | 1.0000 | 1.0000 | + +`CausalTransformerEncoder` applies its causal band at every layer, so the receptive field is `n_layers * context_len`. The study's `context_len=64` over three layers reaches 192 observations, roughly six times what the belief needs. Compute in both the learner and the cached rollout path scales with that reach. + +| point | smallest sufficient context_len | receptive field | +|---|---:|---:| +| `cycle_1` | 4 | 12 | +| `candidate_a` | 8 | 24 | +| `candidate_b` | 16 | 48 | +| `candidate_c` | 8 | 24 | + +Measured on this repository's model: dropping `context_len` from 64 to 16 makes a learner step 3.7x faster and a cached rollout step 3.4x faster, for an end-to-end PPO speed-up of 1.6x once environment stepping, which does not change, is included. + +## Supervised ceiling + +Training the study architecture on next-token prediction to the Bayes cross-entropy, then probing it: + +| point | 1.5k steps | 3k | 4.5k | 6k | +|---|---:|---:|---:|---:| +| `cycle_1` | 0.9567 | 0.9532 | 0.9434 | 0.9318 | +| `candidate_c` | 0.9487 | 0.9478 | 0.9439 | 0.9367 | + +Belief-probe R² falls with continued training at both points while cross-entropy stays at the Bayes floor. This is supervised training, so the decline cycle 1 saw over 20M PPO steps is not caused by reinforcement learning. It is optimiser-driven drift in a representation the task no longer constrains, which makes learning rate, optimiser, and training duration larger influences on the headline metric than most of the differences between arms. diff --git a/experiments/mess3_token_guess_cycle_2/task_parameters/results/20260727T065519Z-ef24bfec/operating_point_grid.png b/experiments/mess3_token_guess_cycle_2/task_parameters/results/20260727T065519Z-ef24bfec/operating_point_grid.png new file mode 100644 index 00000000..9addad8b Binary files /dev/null and b/experiments/mess3_token_guess_cycle_2/task_parameters/results/20260727T065519Z-ef24bfec/operating_point_grid.png differ diff --git a/experiments/mess3_token_guess_cycle_2/task_parameters/results/20260727T065519Z-ef24bfec/run_manifest.json b/experiments/mess3_token_guess_cycle_2/task_parameters/results/20260727T065519Z-ef24bfec/run_manifest.json new file mode 100644 index 00000000..eb0a4aad --- /dev/null +++ b/experiments/mess3_token_guess_cycle_2/task_parameters/results/20260727T065519Z-ef24bfec/run_manifest.json @@ -0,0 +1,75 @@ +{ + "command": [ + "/workspace/.venv/bin/rl-harness", + "experiments.mess3_token_guess_cycle_2.task_parameters.experiment", + "--no-upload-artifacts" + ], + "dependency_lock": { + "file": "/workspace/uv.lock", + "sha256": "986725e46af5252423583fe947d30b2646f541a4863acb8c51d578a7ca1b8220" + }, + "ended_at": "2026-07-27T06:55:55.380651+00:00", + "error": null, + "experiment": { + "file": "/workspace/experiments/mess3_token_guess_cycle_2/task_parameters/experiment.py", + "module": "experiments.mess3_token_guess_cycle_2.task_parameters.experiment", + "source_sha256": "0de4d683e821de377b3cfa9bd0036b37ea6217262f0a43632d152204aac507b3" + }, + "framework_versions": { + "gymnasium": "1.2.2", + "numpy": "2.5.1", + "python": "3.14.6", + "ray": "2.56.0", + "rl-harness": "0.1.0", + "torch": "2.12.1" + }, + "git": { + "commit": "d790369e770a98ac118709c8bed484375415f998", + "dirty": true, + "experiment_repository": { + "commit": "d790369e770a98ac118709c8bed484375415f998", + "dirty": true, + "root": "/workspace" + }, + "library": { + "commit": "7c231655be83affb51c1d24d8cfd475b5d425eb5", + "dirty": false, + "package": "rl-harness", + "root": "/rl-harness-src", + "version": "0.1.0" + } + }, + "hardware": { + "cpu_count": 4, + "cuda_available": false, + "machine": "x86_64", + "mps_available": false, + "platform": "Linux-6.12.94+-x86_64-with-glibc2.39", + "profile": { + "learner_device": "cpu", + "name": "cpu", + "num_env_runners": null, + "num_envs_per_env_runner": 24, + "num_gpus_per_env_runner": 0.0, + "torch_threads": null + } + }, + "run_id": "20260727T065519Z-ef24bfec", + "runtime": { + "artifacts_dir": "/workspace/experiments/mess3_token_guess_cycle_2/task_parameters/artifacts/20260727T065519Z-ef24bfec", + "overrides": { + "hardware_profile": "auto", + "resume_from": null, + "seed": 42, + "smoke": false, + "upload_artifacts": false + }, + "results_dir": "/workspace/experiments/mess3_token_guess_cycle_2/task_parameters/results/20260727T065519Z-ef24bfec", + "resume_from": null, + "seed": 42, + "smoke": false + }, + "schema_version": 2, + "started_at": "2026-07-27T06:55:19.583001+00:00", + "status": "completed" +} diff --git a/experiments/mess3_token_guess_cycle_2/task_parameters/results/20260727T065519Z-ef24bfec/task_parameters.json b/experiments/mess3_token_guess_cycle_2/task_parameters/results/20260727T065519Z-ef24bfec/task_parameters.json new file mode 100644 index 00000000..9f735818 --- /dev/null +++ b/experiments/mess3_token_guess_cycle_2/task_parameters/results/20260727T065519Z-ef24bfec/task_parameters.json @@ -0,0 +1,307 @@ +{ + "candidates": { + "candidate_a": { + "accuracy_ceiling": 0.6639666666666667, + "accuracy_floor": 0.5302666666666667, + "accuracy_range": 0.13370000000000004, + "alpha": 0.7, + "belief_r2_floor": 0.8967379018599352, + "belief_r2_range": 0.10326209814006482, + "context_requirement": { + "16": 0.9968544824392702, + "32": 0.9999964397062356, + "64": 0.9999999999946176, + "8": 0.9417756254662745 + }, + "effective_sample_size": 197.8355422568187, + "integrated_autocorrelation": 151.6411038065937, + "probe_block_bootstrap_ci": [ + 0.8938617556533653, + 0.9012300870289232 + ], + "probe_ci_half_width": 0.0036841656877789064, + "probe_rank": 2, + "self_transition": 0.99, + "smallest_sufficient_context_len": 8, + "smallest_sufficient_receptive_field": 24, + "state_correlation_time": 66.66666666666612 + }, + "candidate_b": { + "accuracy_ceiling": 0.5779333333333333, + "accuracy_floor": 0.4417333333333333, + "accuracy_range": 0.1362, + "alpha": 0.6, + "belief_r2_floor": 0.878248921449702, + "belief_r2_range": 0.12175107855029799, + "context_requirement": { + "16": 0.943060458407508, + "32": 0.9965116891854748, + "64": 0.9999956061509646, + "8": 0.762277706445713 + }, + "effective_sample_size": 103.78188928313486, + "integrated_autocorrelation": 289.06777673082087, + "probe_block_bootstrap_ci": [ + 0.8603267685787276, + 0.8725406679311462 + ], + "probe_ci_half_width": 0.006106949676209283, + "probe_rank": 2, + "self_transition": 0.995, + "smallest_sufficient_context_len": 16, + "smallest_sufficient_receptive_field": 48, + "state_correlation_time": 133.33333333333223 + }, + "candidate_c": { + "accuracy_ceiling": 0.6837666666666666, + "accuracy_floor": 0.5390666666666667, + "accuracy_range": 0.14469999999999994, + "alpha": 0.7, + "belief_r2_floor": 0.906136201034249, + "belief_r2_range": 0.09386379896575103, + "context_requirement": { + "16": 0.995128172467821, + "32": 0.9999892025679498, + "64": 1.0, + "8": 0.923215622292938 + }, + "effective_sample_size": 104.08060835175199, + "integrated_autocorrelation": 288.23813076314525, + "probe_block_bootstrap_ci": [ + 0.9020596893487307, + 0.9111640239783987 + ], + "probe_ci_half_width": 0.004552167314833988, + "probe_rank": 2, + "self_transition": 0.995, + "smallest_sufficient_context_len": 8, + "smallest_sufficient_receptive_field": 24, + "state_correlation_time": 133.33333333333223 + }, + "cycle_1": { + "accuracy_ceiling": 0.6882666666666667, + "accuracy_floor": 0.6731333333333334, + "accuracy_range": 0.015133333333333332, + "alpha": 0.85, + "belief_r2_floor": 0.9671181729028167, + "belief_r2_range": 0.03288182709718335, + "context_requirement": { + "16": 0.9999999999275502, + "32": 1.0, + "64": 1.0, + "8": 0.9999795066409225 + }, + "effective_sample_size": 2324.268140252257, + "integrated_autocorrelation": 12.907288741971072, + "probe_block_bootstrap_ci": [ + 0.9663334399107788, + 0.967472758102929 + ], + "probe_ci_half_width": 0.0005696590960750925, + "probe_rank": 2, + "self_transition": 0.9, + "smallest_sufficient_context_len": 4, + "smallest_sufficient_receptive_field": 12, + "state_correlation_time": 6.6666666666666705 + } + }, + "figure": "/workspace/experiments/mess3_token_guess_cycle_2/task_parameters/results/20260727T065519Z-ef24bfec/operating_point_grid.png", + "grid": { + "0.5_0.9": { + "accuracy_ceiling": 0.3857333333333333, + "accuracy_floor": 0.3646333333333333, + "accuracy_range": 0.021100000000000008, + "alpha": 0.5, + "belief_r2_floor": 0.9533460163985678, + "belief_r2_range": 0.04665398360143225, + "self_transition": 0.9 + }, + "0.5_0.97": { + "accuracy_ceiling": 0.4261333333333333, + "accuracy_floor": 0.3727666666666667, + "accuracy_range": 0.05336666666666662, + "alpha": 0.5, + "belief_r2_floor": 0.8948456505926992, + "belief_r2_range": 0.10515434940730084, + "self_transition": 0.97 + }, + "0.5_0.99": { + "accuracy_ceiling": 0.45253333333333334, + "accuracy_floor": 0.37233333333333335, + "accuracy_range": 0.0802, + "alpha": 0.5, + "belief_r2_floor": 0.8570197263850013, + "belief_r2_range": 0.14298027361499865, + "self_transition": 0.99 + }, + "0.5_0.995": { + "accuracy_ceiling": 0.47273333333333334, + "accuracy_floor": 0.3746333333333333, + "accuracy_range": 0.09810000000000002, + "alpha": 0.5, + "belief_r2_floor": 0.8477653930238022, + "belief_r2_range": 0.15223460697619784, + "self_transition": 0.995 + }, + "0.6_0.9": { + "accuracy_ceiling": 0.4498, + "accuracy_floor": 0.41983333333333334, + "accuracy_range": 0.02996666666666664, + "alpha": 0.6, + "belief_r2_floor": 0.9394608905059447, + "belief_r2_range": 0.060539109494055277, + "self_transition": 0.9 + }, + "0.6_0.97": { + "accuracy_ceiling": 0.5214, + "accuracy_floor": 0.43806666666666666, + "accuracy_range": 0.08333333333333331, + "alpha": 0.6, + "belief_r2_floor": 0.8848229869875442, + "belief_r2_range": 0.11517701301245575, + "self_transition": 0.97 + }, + "0.6_0.99": { + "accuracy_ceiling": 0.5583666666666667, + "accuracy_floor": 0.4363, + "accuracy_range": 0.12206666666666666, + "alpha": 0.6, + "belief_r2_floor": 0.8723894114642732, + "belief_r2_range": 0.12761058853572682, + "self_transition": 0.99 + }, + "0.6_0.995": { + "accuracy_ceiling": 0.5779333333333333, + "accuracy_floor": 0.4417333333333333, + "accuracy_range": 0.1362, + "alpha": 0.6, + "belief_r2_floor": 0.878248921449702, + "belief_r2_range": 0.12175107855029799, + "self_transition": 0.995 + }, + "0.7_0.9": { + "accuracy_ceiling": 0.5384333333333333, + "accuracy_floor": 0.5033333333333333, + "accuracy_range": 0.03510000000000002, + "alpha": 0.7, + "belief_r2_floor": 0.9409244989416722, + "belief_r2_range": 0.05907550105832782, + "self_transition": 0.9 + }, + "0.7_0.97": { + "accuracy_ceiling": 0.6201333333333333, + "accuracy_floor": 0.5280666666666667, + "accuracy_range": 0.09206666666666663, + "alpha": 0.7, + "belief_r2_floor": 0.8988869558203207, + "belief_r2_range": 0.10111304417967926, + "self_transition": 0.97 + }, + "0.7_0.99": { + "accuracy_ceiling": 0.6639666666666667, + "accuracy_floor": 0.5302666666666667, + "accuracy_range": 0.13370000000000004, + "alpha": 0.7, + "belief_r2_floor": 0.8967379018599352, + "belief_r2_range": 0.10326209814006482, + "self_transition": 0.99 + }, + "0.7_0.995": { + "accuracy_ceiling": 0.6837666666666666, + "accuracy_floor": 0.5390666666666667, + "accuracy_range": 0.14469999999999994, + "alpha": 0.7, + "belief_r2_floor": 0.906136201034249, + "belief_r2_range": 0.09386379896575103, + "self_transition": 0.995 + }, + "0.85_0.9": { + "accuracy_ceiling": 0.6882666666666667, + "accuracy_floor": 0.6731333333333334, + "accuracy_range": 0.015133333333333332, + "alpha": 0.85, + "belief_r2_floor": 0.9671181729028167, + "belief_r2_range": 0.03288182709718335, + "self_transition": 0.9 + }, + "0.85_0.97": { + "accuracy_ceiling": 0.7830333333333334, + "accuracy_floor": 0.7164666666666667, + "accuracy_range": 0.06656666666666666, + "alpha": 0.85, + "belief_r2_floor": 0.9378868344834231, + "belief_r2_range": 0.06211316551657686, + "self_transition": 0.97 + }, + "0.85_0.99": { + "accuracy_ceiling": 0.8216, + "accuracy_floor": 0.7285333333333334, + "accuracy_range": 0.09306666666666663, + "alpha": 0.85, + "belief_r2_floor": 0.937784849207427, + "belief_r2_range": 0.062215150792573004, + "self_transition": 0.99 + }, + "0.85_0.995": { + "accuracy_ceiling": 0.8370333333333333, + "accuracy_floor": 0.7351666666666666, + "accuracy_range": 0.10186666666666666, + "alpha": 0.85, + "belief_r2_floor": 0.9442063541758701, + "belief_r2_range": 0.055793645824129934, + "self_transition": 0.995 + }, + "0.95_0.9": { + "accuracy_ceiling": 0.8206, + "accuracy_floor": 0.8206, + "accuracy_range": 0.0, + "alpha": 0.95, + "belief_r2_floor": 0.9948367481576547, + "belief_r2_range": 0.00516325184234534, + "self_transition": 0.9 + }, + "0.95_0.97": { + "accuracy_ceiling": 0.8931333333333333, + "accuracy_floor": 0.8796, + "accuracy_range": 0.013533333333333286, + "alpha": 0.95, + "belief_r2_floor": 0.9810292859935144, + "belief_r2_range": 0.018970714006485645, + "self_transition": 0.97 + }, + "0.95_0.99": { + "accuracy_ceiling": 0.9311666666666667, + "accuracy_floor": 0.8981666666666667, + "accuracy_range": 0.03300000000000003, + "alpha": 0.95, + "belief_r2_floor": 0.974744097401559, + "belief_r2_range": 0.025255902598441016, + "self_transition": 0.99 + }, + "0.95_0.995": { + "accuracy_ceiling": 0.9419, + "accuracy_floor": 0.9030333333333334, + "accuracy_range": 0.038866666666666605, + "alpha": 0.95, + "belief_r2_floor": 0.9756246450262978, + "belief_r2_range": 0.024375354973702246, + "self_transition": 0.995 + } + }, + "seed": 42, + "smoke": false, + "supervised_probe_r2": { + "candidate_c": { + "1500": 0.9487, + "3000": 0.9478, + "4500": 0.9439, + "6000": 0.9367 + }, + "cycle_1": { + "1500": 0.9567, + "3000": 0.9532, + "4500": 0.9434, + "6000": 0.9318 + } + } +} diff --git a/tests/test_mess3_token_guess_audit.py b/tests/test_mess3_token_guess_audit.py new file mode 100644 index 00000000..fe6724db --- /dev/null +++ b/tests/test_mess3_token_guess_audit.py @@ -0,0 +1,68 @@ +"""Tests for the re-analysis of the committed multi-seed token-guess results.""" + +from __future__ import annotations + +from pathlib import Path + +import pytest + +from experiments.mess3_token_guess_cycle_2.audit.experiment import ( + STUDIES, + analyse_study, + collect_study, +) + +REPOSITORY_ROOT = Path(__file__).parents[1] +REFERENCES = { + "belief_r2_floor": 0.9668, + "belief_r2_ceiling": 0.99888, + "accuracy_floor": 0.6732, + "accuracy_ceiling": 0.6883, +} + + +@pytest.mark.parametrize("study", STUDIES) +def test_every_audited_study_committed_three_seeds_per_arm(study): + per_arm = collect_study(REPOSITORY_ROOT, study) + + assert per_arm, f"{study} committed no condition summaries" + for arm, by_seed in per_arm.items(): + assert sorted(by_seed) == [42, 43, 44], arm + + +def test_the_kelly_cycle_3_orderings_do_not_survive_a_holm_correction(): + analysis = analyse_study( + collect_study(REPOSITORY_ROOT, "mess_3_kelly_cycle_3"), + REFERENCES, + ) + comparisons = analysis["comparisons"] + + assert len(comparisons) == 6 + assert not any(values["resolved"] for values in comparisons.values()) + + +def test_only_the_strongest_kelly_cycle_3_arm_clears_the_no_network_floor(): + analysis = analyse_study( + collect_study(REPOSITORY_ROOT, "mess_3_kelly_cycle_3"), + REFERENCES, + ) + conditions = analysis["conditions"] + + assert conditions["conditional_decoupled_kelly_iqn"]["exceeds_no_network_floor"] + for arm in ("ppo", "iqn", "conditional_decoupled_kelly_mean"): + assert not conditions[arm]["exceeds_no_network_floor"], arm + + # Plain PPO sits far below a probe on the raw observations. + assert conditions["ppo"]["fraction_of_usable_range"] < -1.0 + + +def test_every_kelly_cycle_2_kelly_arm_falls_below_the_no_network_floor(): + analysis = analyse_study( + collect_study(REPOSITORY_ROOT, "mess_3_kelly_cycle_2"), + REFERENCES, + ) + conditions = analysis["conditions"] + + for arm, values in conditions.items(): + expected = arm.startswith("correctness") + assert values["exceeds_no_network_floor"] is expected, arm diff --git a/tests/test_mess3_token_guess_operating_point.py b/tests/test_mess3_token_guess_operating_point.py new file mode 100644 index 00000000..ecbcc730 --- /dev/null +++ b/tests/test_mess3_token_guess_operating_point.py @@ -0,0 +1,152 @@ +"""Tests for choosing the MESS3 operating point.""" + +from __future__ import annotations + +import numpy as np +import pytest + +from experiments.mess3_token_guess_cycle_2.operating_point import ( + PROBE_CHAINS, + OperatingPoint, + accuracy_bounds, + belief_r2_floor, + block_bootstrap_interval, + context_requirement, + integrated_autocorrelation, + probe_steps_for_spread, + receptive_field, + replicate_probe_r2, + simulate_parallel, + smallest_sufficient_context_len, +) + +CYCLE_1 = OperatingPoint(alpha=0.85, self_transition=0.90) +SLOW = OperatingPoint(alpha=0.70, self_transition=0.995) + + +def _streams(point, fit=20_000, test=10_000, seed=7): + return ( + simulate_parallel(point, n_steps=fit, seed=seed), + simulate_parallel(point, n_steps=test, seed=seed + 1), + ) + + +def test_operating_point_builds_valid_stochastic_matrices(): + for point in (CYCLE_1, SLOW): + for matrix in (point.transition_matrix, point.emission_matrix): + np.testing.assert_allclose(matrix.sum(axis=1), 1.0) + assert (matrix >= 0.0).all() + assert CYCLE_1.transition_matrix[0, 0] == pytest.approx(0.90) + assert CYCLE_1.emission_matrix[1, 1] == pytest.approx(0.85) + + +def test_slower_chains_persist_for_longer(): + assert SLOW.state_correlation_time > 10 * CYCLE_1.state_correlation_time + assert CYCLE_1.state_correlation_time == pytest.approx(1 / (1 - 0.85), rel=1e-6) + + +def test_slowing_the_chain_lowers_the_no_network_floor(): + cycle_1_floor = belief_r2_floor(*_streams(CYCLE_1)) + slow_floor = belief_r2_floor(*_streams(SLOW)) + + # Cycle 1's parameters leave the metric almost no room. + assert cycle_1_floor > 0.95 + # The slower chain roughly triples what the metric can resolve. + assert (1 - slow_floor) > 2.0 * (1 - cycle_1_floor) + + +def test_slowing_the_chain_also_widens_the_accuracy_range(): + _, cycle_1_test = _streams(CYCLE_1) + _, slow_test = _streams(SLOW) + cycle_1_low, cycle_1_high = accuracy_bounds(CYCLE_1, cycle_1_test) + slow_low, slow_high = accuracy_bounds(SLOW, slow_test) + + assert cycle_1_high - cycle_1_low < 0.03 + assert slow_high - slow_low > 3.0 * (cycle_1_high - cycle_1_low) + # The task has to stay learnable: well above chance. + assert slow_high > 0.5 + + +def test_a_64_observation_context_suffices_at_every_candidate(): + for point in (CYCLE_1, SLOW): + _, test = _streams(point) + scores = context_requirement(point, test, context_lengths=(8, 32, 64)) + assert scores[64] > 0.999, point + assert scores[64] >= scores[32] >= scores[8] + + +def test_receptive_field_multiplies_the_band_by_the_layer_count(): + # CausalTransformerEncoder applies the band at every layer, so the study's + # context_len of 64 over three layers reaches 192 observations. + assert receptive_field(n_layers=3, context_len=64) == 192 + assert receptive_field(n_layers=3, context_len=16) == 48 + assert receptive_field(n_layers=1, context_len=16) == 16 + for bad in ({"n_layers": 0, "context_len": 8}, {"n_layers": 3, "context_len": 0}): + with pytest.raises(ValueError): + receptive_field(**bad) + + +def test_the_study_context_len_is_far_larger_than_the_belief_needs(): + for point in (CYCLE_1, SLOW): + _, test = _streams(point) + smallest = smallest_sufficient_context_len(point, test) + assert smallest <= 16, point + # Well inside the 192-observation reach the study currently pays for. + assert receptive_field(n_layers=3, context_len=smallest) <= 48 + + +def test_smallest_sufficient_context_len_needs_a_wide_enough_window(): + _, test = _streams(SLOW) + narrow = simulate_parallel(SLOW, n_steps=2_000, seed=3, window=8) + assert smallest_sufficient_context_len(SLOW, test) >= 4 + with pytest.raises(ValueError): + smallest_sufficient_context_len(SLOW, narrow, candidates=(16, 20)) + + +def test_replication_confirms_the_slow_chain_is_the_noisier_probe(): + fast = replicate_probe_r2( + CYCLE_1, fit_steps=20_000, test_steps=10_000, replicates=8, seed=1 + ) + slow = replicate_probe_r2( + SLOW, fit_steps=20_000, test_steps=10_000, replicates=8, seed=1 + ) + assert slow.std(ddof=1) > 3.0 * fast.std(ddof=1) + + +def test_probe_steps_for_spread_follows_the_square_root_law(): + # Measured by replication: +/-0.0025 at 30,000 steps on the slow chain. + assert probe_steps_for_spread( + 0.0025, measured_steps=30_000, target_spread=0.0005 + ) == pytest.approx(750_000, rel=0.01) + # Never asks for fewer steps than were already run. + assert ( + probe_steps_for_spread(0.0005, measured_steps=30_000, target_spread=0.002) + == 30_000 + ) + with pytest.raises(ValueError): + probe_steps_for_spread(0.001, measured_steps=30_000, target_spread=0.0) + + +def test_slowing_the_chain_costs_independent_probe_samples(): + fast = integrated_autocorrelation(CYCLE_1, n_steps=60_000, seed=3) + slow = integrated_autocorrelation(SLOW, n_steps=60_000, seed=3) + + assert fast < 30.0 + # The precision cost is what a wider metric is traded against. + assert slow > 5.0 * fast + + +def test_a_realistically_collected_probe_has_a_much_wider_interval(): + # Sixteen long trajectories is how `collect_probe_data` gathers a rollout. + # Four thousand short ones is nearly independent sampling, and flatters the + # interval by hiding the correlation a real probe actually incurs. + def interval(n_chains): + fit = simulate_parallel(SLOW, n_steps=20_000, seed=7, n_chains=n_chains) + test = simulate_parallel(SLOW, n_steps=10_000, seed=8, n_chains=n_chains) + low, high = block_bootstrap_interval( + fit, test, window=32, seed=5, resamples=120 + ) + return (high - low) / 2.0 + + assert interval(PROBE_CHAINS) > 2.0 * interval(4_000) + assert interval(PROBE_CHAINS) > 0.005 diff --git a/tests/test_mess3_token_guess_references.py b/tests/test_mess3_token_guess_references.py new file mode 100644 index 00000000..995ec387 --- /dev/null +++ b/tests/test_mess3_token_guess_references.py @@ -0,0 +1,116 @@ +"""Tests for the token-guess metric reference points.""" + +from __future__ import annotations + +import numpy as np +import pytest + +from envs.mess3.model import passive_model +from experiments.mess3_token_guess_cycle_1.comparison.experiment import ( + BASE_MODEL_CONFIG, + ENV_CONFIG, +) +from experiments.mess3_token_guess_cycle_2.metric_references import ( + ALPHA, + bayes_accuracy_by_context, + compute_references, + normalise, + raw_token_window_r2, + simulate_stream, + untrained_module_r2, +) + + +def test_simulated_belief_matches_the_delay_one_transducer_update(): + model = passive_model(alpha=ALPHA) + transition = np.asarray(model.transition_matrix) + emission = np.asarray(model.emission_matrix) + stream = simulate_stream(n_steps=64, seed=3, window=8) + + for index in range(len(stream.beliefs) - 1): + belief = stream.beliefs[index] + token = stream.tokens[index] + # delay=1 composes the observation operator before the transition, so + # the update is diag(P(y|s)) @ T applied in row-vector convention. + operator = np.diag(emission[:, token]) @ transition + expected = belief @ operator + expected /= expected.sum() + np.testing.assert_allclose(stream.beliefs[index + 1], expected, atol=1e-12) + + +def test_simulated_beliefs_are_probability_vectors(): + stream = simulate_stream(n_steps=256, seed=11, window=8) + np.testing.assert_allclose(stream.beliefs.sum(axis=1), 1.0, atol=1e-12) + assert (stream.beliefs >= 0.0).all() + + +def test_raw_token_probe_saturates_well_below_the_supervised_ceiling(): + fit = simulate_stream(n_steps=20_000, seed=100) + test = simulate_stream(n_steps=10_000, seed=101) + scores = raw_token_window_r2(fit, test, context_lengths=(1, 2, 4, 8, 16)) + + # More observations never hurt an affine probe by a meaningful margin. + ordered = [scores[k] for k in (1, 2, 4, 8, 16)] + assert all(later >= earlier - 5e-3 for earlier, later in zip(ordered, ordered[1:])) + + # A single observation is already most of the way there, and the window + # saturates by eight observations, short of the supervised 0.9989. + assert 0.78 < scores[1] < 0.83 + assert scores[8] == pytest.approx(scores[16], abs=5e-3) + assert 0.95 < scores[8] < 0.98 + + +def test_raw_token_floor_exceeds_several_published_cycle_1_scores(): + fit = simulate_stream(n_steps=20_000, seed=100) + test = simulate_stream(n_steps=10_000, seed=101) + floor = raw_token_window_r2(fit, test, context_lengths=(8,))[8] + + # Cycle 1 reported these for reward-only, max-entropy, and predictive-loss + # PPO. A probe on the untransformed observations beats all three. + for reported in (0.8552, 0.8558, 0.9319): + assert reported < floor + + +def test_bayes_accuracy_spans_a_narrow_band_above_repeat_previous_token(): + stream = simulate_stream(n_steps=40_000, seed=102) + accuracy = bayes_accuracy_by_context(stream, context_lengths=(1, 2, 4, 8, 64)) + + # One observation is the repeat-the-previous-token rule. + assert 0.66 < accuracy[1] < 0.69 + assert accuracy[8] == pytest.approx(accuracy[64], abs=5e-3) + assert accuracy[64] - accuracy[1] < 0.03 + + +def test_normalise_reports_position_within_the_usable_range(): + assert normalise(0.9670, floor=0.9670, ceiling=0.9989) == pytest.approx(0.0) + assert normalise(0.9989, floor=0.9670, ceiling=0.9989) == pytest.approx(1.0) + assert normalise(0.8552, floor=0.9670, ceiling=0.9989) < 0.0 + with pytest.raises(ValueError): + normalise(0.5, floor=0.9, ceiling=0.9) + + +def test_untrained_module_is_probed_with_the_study_architecture(): + result = untrained_module_r2( + env_config=ENV_CONFIG, + model_config=BASE_MODEL_CONFIG, + seed=5, + fit_steps=2_000, + test_steps=1_000, + ) + assert result["n_fit"] == 2_000 + assert result["n_test"] == 1_000 + assert -1.0 <= result["r_squared"] <= 1.0 + assert 0.0 <= result["token_accuracy_greedy"] <= 1.0 + + +def test_compute_references_reports_a_floor_below_the_supervised_ceiling(): + references = compute_references(seed=7, fit_steps=8_000, test_steps=4_000) + floor = references["belief_r2_floor"] + low, high = references["belief_r2_probe_noise_95ci"] + + assert 0.94 < floor < 0.99 + assert low < floor < high + assert references["accuracy_floor_repeat_previous_token"] < ( + references["accuracy_ceiling_bayes"] + ) + assert "untrained_module" not in references diff --git a/tests/test_mess3_token_guess_statistics.py b/tests/test_mess3_token_guess_statistics.py new file mode 100644 index 00000000..ba66d342 --- /dev/null +++ b/tests/test_mess3_token_guess_statistics.py @@ -0,0 +1,104 @@ +"""Tests for the token-guess seed-level statistics.""" + +from __future__ import annotations + +import numpy as np +import pytest + +from experiments.mess3_token_guess_cycle_2.statistics import ( + compare, + holm_adjust, + seeds_for_power, + summarise, + t_test_power, +) + +# Kelly cycle 3 belief-probe R², seeds 42/43/44, as recorded in each arm's +# results//condition_summary.json. +CYCLE_3_PPO = {42: 0.860618, 43: 0.932631, 44: 0.829447} +CYCLE_3_IQN = {42: 0.974336, 43: 0.975997, 44: 0.956050} + + +def test_summarise_reports_a_wider_interval_than_the_published_plus_minus(): + values = list(CYCLE_3_PPO.values()) + estimate = summarise(values) + + population_sd = float(np.std(values, ddof=0)) + assert estimate.sample_sd > population_sd + + # The published "±" was a population standard deviation. A 95% interval on + # the mean of three seeds is several times wider. + assert estimate.half_width > 2.0 * population_sd + assert estimate.ci_low < estimate.mean < estimate.ci_high + assert estimate.n == 3 + + +def test_summarise_rejects_a_single_seed(): + with pytest.raises(ValueError): + summarise([0.97]) + + +def test_seeds_for_power_scales_with_the_squared_noise_to_signal_ratio(): + assert seeds_for_power(0.02, 0.01) < seeds_for_power(0.01, 0.01) + assert seeds_for_power(0.01, 0.02) > seeds_for_power(0.01, 0.01) + # An unpaired design needs more seeds than a paired one. + assert seeds_for_power(0.01, 0.01, paired=False) > seeds_for_power( + 0.01, 0.01, paired=True + ) + assert seeds_for_power(0.0, 0.01) >= 10_000 + + +def test_power_rises_with_seeds_and_the_planner_agrees_with_it(): + difference, sample_sd = 0.010, 0.008 + powers = [t_test_power(n, difference, sample_sd) for n in (3, 6, 12, 24)] + assert powers == sorted(powers) + assert powers[0] < 0.8 < powers[-1] + + needed = seeds_for_power(difference, sample_sd) + assert t_test_power(needed, difference, sample_sd) >= 0.8 + assert t_test_power(needed - 1, difference, sample_sd) < 0.8 + + +def test_the_headline_ppo_versus_iqn_gap_is_unresolved_at_three_seeds(): + # Kelly cycle 3 reported PPO 0.8742 ± 0.0432 against IQN 0.9688 ± 0.0090, + # which reads as decisive. Paired across the seeds both arms actually ran, + # the interval still contains zero. + comparison = compare(CYCLE_3_IQN, CYCLE_3_PPO) + + assert comparison.n == 3 + assert comparison.difference > 0.09 + assert not comparison.resolved + assert comparison.ci_low < 0.0 < comparison.ci_high + assert comparison.p_value > 0.05 + assert comparison.seeds_for_power > 3 + + +def test_a_difference_smaller_than_its_noise_is_not_resolved_by_three_seeds(): + # Two conditions separated by less than their seed-to-seed spread. + left = {42: 0.9720, 43: 0.9660, 44: 0.9770} + right = {42: 0.9700, 43: 0.9700, 44: 0.9700} + comparison = compare(left, right) + + assert abs(comparison.difference) < 0.005 + assert not comparison.resolved + assert comparison.seeds_for_power > 3 + + +def test_compare_requires_shared_seeds(): + with pytest.raises(ValueError): + compare({42: 0.9}, {43: 0.8}) + + +def test_holm_adjustment_is_monotonic_and_never_shrinks_a_p_value(): + raw = {"a": 0.001, "b": 0.02, "c": 0.04, "d": 0.5} + adjusted = holm_adjust(raw) + + assert set(adjusted) == set(raw) + for name, value in raw.items(): + assert adjusted[name] >= value + ordered = sorted(raw, key=lambda name: raw[name]) + values = [adjusted[name] for name in ordered] + assert values == sorted(values) + # The smallest of four p-values is multiplied by four. + assert adjusted["a"] == pytest.approx(0.004) + assert holm_adjust({}) == {}