mixle.experimental.context_spine module

E1: the chunked-recurrent training spine every Track-E long-context mechanism plugs into.

See notes/designs/E1.md for the design decisions (RoPE over learned-absolute position embeddings, windowed-mask derivation, why detach_horizon only ever cuts the backward graph and never the forward one, and why this ships its own tiny TBPTT driver instead of routing through GradLeaf).

mixle.models.transformer.CausalLM and mixle.models.streaming_transformer_leaf.StreamingTransformer both train bounded, independent micro-batches with a learned position table capped at block tokens – neither carries state across calls. ContextMechanism is the minimal protocol that adds streaming + carried state + truncated-backprop-through-time (TBPTT) on top, without touching either of those modules. SlidingWindowSpine is the E1 baseline mechanism (Transformer-XL-style stop-gradient KV carry); E2-E6 differ only in what step’s carried state contains.

class ContextMechanism(*args, **kwargs)

Bases: Protocol

The substrate contract every Track-E long-context mechanism implements.

step is per-position teacher-forced (returns the mean loss over every position in the chunk, not just the last one – unlike CausalLM.forward, which returns only the last position’s logits).

init_state(batch_size, *, device='cpu')

A fresh state for batch_size independent streams (empty cache / zero memory).

Parameters:

• batch_size (int)

• device (str)

Return type:

Any

step(state, chunk)

chunk = (x, y), (batch, T) long tensors. Returns (new_state, mean_loss).

Parameters:

• state (Any)

• chunk (tuple[Any, Any])

Return type:

tuple[Any, Any]

detach(state)

Stop-gradient the carried state (cuts the TBPTT backward graph at this point).

Parameters:

state (Any)

Return type:

Any

class SlidingWindowState(cache_k=<factory>, cache_v=<factory>, pos=0, cache_positions=<factory>, batch_size=None, n_head=None, head_dim=None, device=None)

Bases: object

Raw per-layer KV caches and their explicit absolute positions.

Parameters:

• cache_k (list[Any])

• cache_v (list[Any])

• pos (int)

• cache_positions (list[Any])

• batch_size (int | None)

• n_head (int | None)

• head_dim (int | None)

• device (Any | None)

train_tbptt(mechanism, state, chunks, opt, *, detach_horizon=1, extra_loss=None)

Stream chunks through mechanism, TBPTT-training with the given optimizer.

Every detach_horizon chunks (or at end of stream, whichever comes first): backward the mean accumulated loss, step the optimizer, then mechanism.detach(state) to cut the graph before continuing. detach_horizon=1 is literal per-chunk stop-gradient (Transformer-XL); a horizon spanning the whole stream means no mid-stream detach happens at all (see notes/designs/E1.md). Returns {"losses": [float, ...], "state": final_state} – one loss per chunk, detached telemetry.

extra_loss is an optional live tensor added to the FIRST backward only – a term earned outside chunks (e.g. an auxiliary regularizer from a warm-up prefix this call does not train the LM on). Folding it in here rather than backpropagating it separately keeps the optimizer-step count independent of whether the term is present, so an ablation on it is not confounded by step count. It is excluded from the returned losses telemetry, which stays one entry per chunk.

Parameters:

• mechanism (ContextMechanism)

• state (Any)

• chunks (Any)

• opt (Any)

• detach_horizon (int)

• extra_loss (Any | None)

Return type:

dict[str, Any]