Struct Layers

pub struct Layers<M>where
    M: Module,{
    pub n_real_layers: usize,
    pub n_virtual_layers: Option<(usize, Schedule)>,
    pub real_layers: Vec<Layer<M>>,
    pub ignore_first_residual: bool,
    pub ignore_last_residual: bool,
    pub residuals: Residuals,
    pub class_latents: Vec<ClassLatent>,
    pub class_latents_emb: Option<Param<Tensor<2>>>,
    pub grad_horizon: Option<usize>,
}

Expand description

A stack of Layers with optional virtual-layer scheduling — one struct for every Block family.

Fields§

§n_real_layers: usize

Number of real (weight-bearing) layers.

§n_virtual_layers: Option<(usize, Schedule)>

Optional (n_virtual_layers, schedule) for weight-sharing.

§real_layers: Vec<Layer<M>>

The weight-bearing layers, length n_real_layers.

§ignore_first_residual: bool

Zero the first virtual layer’s residual when true.

§ignore_last_residual: bool

Zero the last virtual layer’s residual when true.

§residuals: Residuals

How residuals are threaded between layers (plain additive vs Multi-Gate).

§class_latents: Vec<ClassLatent>

Positions of the stack-level class latents, spliced into the sequence once before the first virtual layer (independent of any per-Layer class latents). Empty ⇒ none.

§class_latents_emb: Option<Param<Tensor<2>>>

The stack-level class-latent embeddings, [num_class_latents, d_model].

§grad_horizon: Option<usize>

Back-propagate only the last K virtual layers; everything below runs without building an autodiff graph. None (the default) tracks the whole stack.

This is the truncated-BPTT knob of TRM/HRM-style deep recursion: with n_virtual_layers far above n_real_layers, tracking every pass is what runs out of memory, and both papers back-propagate only a suffix (TRM one full recursion, HRM-Text a horizon K warmed from 2 to 5). K is counted from the top so it stays meaningful when the stack depth changes, and so a training loop can move it per step.

Under weight sharing the same real layer serves both sides of the cut; the prefix runs an inner-backend copy, so each weight still receives gradient — from its tracked applications only.

The stack input is the exception, and deliberately so. It enters at the bottom and rides the residual stream upward, so a cut would sever its only path and a network’s in_proj (or a vocab net’s embedding) would never train at all — silently. TRM and HRM never meet this because they re-inject the input at every recursion; this stack reads it once. The boundary therefore re-attaches it straight-through: a value-zero term restores an identity gradient path, which under Residuals::Standard is not a guess but the exact leading term of ∂(x + Σ F_l)/∂x, the rest being precisely the prefix one chose not to differentiate. Under MultiGate the residual lives in the depth-streams rather than the token, so every carrier gets the identity path — the seed stream is the input and the pool is convex, so an identity prefix leaves all k streams equal to it. Correcting only the pooled token would leave the streams’ contribution out of the input’s gradient, and under the carry-biased gate init MGR is built for that is most of it. Values are untouched in every case.

Every class embedding trains, at all three levels and on both sides of the cut: a network’s ClassTokens and this stack’s own ClassLatents ride the carry because it is taken after they are spliced, and a per-Layer latent below the cut gets a ghost row in the carry (value zero, taken from the tracked table). They are learnable input rows, not part of a layer’s transform — which is what stays undifferentiated below the cut. Anything else would leave a silently dead parameter.

K >= n_virtual behaves exactly like None, and so does any value at all off the autodiff backend. Honoured by Self::forward, Self::step and Self::prime alike, so a cut stack decodes under the same truncation it trains under.

Implementations§

§

impl<M> Layers<M>
where M: Block, <M as Block>::Options: Clone,

pub fn class_latent_output_indices(&self, orig_len: usize) -> Vec<usize>

Output positions of the stack-level class latents for an orig_len input.

A marker that never lands (a Custom at or past the end) reports a position past the emitted sequence — compare against its length.

pub fn n_virtual_count(&self) -> usize

Number of (virtual) layers this stack runs.

pub fn forward( &self, x: Tensor<3>, caches: Option<<M as Block>::Caches>, options: <M as Block>::Options, class: Option<&mut ClassCursors>, ) -> (Tensor<3>, <M as Block>::Caches)

Full-sequence pass through every (virtual) layer.

Layer returns only its delta — F_l = Block(RMSNorm(·)), plus the feed-forward sub-block’s contribution when the layer has one; the outer residual is added here. With Residuals::Standard each layer adds the input skip (unless suppressed). With Residuals::MultiGate the skip is dropped and up to n_stream parallel streams — seeded with x as the first one — carry the residual: each layer reads their attention-pooled aggregate as input, and its output either becomes a new stream (while fewer than n_stream exist) or is gated into every stream (see MultiGate).

ignore_first/last_residual apply to both paths: skipping the first restarts the residual carry from the first layer’s output (the input is read but not carried); skipping the last makes the stack output the last layer’s transform F_l alone (no input-dependent carry).

class places the stack-level and the per-layer class latents; None takes x for the whole sequence (so every latent lands in this call). Passing the same [ClassCursors] to consecutive chunks splits the sequence without moving a single latent — see [ClassCursors]. Both residual paths host them: a per-layer latent is spliced into the token sequence and, under MultiGate, into every carried stream too (the aggregator over the resulting identical streams reproduces the row, so the layer above reads it back exactly as the additive skip hands it on).

pub fn step( &self, x: Tensor<2>, caches: Option<<M as Block>::Caches>, class: Option<&mut ClassCursors>, ) -> (Tensor<2>, <M as Block>::Caches)

Single-token step through every (virtual) layer.

class drives two independent class-latent levels — the stack-level Self::class_latents (class.stack, spliced once below the first layer, exactly as in forward) and the per-Layer latents (class.per_layer[i], one cursor per virtual layer).

Because a layer’s class latents grow the sequence the next layer sees (exactly as in forward), a single user step is a cascade: the bottom input stream (the stack latents falling on this step, plus the user token) is threaded up the stack, each layer expanding it with its own class latents. Every layer’s recurrence therefore sees the same token order as forward, so forward and step agree.

The step returns the (fully propagated) output of the last token of that stream — the user token, unless an End latent (the one kind that closes the sequence rather than preceding a token) follows it. Latents emitted before the user token are stepped for their effect on the state alone.

None injects nothing at either level (and Middle/End latents panic, as they do without a [ClassCursors::full_len] hint).

Self::grad_horizon applies here exactly as in Self::forward, on the same virtual layers, so a stack decodes under the truncation it trains under. Note the cut rebuilds an inner-backend view of the stack once per call, which is per token here rather than per sequence — negligible against a training step, but not something to leave set for plain decoding (where it is inert anyway, the model then being off the autodiff backend).

pub fn prime( &self, batch: usize, caches: Option<<M as Block>::Caches>, class: Option<&mut ClassCursors>, ) -> (Option<Tensor<2>>, Option<<M as Block>::Caches>)

Step the class latents the stack has waiting for its next user token — with no user token, so nothing but class data is consumed.

This is Self::step’s opening half on its own: the stack-level latents due now open the bottom stream (empty when none are), which then goes up the stack through the very cascade step runs, with every layer additionally flushing the latents its next token was going to be preceded by. A prime followed by a step therefore runs exactly the sequence that step alone would have. End latents are never primed: closing the sequence, they belong to the step carrying its last user token (which is why that step returns them). A cursor already at the announced end therefore primes nothing.

Returns the fully propagated output of the last latent emitted, or None when none were waiting — the seedless-generation entry point: prime → sample → step → sample → … batch sizes the latent rows, which are the only inputs there are.

The caches come back as they went in when nothing ran (None included); a partly primed stack is completed with zero caches for the layers that stepped nothing, which is exactly the state they hold. None cursors inject nothing at all (Middle/End latents then panic, as in step).

pub fn step_infinite(&self, x: Tensor<2>) -> Tensor<2>

Stationary fixed point of the whole stack under a constant token, with no caches involved: under a constant input each layer’s output converges (its decay damps the transient, and the readout phase of the rotation cancels), so the downstream layer’s input converges too and the limit composes exactly, layer by layer — even though every layer’s SSM state keeps rotating forever. Residual handling mirrors Self::step; cursorless (class latents are not injected).