Sampled-node phase RoPE

sampled-node-phase-rope

Status: concluded, positive (reviewed 2026-09-28; this line previously read "active"). The answer is in the front matter above.

Trainer note: post-fix

Question

Can sampled corpus-position / phase information improve the depth-16 CUDAX AGPT recipe by presenting trie paths at RoPE positions that match where their prefixes occur in the corpus?

The motivating concern is that the ordinary prefix trie pools all occurrences of the same prefix into one node:

mass(P) = sum_q mass(P, q)
target(P) = global next-token counts after P

If continuation statistics depend on local corpus phase, that pooling discards information that RoPE could use. This experiment tests several ways to expose phase information without physically materializing a full phase-conditioned trie.

See phase-unrolled-tree.md for the design sketch and phase convention.

Code Context

Branch: sampled-node-phase-rope

The run metadata records base commit 0e118a213e000e7d6fedc002470ef75549f67c84 with a dirty worktree. The relevant implementation changes live in the CUDAX v2 trainer path and position-data tooling:

• src/cudax/*_v2.cuh
• src/cudax/agpt_train_v2.cu
• src/tools/build_position_table.cr
• phase inspection / validation tools in src/tools/

The experiment artifacts are in rnd/sampled-node-phase-rope/.

Protocol

Shared settings unless noted:

• Corpus: data/input.txt, sampled heldout carve 5% x 10 chunks, seed 42
• Trainer: bin/agpt_train_v2
• Model: data/seeds/shake-d64L2-h4-dff256-s128-seed42.model
• Architecture: d_model=64, n_layers=2, n_heads=4, d_ff=256
• Effective AGPT depth/window: 16
• Optimizer: Adam lr=0.0015, beta1=0.9, beta2=0.999
• Training: pd=1, anc_grad=true, chunk_queries=50000
• Eval: multi-chunk heldout, CPU, batch_size=1
• Position data: /tmp/agpt_snp_25ep_position_data

The main tested experimental.rope_position_mode values were:

• sampled-unit-phase: sample RoPE phase from prefix occurrence statistics.
• phase-sweep: control that changes presentation phase without changing target weighting.
• phase-weighted: weight prefix losses by phase mass while keeping global next-token targets.
• phase-conditioned: use phase-conditioned target distributions.

Later runs moved from wrapped / non-wrapped presentations to the current unwrapped interpretation: RoPE positions are presented without wrapping the local attention path, while phase mass is still read modulo the phase window.

Results

Rows report final heldout metrics for each run. Most are from result.json; the RunPod d128/L6 phase-conditioned rows are from locally generated eval_raw_epoch_*.json files, because the pod image lacked Python for the post-training eval step. Lower is better.

Current Best

Best fixed-token result in this experiment:

Best standard chunked rolling byte result in this experiment:

The first major finding in this experiment was the d128/L6 phase-weighted unwrapped run. It was the first result in this line to beat the previous larger-model static best (fixed_token_ppl=4.1705, byte PPL 4.7086), and it showed that phase-aware training could keep improving past the point where the static d128/L6 run had already peaked. That made phase-weighted the first clear evidence that the phase signal was useful, not just a presentation artifact.

The later d128/L6 phase-conditioned direct runs are now the numerical best, with shuffled phase order improving both fixed-token and standard chunked rolling byte PPL. This should be read as a follow-up on the phase-weighted finding, not as a replacement for it: phase-weighted established the result and remains the simpler/scalability-friendlier baseline, while phase-conditioned direct shows that phase-specific targets can squeeze out a little more on this setup. The shuffled phase-order result suggests the exact phase presentation curriculum matters too.

Unlike the static d128/L6 runs, which peaked near epoch 128 and then degraded, the phase-aware heldout curves improved through epoch 512. The first d128/L6 phase-weighted run improved as follows:

The 512-epoch run used the same setup but a longer cosine schedule, so same-epoch checkpoints are not a pure continuation of the 256-epoch schedule. It still improved beyond the 256 run by the later checkpoints:

The 512-epoch phase-weighted unwrapped run was still improving at the recorded checkpoints:

The d64/L2 512-epoch phase-conditioned direct run also improved monotonically over the recorded late checkpoints, but stayed behind d64/L2 phase-weighted:

The d128/L6 phase-conditioned direct run was launched on an A100 via RunPod with chunk_queries=50000 after 8-epoch smoke checks at 25k, 50k, and 100k. The smoke runs had identical trained-query counts and effectively identical heldout metrics; 50k was marginally best at epoch 8. The 512-epoch run took 6385.1s of trainer wall time and continued improving through the final checkpoint:

A follow-up d128/L6 phase-weighted run matched chunk_queries=50000 and used a longer 768-epoch cosine schedule to test whether the simpler phase-weighted method would catch up under the same chunking and more epochs. It did not catch the phase-conditioned result. Its best late fixed-token checkpoint was epoch 512, while rolling byte PPL was lowest at epoch 768:

A matched d128/L6 phase-conditioned direct 768-epoch run used the same chunk_queries=50000 and checkpoint schedule. Unlike the phase-weighted 768 run, the late fixed-token curve kept improving through epoch 768. It did not beat the earlier 512-epoch phase-conditioned run on fixed-token PPL, but it did produce the best standard chunked rolling byte PPL recorded in this experiment:

A matched d128/L6 phase-conditioned direct 512-epoch run then changed only the phase presentation order from the regular sequential cycle to a deterministic shuffle (experimental.phase_order: shuffle, seed 42). It covered each phase once per 64-epoch cycle, but in permuted order. This was the first run to move the best fixed-token PPL below 3.92 and the best standard chunked rolling byte PPL below 4.50. The best checkpoint was epoch 448; epoch 512 regressed slightly but still beat the earlier sequential 512 run:

A second shuffled phase-order seed (phase_order_seed: 123) was run as a variance check. It improved over the earlier sequential 512-epoch run on fixed-token PPL at its best checkpoint, but did not reproduce the seed-42 best:

For reference, the best recorded larger-model static baseline from rnd/baseline-calibration-v2-static-sampled was:

External Comparisons

The strongest current result is now in the range of published Tiny Shakespeare character-model baselines, but the comparisons are not all strict apples-to-apples.

nanoGPT's reported Tiny Shakespeare character run uses the canonical data/shakespeare_char corpus (1,115,394 chars, 65-char vocab), but a different validation protocol from this experiment: first 90% train / last 10% validation, block_size=256, n_layer=6, n_head=6, n_embd=384, and reported best validation loss 1.4697. Since PyTorch cross entropy is in natural-log units, the corresponding perplexity is exp(1.4697)=4.3479; it does not become 2^1.4697 unless the loss was already measured in bits.

The clean claim is therefore not that AGPT has beaten nanoGPT under an identical protocol. The defensible claim is that the best phase-aware AGPT run has reached nanoGPT-class Tiny Shakespeare character perplexity while using a 16-character explicit window, versus nanoGPT's 256-character block context. The strict same-split baseline in hand is KenLM: AGPT's fixed-token PPL 3.9070 is below the current-split KenLM modified-KN best of 4.1183.

The phase-weighted concatenated-heldout sanity eval uses the same heldout bytes as the chunked run, but treats them as one document. It is close to the chunked result: fixed-token PPL moves from 3.9452 to 3.9540, while rolling byte PPL moves from 4.5201 to 4.4961. This suggests the headline result is not an artifact of the heldout chunk boundary convention.

Scaling Notes

The phase variants differ materially in how they should scale:

• phase-weighted is the practical path. It needs prefix/node phase-mass statistics and changes how sampled losses are weighted, but it does not multiply the target table by every phase.
• phase-conditioned is much more expensive if materialized directly, because targets become prefix-by-phase distributions. The storage and sampling pressure grow with the number of active (prefix, phase) pairs, not just the number of prefixes.
• As corpus size grows, rare prefixes spread out and many phase bins will have little or no mass. That sparsity is helpful if we keep the representation sparse, but wasteful if we allocate dense phase state.
• The likely scalable recipe is to keep a global prefix target, store compact phase summaries only where there is enough mass, and sample phase-aware presentations from those summaries. Bigger corpora should give better phase estimates for common prefixes, but deeper/rarer prefixes will still need smoothing or fallback to global targets.

Open scaling question: whether the useful phase signal is local modulo-position structure that saturates quickly, or whether larger corpora expose enough stable phase-conditioned continuation structure to justify richer sparse phase-conditioned targets.

Interpretation

• Phase-aware presentation did eventually improve the d64/L2 line, but the useful variant was not the first sampled-unit-phase or wrapped presentation.
• phase-weighted with unwrapped RoPE positions was the first strong result in this experiment. It established that phase-aware training could beat the static d128/L6 baseline and keep improving through 512 epochs.
• The d64/L2 phase-conditioned direct run improved through 512 epochs but did not beat phase-weighted. The d128/L6 phase-conditioned direct run did beat the d128/L6 phase-weighted result by a small margin at epoch 512, making it the current numerical best.
• Matching chunk_queries=50000 and extending d128/L6 phase-weighted to 768 epochs did not close that gap. This supports, but does not yet fully prove, the hypothesis that phase-conditioned targets add useful information beyond phase-weighted global targets.
• Combining phase-aware training with d128/L6 beat the previous static d128/L6 project best and did not show the static run's post-128 heldout degradation by epoch 512.
• The d128/L6 phase-weighted curve crossed the 4.0 fixed-token threshold and was still improving in the 512-epoch cq25k run, though the gains from 384 to 512 were tapering. Under the 768-epoch cq50k schedule, its fixed-token PPL was already worse by epoch 576, while byte PPL only recovered slightly by epoch 768.

Reproduction

The configs are checked into this directory. Example:

XDG_CACHE_HOME=/tmp/agpt-xdg-cache \
HF_DATASETS_CACHE=/tmp/agpt-hf-datasets \
HF_HOME=/tmp/agpt-hf-home \
bin/agpt_experiment \
  --config rnd/sampled-node-phase-rope/d64L2-depth16-pd1-adam-lr0015-512ep-phase-weighted-unwrapped.yml \
  --trainer v2

Prerequisites:

• Build bin/agpt_train_v2 and bin/agpt_experiment.
• Ensure the trie path in the config exists or is rebuilt.
• Ensure /tmp/agpt_snp_25ep_position_data exists. Use the local build_position_table / phase validation tools from this branch to rebuild the position data if needed.

Artifact Notes

The crash on 2026-06-03 interrupted post-training evaluation for 20260603T112551-d64l2-depth16-pd1-adam-lr0015-512ep-phase-weighted-unwrapped. The checkpoint files were already present through epoch 512. The missing HF conversion/eval artifacts, result.json, and runs.json were reconstructed from the saved checkpoints using the same conversion/eval commands as src/tools/agpt_experiment.cr.

The durable records are:

• Per-run result.json
• Per-run meta.json
• rnd/sampled-node-phase-rope/runs.json

Large model checkpoints and HF conversion directories are output artifacts, not the core experiment record.