Trie packing

Three sequences, one shared tree

A trie is a prefix tree. Trie packing merges identical prefixes across training sequences into a single set of token nodes. Each root-to-leaf path still represents an original sequence. A token is shareable only when its entire preceding context is identical.

Sequence 1: A B C D
Sequence 2: A B C E
Sequence 3: A B F G

The formulas stay; the layout and visibility change

Step 1 · Store each node once

packed tokens: [A, B, C, D, E, F, G]
packed index:   0  1  2  3  4  5  6
position IDs:  [0, 1, 2, 3, 3, 2, 3]

input_ids: [1, 7]
hidden:    [1, 7, hidden_dim]

Position IDs use the node’s depth on its original path, not its flattened index. F is at packed index 5 but position 2 in ABFG. Use these positions for RoPE or learned position embeddings.

Step 2 · Allow only ancestors and self

Apply this visibility at every attention layer. Rows are queries; columns are keys. A 1 means visible.

E sees ABCE, never D. G sees ABFG, never CDE. This is not the ordinary lower-triangular mask of a length-7 sequence.

Step 3 · Run the ordinary transformer operations

M_tree is 0 for allowed connections and −∞ elsewhere. Q/K/V projections, FFNs, residuals, and token-wise normalization use their usual formulas. The FFN processes 7 rows rather than 12.

Why this works: identical initial representations plus identical visible contexts produce identical layer outputs. Repeating this argument across layers makes every packed node equivalent to its unpacked occurrences, under deterministic, compatible model behavior.

Share logits, preserve every prediction

At node u, the LM head produces logits z_u = h_uW_LM: one raw score per vocabulary token. Softmax gives the next-token distribution p_u. The logits after C predict D or E; they do not predict C itself.

This correspondence is the loss mapping. It preserves nine prediction occurrences using four distinct logit rows.

Same chain rule, shared computation graph

Below, G_X means ∂L/∂X. We show derivatives for L_sum; divide them by 9 for L_mean.

Step 1 · Combine gradients at the shared logits

For target y, ℓ = −log p(y) = −z_y + log Σ_j exp(z_j). Differentiating gives G_z = p − e_y, where e_y is the target’s one-hot vector.

At C, the gradients for targets D and E add: (p_C − e_D) + (p_C − e_E). Different target contributions are preserved; they are not replaced with an arbitrary branch’s gradient.

Step 2 · Run the usual LM-head and FFN backward

The shared row’s incoming gradient is already accumulated. FFNs and normalization need no tree-specific adjustment. Residual branches add their gradients as usual.

For any shared linear operation Y = XW, two incoming contributions G₁ and G₂ can be combined before the matrix multiplications:

This principle also holds for nonlinear layers: at fixed forward activations, the backward Jacobian-vector operation is linear in its incoming gradient.

Step 3 · Attention collects descendant contributions

From O = PV, ordinary matrix backward gives:

For shared value row B, the final attention layer receives:

Queries B, C, and F all use the same V_B, so their gradients add. Queries D, E, and G contribute zero in this no-EOS example because their output activations do not affect any loss. With EOS targets, their contributions would be included too.

Through softmax and the QK product, for score matrix S = QK^T/√d_k + M_tree:

Masked edges contribute zero. These are single-head equations before any RoPE transformation; a full block also backpropagates through RoPE, output projection, normalization, and residual paths normally.

Step 4 · Repeat through earlier layers

The accumulated input gradient flows into the preceding layer. Repeat ordinary backward until reaching the embeddings. A loss-masked prompt node can still receive gradients from descendants that attend to it.

The training-loop change is small

This PyTorch sketch assumes a model adapter that accepts the tree mask with the correct backend-specific shape and convention. It is not a drop-in configuration for every model.

import torch
import torch.nn.functional as F

# A ... G are vocabulary token IDs, not packed row indices.
ids = torch.tensor([[A, B, C, D, E, F, G]], device=device)
positions = torch.tensor([[0, 1, 2, 3, 3, 2, 3]], device=device)

outputs = model(
    input_ids=ids,
    position_ids=positions,
    attention_mask=tree_mask,  # Correctly adapted for this backend
)
logits = outputs.logits[0]     # [7, vocab_size]; select batch item 0

# One entry per ORIGINAL prediction occurrence.
source_rows = torch.tensor(
    [0, 1, 2,  0, 1, 2,  0, 1, 5], device=device
)
targets = torch.tensor(
    [B, C, D,  B, C, E,  B, F, G], device=device
)

log_probs = F.log_softmax(logits, dim=-1)
token_logps = log_probs[source_rows, targets]  # [9]
loss = -token_logps.mean()                    # Original denominator

optimizer.zero_grad(set_to_none=True)
loss.backward()                              # Unchanged
optimizer.step()                             # Unchanged

Why indexing matters: differentiable indexing scatters and adds gradients into repeated source entries during backward. It preserves repeated targets without requiring custom hooks. Gathering target log-probabilities also avoids constructing nine full vocabulary-sized logit rows.

For PPO/GRPO, gather the same per-occurrence log-probabilities, then apply each occurrence’s own advantage, old log-probability, clipping rule, mask, and weight. Do not generally average advantages before applying a nonlinear clipped loss.

What changes, what does not

Memory: a dense mask is not enough

In this example, dense scores grow from 3 × 4² = 48 entries per head to 7² = 49, even though allowed query–key pairs fall from 30 to 21. Larger trees can amplify the dense-matrix overhead.

Use a tree-capable tiled attention implementation to avoid storing the full score matrix, and block sparsity to skip forbidden work. Compact mask construction and a unique-token budget per pack also matter. Peak memory can still exceed a baseline that processes one trajectory at a time.

Verify equivalence before measuring speed

1. Start with the same parameters and deterministic forward behavior. Disable dropout for the first comparison.
2. Compare every mapped packed logit row with its corresponding unpacked row.
3. Compare the loss with identical reduction, masking, and sample weighting.
4. Compare every parameter gradient, then one optimizer update, within dtype-appropriate numerical tolerances.
5. Test forks, repeated targets, EOS, padding, masked prompts, unequal sequence lengths, and multiple packed trees.

Watch the assumptions: detached prefix K/V breaks gradient flow. Independent dropout masks prevent exact sample-wise identity. MoE capacity limits and auxiliary losses may depend on original token multiplicities. Relative/windowed attention must preserve the original path semantics. Distributed loss scaling and checkpoint recomputation must remain consistent.

Implementation references

The walkthrough above is a mathematical reference design. Framework APIs and optimized kernels differ.

• Tree Training: Accelerating Agentic LLMs Training via Shared Prefix Reuse — tree packing and gradient restoration.
• AReaL: Tree Training — configuration and backend support.
• TRL: experimental asynchronous GRPO — prefix-tree layout and per-occurrence loss mapping.
• PyTorch: FlexAttention — custom attention masks and block sparsity.