Seq_len decoupling

Seq_len decoupling — Phase 0: position → contributing-nodes map

Goal: build the bookkeeping needed for the decoupled-attention work in notes/seq-len-extension/shared_key_rope.md. For each corpus position p, find the radix-trie nodes whose path terminates at corpus position p. Equivalently: position p has a representation at every node R where R's path-from-root spells out a backward-suffix of length 1 to d ending at p.

Phase 0 framing (clarified by user 2026-05-11)

A corpus position p is represented in the trie at multiple nodes, one per backward-suffix length:

So position p has up to d representations spread across d different unigram subtrees (in the full uncompressed view). In a radix-compressed trie, the count can be less than d because multiple consecutive depth levels may collapse onto one radix edge.

Position 0 has exactly 1 representation (just corpus[0] in its unigram root). Position 1 has 2. Up to position d-1 the count ramps; after that every position has up to d.

What this is not

• Not a single landing node per position (my first iteration).
• Not a single root-to-leaf walk per position with all touched nodes attributed to the starting position (my second iteration — also wrong).
• The correct framing: nodes are attributed to the corpus position where their path terminates, not where the walk starts.

Algorithm

For each starting position s ∈ [0, N-1]:

1. Walk forward up to d chars from root through the radix trie.
2. At each radix node touched, attribute that node to the corpus position where its path ends: terminal_pos = s + cumulative_match_len − 1.

Each walk emits ~9 contributions (one per radix node visited), distributed across positions [s, s+d-1]. Aggregated over all walks, each position p collects contributions from walks starting at s ∈ [max(0, p-d+1), p].

Tool

bin/agpt_position_map (source: src/tools/agpt_position_map.cr):

bin/agpt_position_map \
  --trie /home/trans/agpt-tries/gutenberg_5m_d32_radix_corpus \
  --file data/gutenberg_5m.txt \
  --out rnd/seq-len-decouple/gutenberg_5m_d32_pos_to_node.bin

Output binary format: per corpus position p, a record [k:int32] [nid_1, ..., nid_k:int32] where k = number of radix nodes representing p, and the nid_i are sorted by walk (deepest first = longest backward suffix first).

Results — Shakespeare 1M, d=8

Corpus starts "First Citizen:..." — first 8 positions verified manually:

Position 0: 1 representation. ✓ (matches the verification claim)

Aggregate stats:

Results — Gutenberg 5M, d=32

Contributions-per-position distribution (k = radix nodes per terminal position)

Mean 9.24, modal 8, very long tail. The k=32 tail are positions where the walks happen to terminate at every char-depth (no radix compression along any of the 32 backward-suffix paths).

Top contributing nodes

node_id   contribs  mass     depth  edge_len
149117    821,687   821,687  1      1
2772912   477,225   477,225  1      1
6552035   339,827   339,827  1      1
1839620   306,058   306,058  1      1
5318063   290,452   290,452  1      1

The most-contributed-to nodes are depth-1 unigram root children: every corpus position contributes to some unigram root (the one matching its char). Top node (149117, mass 821k) is the most common starting char in the corpus — likely space (32 in token-id terms).

Mass consistency

For each radix node N, the count of corpus positions attributing to N exactly equals the trie's stored edge_mass for N. Zero mismatches across all 7.5M nodes and 46M contributions. This verifies the walk records exactly the same contributions the trie builder did.

Implications for seq_len decoupling

Each corpus position has ~9 radix-node identities at d=32, ranging from broad (depth-1 root child, mass hundreds-of-thousands) to narrow (depth-32 leaf, mass 1). For shared-key RoPE attention at query position q:

K_p^(k) = RoPE(k_{n_p^(k)}, p)    where n_p^(k) is the radix node
                                  representing p at backward-length k

The model has 9 plausible vectors per position. Phase 1 candidates:

1. Always-leaf: use the deepest backward-suffix identity (longest context, lowest mass).
2. Always-shallow-best: use the deepest with mass ≥ τ. Falls back to shorter contexts when leaf is a singleton.
3. Stack-attention: emit all of a position's identities as separate attention keys.
4. Adaptive per-query: choose backward-length based on what the query needs.

Files

rnd/seq-len-decouple/
├── README.md                                       (this file)
├── gutenberg_5m_d32_pos_to_node.bin                (205 MB, gitignored)
└── logs/
    ├── gutenberg_d32.log                           (v1: leaf-only, wrong)
    ├── gutenberg_d32_full_contrib.log              (v2: full but mis-attributed)
    └── gutenberg_d32_terminal_attr.log             (v3: correct)

Phase 1A — Naive seq_len extension (FAILED, expected)

Diagnostic test: take an existing d=16-trained AGPT model, evaluate perplexity at increasing seq_len, see whether the model can naturally use context past d=16 via standard transformer RoPE extrapolation.

Model: /tmp/agpt-gut-d16-pd1.model (30 SE pd=1 d=16 Gutenberg, mass-weight=off, PPL@16 = 8.01). Eval on data/gutenberg_5m.txt.

Conclusion: the d=16 model cannot use positions beyond seq=16 through naive extrapolation. PPL rises monotonically with seq_len beyond d, doubling by seq=48. Standard RoPE extrapolation does not work at this scale — the model never saw those position rotations during training, and there's no latent capacity to read them.

This rules out the cheapest hypothesis (A: extend seq_len with the existing architecture and hope) and establishes that breaking the d-shackle requires architectural change: option B, shared-key RoPE with explicit trie-derived identities per position, as described in notes/seq-len-extension/shared_key_rope.md.

Phase 1B — Shared-key RoPE architecture (next step)

Required: at attention time for query position q, build K vectors for the seq_len positions in the window using the trie's per-node base keys plus per-position RoPE, instead of the model's W_K · residual.

Components:

• Per-radix-node learned key vectors (one per node, ~7M × D extra parameters for d=16). Could be initialized from the trie's empirical P_out distribution or randomly.
• Inference path that, for each attention query, gathers the K vectors via the position→node binary map (already built in Phase 0).
• Training path that updates the per-node key vectors.

Open questions for design:

• Use which of the ~9 node identities per position? (always leaf, shallow-best, stack-all, adaptive)
• Share base keys between forward and backward direction? (unified DAG)
• Train from scratch or fine-tune existing AGPT model?

Followups

• Confirm node 149117 is the unigram root for space (id 32 = ' ' in the dataset's char_to_id).
• Implement inverse map dump (--out-inverse PATH) — per radix node, list the corpus positions terminating there. Decoupled attention needs this to gather K vectors at training/inference time.
• Phase 1: shared-key RoPE inference at seq_len > d on the AGPT 4.90 PPL model. Use the binary map to drive position-to-node lookup during attention.