README.md

---
license: apache-2.0
base_model: Qwen/Qwen3-4B-Base
tags:
  - compositional-generalization
  - interpretability
  - recall-then-assemble
  - grpo
language:
  - en
---

# Compositional-generalization checkpoints (recall-then-assemble, Qwen3-4B-Base)

## TL;DR

15 checkpoints of Qwen3-4B-Base from **one controlled study of whether a model can compose skills it has
only ever seen in isolation**. Every checkpoint is one cell of that study; cells differ in exactly one
training variable, so pairs of folders isolate one effect each. One subfolder = one complete HF model
directory (fp32 `model.safetensors` ≈ 16 GB, tokenizer, chat template).

```python
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
repo, sub = "Physicsru/compgen-checkpoints", "sft-eco-s1"   # any folder from the table below
tok = AutoTokenizer.from_pretrained(repo, subfolder=sub)
model = AutoModelForCausalLM.from_pretrained(repo, subfolder=sub, torch_dtype=torch.float32)
```

**Which two folders compare what** (held-out compositional accuracy at depth 4 / depth 8 unless stated):

| contrast | without | with | effect |
|---|---|---|---|
| **Co-occurrence atomic data** (several independent single-skill tasks in one prompt/answer, width 2-4) — cleanest version: identical task stream and steps | `sft-matched-isolated-s1` | `sft-matched-grouped-s1` | 0.65 / 0.16 → 0.93 / 0.47 |
| **Co-occurrence atomic data** — the paper's two models (same compositions; E-co has fewer rows) | `sft-v1-s1` | `sft-eco-s1` | 0.53 / 0.09 → 0.97 / 0.73 |
| **Composition demos (depth 2-4, train skills only)**, atomics single-task | `sft-d1-s1` | `sft-v1-s1` | 0.00 / 0.00 → 0.53 / 0.09 |
| **Composition demos**, atomics in co-occurrence groups | `sft-c1-s1` | `sft-eco-s1` | 0.66 / 0.08 → 0.97 / 0.73 |
| **Depth of the demos** (depth-2 only vs depth 2-4, same row count) | `sft-c4-s1` | `sft-c4b-s1` | 0.70 / 0.12 → 0.90 / 0.62 |
| **RL** (GRPO on train-skill chains + atomic prompts; held-out skills never composed) — good outcome | `sft-eco-s1` | `rl-mdb-step100` | held-out chains d20 0.50 → 0.95 |
| **RL** — drifted outcome (one held-out skill loses a parameter) | `sft-eco-s1` | `rl-r1a-step100` | d20 0.50 → 0.61; `func_24` arity error 25% |
| **Seed only** (same RL pool, same init) | `rl-r1a-step100` | `rl-r1b-step100` | d20 0.61 vs 0.89; arity error 25% vs 1% |
| **Same RL run over time** (drift switches on, then off) | `rl-mda-step5` → `rl-mda-step50` | → `rl-mda-step100` | arity error 6% → 44% → 0.3%; d20 0.82 → 0.57 → 0.93 |
| **The recall-then-assemble format itself** | `stage15b-init` | any `sft-*` | 0.03 / 0.00 → … |

The four SFT cells `sft-d1-s1`, `sft-v1-s1`, `sft-c1-s1`, `sft-eco-s1` form a 2×2 (co-occurrence no/yes ×
demos no/yes); all SFT cells start from `stage15b-init`; all RL runs start from `sft-eco-s1`.

## The task

A library of 25 string operators with opaque names `func_0` … `func_24` (e.g. `func_3` = sort the
characters; `func_24` = a backtracking palindrome transform with a semantically unused `depth` parameter).
The model first learns every operator **as an atomic skill** (closed-book: given `def main_solution(x):
return func_3(x)` and an input, recall the body of `func_3` and write a program). It is then trained on
**compositions** (`func_5(func_3(func_1(x, 2)), 'ab')`) built **only from the 13 train operators**; the 12
**held-out operators** are never composed in any training data. The test asks the model to compose the
held-out operators. The answer format is *recall-then-assemble* (RA): one `Recall func_N: <gloss>` block
with the definition per operator, then one `Assemble:` block with all definitions plus `main_solution`.

Metrics: **held-out d4 / d8** = program accuracy on held-out-operator programs of expression depth 4 / 8
(256 per depth, `eval/heldout_trees_d1-8.parquet`); **chains d20** = pure nesting chains f1(f2(…f20(x)))
over held-out operators (`eval/heldout_chains_d1-20.parquet`); **arity error** = fraction of `Recall
func_24` blocks that drop the second parameter (`def func_24(s):`), the single failure that separates good
and bad RL outcomes. Greedy decoding, 3,072 new tokens.

## Lineage

```
Qwen3-4B-Base
 └─ stage 1: atomic skills with bodies shown (RFT on code execution)
     └─ stage 1.5 "closed-book": 20k depth-1 atomic tasks over all 25 ops + 12k depth-2..4 TRAIN-op
        compositions, prose + one code block (no RA format)                      → stage15b-init
         ├─ RA bootstrap SFT cells (2 epochs, batch 128, LR 2e-5, seed 1)         → sft-*
         │    all cells share the SAME 10,000 atomic tasks (400 per op, all 25 ops) and differ in
         │    (a) how atomic tasks are packed into answers and (b) which composition rows are added
         └─ (from sft-eco-s1) GRPO, KL 0.01 to the init, 100 steps                → rl-*
```

## Checkpoints

| folder | init | atomic side (10k single-op tasks) | composition side | RL pool | steps | held-out numbers |
|---|---|---|---|---|---|---|
| `stage15b-init` | stage 1 | 20k closed-book depth-1 rows, prose + one code block | 12k depth-2..4 train-op comps (same format) | — | 500 | d4 / d8 = 0.03 / 0.00 (no RA format yet) |
| `sft-d1-s1` | stage15b-init | **single-task**: one op per answer | **none** | — | 152 | 0.00 / 0.00 |
| `sft-v1-s1` | stage15b-init | single-task | 15,979 mixed-structure train-op comps, depth 2-4 | — | 400 | 0.53±0.15 / 0.09±0.05 (3 seeds; this seed 0.75 / 0.15) |
| `sft-c1-s1` | stage15b-init | **co-occurrence groups**: 1-4 independent single-op tasks in one answer (all ops) | **none** | — | 58 | 0.66±0.14 / 0.08±0.04 |
| `sft-eco-s1` ("E-co") | stage15b-init | co-occurrence groups | the same 15,979 comps | — | 308 | 0.97±0.02 / 0.73±0.15 (this seed 0.98 / 0.83); chains d20 0.50 |
| `sft-c4-s1` | stage15b-init | co-occurrence groups | 12,452 mixed comps, **depth 2 only** | — | 252 | 0.70±0.03 / 0.12±0.05 |
| `sft-c4b-s1` | stage15b-init | co-occurrence groups | 12,436 mixed comps, **depth 2-4** (same count as c4) | — | 252 | 0.90±0.04 / 0.62±0.13 |
| `sft-matched-isolated-s1` | stage15b-init | per step: 48 atomic tasks as 48 single-task rows | per step: 76 comps (same stream as grouped) | — | 400 | 0.65±0.08 / 0.16±0.02 (3 paired seeds) |
| `sft-matched-grouped-s1` | stage15b-init | per step: the SAME 48 atomic tasks packed into 20 grouped rows | the same 76 comps | — | 400 | 0.93±0.03 / 0.47±0.04 |
| `rl-r1a-step100` | sft-eco-s1 | — | — | 1:1 mix of train-op serial chains d10-20 and all-op multi-task atomic prompts (widths 1-12) | 100 | chains d8/d12/d20 0.85 / 0.63 / 0.61; arity error 226/903 = 25% |
| `rl-r1b-step100` | sft-eco-s1 | — | — | same pool, different seed | 100 | 0.94 / 0.97 / 0.89; 9/903 = 1% |
| `rl-mda-step5` | sft-eco-s1 | — | — | **mixed-depth pool**: chains + prompts holding one deep train-op chain task followed by 1-2 long-literal atomic tasks over all ops | 5 | 0.95 / 0.92 / 0.82; 57/903 = 6% |
| `rl-mda-step50` | sft-eco-s1 | — | — | same run | 50 | 0.91 / 0.63 / 0.57; 396/903 = **44%** |
| `rl-mda-step100` | sft-eco-s1 | — | — | same run | 100 | 0.97 / 0.97 / 0.93; 3/903 = 0.3% |
| `rl-mdb-step100` | sft-eco-s1 | — | — | mixed-depth pool, second replicate | 100 | 0.99 / 0.98 / 0.95; 33/903 = 3.7% |

"Mixed-structure comps" = train-op compositions whose expression trees mix unary nesting, binary operators
and literals (the benchmark's native generator). "Co-occurrence groups" put several *independent* atomic
tasks into one prompt/answer so each definition is recalled next to others (no data flow between them).
Held-out operators appear in every cell only as depth-1 atomic tasks. All steps: batch 128; SFT 2 epochs.

## What each contrast shows (for choosing an analysis)

- **Co-occurrence** (`isolated` vs `grouped`, `v1` vs `eco`): the failure of the "without" model is a
  *chimera* — a held-out name written with a train operator's body (1,102 of 1,175 wrong bodies in v1).
  Decomposed inference (each helper recalled in its own forward pass, then assembled mechanically) scores
  1.00 for v1, so the whole gap is name→definition retrieval under multi-definition load, not missing
  knowledge. In the matched pair the gap opens from step 100.
- **Composition demos** (`d1` vs `v1`, `c1` vs `eco`): without co-occurrence the demos are what makes any
  multi-helper program possible (d1 never writes one); with co-occurrence the model already composes at
  mid depth without a single demo (c1 0.66) and the demos buy depth (d8 0.08 → 0.73). Demo *diversity*
  matters (c4 vs c4b), demo *volume* does not (2k vs 12k depth-2 demos: 0.77 vs 0.70).
- **RL** (`eco` vs `mdb` / `r1a`, `r1a` vs `r1b`, `mda` over time): RL on train-operator chains lifts deep
  serial execution for every operator (d20 0.50 → 0.9+), but in some seeds one held-out operator
  (`func_24`, whose dropped parameter is semantically inert) loses its second parameter when its call is
  nested among other library calls. The drift is latent in the SFT init (E-co writes `func_24(s)` in
  0 / 5 / 18% of blocks at name→argument distance <40 / 40-150 / 150+ chars), tracks the number of other
  library calls' argument tails between the name and the own argument, is not triggered by builtin-method
  tails, long literals, definition load or a preceding chain (probe families 1b/1d/1f/1g in `eval/`), and
  reverses inside one run (`mda` 44% → 0.3%). KL 0.01 keeps RL weights close to `sft-eco-s1`, so
  weight-difference analyses between `eco` and any `rl-*` folder, or between `mda` steps, are meaningful.

## Evaluation inputs (`eval/`)

| file | rows | what |
|---|---|---|
| `heldout_trees_d1-8.parquet` | 2,048 | held-out-operator programs, depths 1-8 × 256 (the d4 / d8 numbers) |
| `trainop_trees_d1-8.parquet` | 2,048 | the same over the 13 train operators |
| `heldout_chains_d1-20.parquet` | 2,816 | pure nesting chains over held-out operators, d ∈ {1,2,4,…,20} × 256, outputs ≤ 1,000 chars |
| `trainop_chains_d1-20.parquet` | 2,816 | the same over train operators |
| `probe_1b.parquet` | 6,912 | definition-load probe: target op recalled after r other recalls (teacher-forced prefix) |
| `probe_1d.parquet` | 1,792 | chain-position probe: target op at the outer / middle / inner position of a deep chain |
| `probe_1f.parquet` | 2,048 | builtin-method argument tails in front of the op's own argument |
| `probe_1g.parquet` | 1,280 | long literal (8-200 chars) as the op's first argument |

Rows are in the verl code-execution format: `prompt` (chat messages), `reward_model.ground_truth` (JSON
with `ref_code`, `ref_input`, `ref_output`), `extra_info` (depth, k = distinct ops, probe condition).

## Provenance

Every number is traceable in the project ledger `verl/analysis/RESULTS_PROVENANCE.md` (sections MAIN
RESULT, C1–C5 CAUSAL TABLE, MATCHED SFT, DEEP CHAINS, Mixed-depth arm RESULT, PROBES) to a job id, a
sweep directory and a CI / classification report; 3-seed bands are mean ± sd over trainer seeds 1 / 7 / 123.