mixle.experimental.summary_tree module

E4: hierarchical summary tree + multi-scale objective – see notes/designs/E4.md for the full derivation (persistent append-only tree over evicted tokens, tree-path positional encoding, the predict-the-summary auxiliary loss, the stop-gradient horizon receipt). This module implements that note section-by-section; see the note’s “Implementation notes vs. this design” section for the small, honestly-documented places this module simplifies the note’s scheme for tractability.

What this is. E1’s SlidingWindowSpine keeps an exact but bounded KV window; anything evicted from that window is gone. E4 keeps it: every evicted token is folded, one at a time, into a persistent tree of learned summaries via mixed-radix carry propagation (the fast-multipole-method structure – near field exact, far field via a bounded number of increasingly coarse representatives – applied to token history). Tree depth grows only as log_fanout(evicted_count), so the far-field attention set stays bounded regardless of how much history has streamed through.

Positional encoding. RoPE’s q . k dependence on i - j is well-conditioned only when i - j is a small, well-scaled number (E1’s window); a far-field summary node represents a range of possibly billions of original positions, and there is no single j to rotate by. E4 replaces RoPE for the far field with (a) a content channel – a level embedding plus a sibling-slot embedding summed into the node’s summary before it’s used as an attention key – and (b) a relative bias channel – a learned scalar indexed by tree distance (lca_depth, see below), the ALiBi/T5-bias shape of mechanism but indexed by tree distance instead of linear offset. Near-field window tokens keep ordinary RoPE unchanged (E1’s _rope_angles/_apply_rope, reused verbatim).

Predict-the-summary auxiliary loss. Every node, the moment it’s finalized, is scored against the exact additive token-id histogram of the leaves it covers via one shared linear head (d_model -> vocab) and cross-entropy against the normalized histogram. This is direct supervision independent of whether any future query ever attends to that node – for a node many levels up the tree that a training run may never query again, this is its compressor’s only gradient.

Stop-gradient horizon (receipted). A node moves from live to archived once H further nodes have finalized at its own level after it; .summary is .detach()-ed at that moment and never un-detaches. Archived nodes stay forward-visible (attendable) but stop receiving gradient. See SummaryTreeSpine.detach() and mixle/tests/summary_tree_test.py for the exact-accounting receipt.

class TreeNode(summary, histogram, path, level, g, finalized_step, finalized_index_within_level, detached=False, detached_at_finalized_count=None)[source]

Bases: object

One finalized node of the persistent summary tree (E4.md’s TreeNode).

summary carries one entry per layer (each (batch, d_model)) since every layer’s tree is built from that layer’s own (k, v) via that layer’s own qkv projection, but the bookkeeping fields below (histogram/path/level/g) are identical across layers – they describe which evicted tokens this node covers, not any layer-specific content – so they’re stored once.

Parameters:

  • summary (list[Any])

  • histogram (Any)

  • path (tuple[int, ...])

  • level (int)

  • g (int)

  • finalized_step (int)

  • finalized_index_within_level (int)

  • detached (bool)

  • detached_at_finalized_count (int | None)

class SummaryTreeState(window, cached_ids, pending_leaf, pending, live, archived, level_finalized_count, evicted_count=0)[source]

Bases: object

ContextMechanism carried state: E1’s exact near field plus the persistent far-field tree.

Parameters:

  • window (SlidingWindowState)

  • cached_ids (Any | None)

  • pending_leaf (list)

  • pending (list[list[TreeNode]])

  • live (list[list[TreeNode]])

  • archived (list[list[TreeNode]])

  • level_finalized_count (list[int])

  • evicted_count (int)

class SummaryTreeSpine(vocab, *, d_model=32, n_layer=2, n_head=2, window=16, fanout=4, max_level_cap=24, detach_horizon_nodes=2, aux_weight=0.1)[source]

Bases: Module

E4: SlidingWindowSpine’s exact near field plus a persistent, bounded far-field tree of learned summaries, merged into one joint softmax per layer (E4.md’s “Far-field attention: merging into E1’s score”). See the module docstring for the positional encoding and auxiliary loss this adds on top of E1 unchanged.

Parameters:

  • vocab (int)

  • d_model (int)

  • n_layer (int)

  • n_head (int)

  • window (int)

  • fanout (int)

  • max_level_cap (int)

  • detach_horizon_nodes (int)

  • aux_weight (float)

log_density(x, y)[source]

x, y: (n, T) long tensors. Returns -mean_per_position_nll for each of the n sequences, each scored independently (state re-initialized per row) – one non-streaming forward per row, computed by calling init_state + step once per row exactly as a length-T, single-chunk stream would (E5.md’s “GradLeaf citizenship”).

Parameters:
Return type:

Any

digits_of(n, base)[source]

Base-base digits of n, least-significant first. digits_of(0, base) == (0,).

Parameters:
Return type:

tuple[int, …]

lca_depth(query_pos, node_level, node_g, fanout, *, max_climb=64)[source]

Tree distance between a query at absolute position query_pos and a node covering the contiguous leaf range [g * fanout**level, (g + 1) * fanout**level) (g = the node’s own 0-based sequential index among nodes finalized at its level – see E4.md’s carry-propagation section). Returns how many levels above node_level the query’s ancestor chain must climb before it lands in the same subtree as the node – see notes/designs/E4.md’s “Implementation notes” section for why this integer recurrence is the exact algebraic equivalent of “matching leading digits of the two paths” without materializing query_pos’s (potentially ~1e9-long) digit expansion. Pure function of (query_pos, node_level, node_g, fanout) – no dependence on how the stream was chunked, which is what makes it stable under re-chunking (Acceptance §3).

Parameters:

  • query_pos (int)

  • node_level (int)

  • node_g (int)

  • fanout (int)

  • max_climb (int)

Return type:

int