"""Recursive (looped-depth) decoder-only transformer. Three sections — entry / body / exit — where the body is a small stack of shared weights applied ``num_loops`` times per forward. The loop update is ``h_{t+1} = A * h_t + B * e + R(h_t + e)`` with per-channel injection gates A, B initialised to zero (so the loop starts as a weight-shared transformer stack on h+e). Optional cross-loop expert diversity for shared MoE routers via ``moe_diversity_factor``. """ from __future__ import annotations from dataclasses import dataclass from typing import List, Optional, Tuple, Dict, Any import torch import torch.nn as nn import torch.nn.functional as F from .lm_loss import ( lm_cross_entropy_from_logits, token_superposition_attention_mask, token_superposition_embeddings, ) from .baseline import ( BaselineConfig, RMSNorm, TransformerBlock, MoELayer, combine_lm_and_aux_loss, init_moe_router_weights, count_parameters, model_summary, ) @dataclass class RecursiveConfig(BaselineConfig): # Auto-derived from entry + body + exit in __post_init__. num_layers: int = 0 num_entry_layers: int = 2 num_body_layers: int = 4 num_exit_layers: int = 2 num_loops: int = 4 # Std of the random init for the per-channel A gate. 0 (default) # leaves the loop's residual mixing inert at step 0; small positive # values (e.g. 0.02) break that symmetry. B always starts at zero. body_gate_init_std: float = 0.0 def __post_init__(self): super().__post_init__() if self.num_body_layers <= 0 or self.num_loops <= 0: raise ValueError( "num_body_layers and num_loops must both be > 0; set " "num_entry_layers / num_exit_layers to 0 if you want a " "purely-body model." ) if self.body_gate_init_std < 0: raise ValueError("body_gate_init_std must be >= 0") self.num_layers = ( self.num_entry_layers + self.num_body_layers + self.num_exit_layers ) class RecursiveBlock(nn.Module): """One iteration of the body loop: ``h_{t+1} = A*h + B*e + R(h+e)``. The body's transformer blocks are reused ``num_loops`` times, so MoE layers carry per-loop bias rows and the cross-loop diversity term. """ def __init__(self, config: RecursiveConfig): super().__init__() self.blocks = nn.ModuleList([ TransformerBlock(config, num_loops=config.num_loops) for _ in range(config.num_body_layers) ]) if config.body_gate_init_std > 0: self.A = nn.Parameter( torch.randn(config.d_model) * config.body_gate_init_std ) else: self.A = nn.Parameter(torch.zeros(config.d_model)) self.B = nn.Parameter(torch.zeros(config.d_model)) def forward( self, h: torch.Tensor, e: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, is_causal: bool = True, loop_idx: int = 0, ) -> Tuple[torch.Tensor, torch.Tensor, List[Optional[torch.Tensor]]]: x = h + e aux_loss = torch.zeros((), device=x.device, dtype=x.dtype) topk_list: List[Optional[torch.Tensor]] = [] for block in self.blocks: x, block_aux, block_topk = block( x, attention_mask=attention_mask, is_causal=is_causal, loop_idx=loop_idx, ) aux_loss = aux_loss + block_aux topk_list.append(block_topk) h_next = self.A * h + self.B * e + x return h_next, aux_loss, topk_list class RecursiveTransformer(nn.Module): def __init__(self, config: RecursiveConfig): super().__init__() self.config = config self.token_emb = nn.Embedding(config.vocab_size, config.d_model) self.entry = nn.ModuleList([ TransformerBlock(config) for _ in range(config.num_entry_layers) ]) self.body = RecursiveBlock(config) self.exit = nn.ModuleList([ TransformerBlock(config) for _ in range(config.num_exit_layers) ]) self.final_norm = RMSNorm(config.d_model, eps=config.norm_eps) self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False) self.lm_head.weight = self.token_emb.weight self._init_weights() def _init_weights(self): # ``RecursiveBlock.A`` and ``.B`` stay at their zero init — they are # nn.Parameter (not Linear/Embedding) and so are skipped by this pass. for module in self.modules(): if isinstance(module, nn.Linear): torch.nn.init.normal_(module.weight, mean=0.0, std=0.02) if module.bias is not None: torch.nn.init.zeros_(module.bias) elif isinstance(module, nn.Embedding): torch.nn.init.normal_(module.weight, mean=0.0, std=0.02) init_moe_router_weights(self, self.config.router_init_std) def forward( self, input_ids: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, labels: Optional[torch.Tensor] = None, is_causal: bool = True, token_superposition_bag_size: int = 1, ) -> Dict[str, Any]: x = token_superposition_embeddings( self.token_emb, input_ids, token_superposition_bag_size, ) attention_mask = token_superposition_attention_mask( attention_mask, token_superposition_bag_size, ) aux_loss = torch.zeros((), device=input_ids.device, dtype=x.dtype) topk_indices_list: List[Optional[torch.Tensor]] = [] for layer in self.entry: x, layer_aux, layer_topk = layer( x, attention_mask=attention_mask, is_causal=is_causal ) aux_loss = aux_loss + layer_aux topk_indices_list.append(layer_topk) e = x h = torch.zeros_like(e) for loop_idx in range(self.config.num_loops): h, block_aux, block_topks = self.body( h, e, attention_mask=attention_mask, is_causal=is_causal, loop_idx=loop_idx, ) aux_loss = aux_loss + block_aux topk_indices_list.extend(block_topks) x = h for layer in self.exit: x, layer_aux, layer_topk = layer( x, attention_mask=attention_mask, is_causal=is_causal ) aux_loss = aux_loss + layer_aux topk_indices_list.append(layer_topk) x = self.final_norm(x) logits = self.lm_head(x) lm_loss: Optional[torch.Tensor] = None if labels is not None: lm_loss = lm_cross_entropy_from_logits( logits, labels, token_superposition_bag_size=token_superposition_bag_size, ignore_index=-100, ) loss = combine_lm_and_aux_loss( lm_loss, aux_loss if self.config.use_moe else None, self.training, ) return { "logits": logits, "loss": loss, "lm_loss": lm_loss, "aux_loss": aux_loss if self.config.use_moe else None, "topk_indices": topk_indices_list if self.config.use_moe else None, } def update_router_biases(self, topk_indices_list: List[Optional[torch.Tensor]]) -> None: """Apply DeepSeek-style bias updates. Index layout: [entry_0..E-1, loop_0.b_0..B-1, loop_1.b_0..B-1, ..., loop_{L-1}.b_0..B-1, exit_0..X-1] Each body block is updated once per parameter set with all its loop iterations grouped, so the cross-loop diversity term sees them together. """ if not self.config.use_moe: return n_entry = self.config.num_entry_layers n_body = self.config.num_body_layers n_loops = self.config.num_loops for i, layer in enumerate(self.entry): topk = topk_indices_list[i] if topk is not None and isinstance(layer.ffn, MoELayer): layer.ffn.update_bias(topk, loop_idx=0) body_offset = n_entry for r, block in enumerate(self.body.blocks): if not isinstance(block.ffn, MoELayer): continue topk_per_loop: List[torch.Tensor] = [] valid = True for l in range(n_loops): idx = body_offset + l * n_body + r topk = topk_indices_list[idx] if topk is None: valid = False break topk_per_loop.append(topk) if valid: block.ffn.update_bias_per_loop(topk_per_loop) exit_offset = n_entry + n_loops * n_body for i, layer in enumerate(self.exit): topk = topk_indices_list[exit_offset + i] if topk is not None and isinstance(layer.ffn, MoELayer): layer.ffn.update_bias(topk, loop_idx=0) @torch.no_grad() def get_balance_stats(self) -> Dict[str, float]: """One entry per parameter set — body sub-blocks appear once each (not ``num_loops`` times).""" if not self.config.use_moe: return {} stats: Dict[str, float] = {} def _record(name: str, ffn: nn.Module) -> None: if hasattr(ffn, "bias"): bias = ffn.bias stats[f"{name}_bias_mean"] = bias.abs().mean().item() stats[f"{name}_bias_max"] = bias.abs().max().item() for idx, layer in enumerate(self.entry): _record(f"entry{idx}", layer.ffn) for idx, block in enumerate(self.body.blocks): _record(f"body{idx}", block.ffn) for idx, layer in enumerate(self.exit): _record(f"exit{idx}", layer.ffn) return stats @torch.no_grad() def generate( self, input_ids: torch.Tensor, max_new_tokens: int = 100, temperature: float = 1.0, top_k: Optional[int] = None, attention_mask: Optional[torch.Tensor] = None, eos_token_id: Optional[int] = None, ) -> torch.Tensor: self.train(False) batch_size = input_ids.size(0) for _ in range(max_new_tokens): outputs = self.forward( input_ids, attention_mask=attention_mask, is_causal=True, ) logits = outputs["logits"][:, -1, :] / temperature if top_k is not None: v, _ = torch.topk(logits, min(top_k, logits.size(-1))) logits[logits < v[:, [-1]]] = -float("Inf") probs = F.softmax(logits, dim=-1) next_token = torch.multinomial(probs, num_samples=1) input_ids = torch.cat([input_ids, next_token], dim=-1) if attention_mask is not None: attention_mask = torch.cat([ attention_mask, torch.ones( (batch_size, 1), device=attention_mask.device, dtype=attention_mask.dtype, ), ], dim=-1) if eos_token_id is not None and (next_token == eos_token_id).all(): break return input_ids __all__ = [ "RecursiveConfig", "RecursiveBlock", "RecursiveTransformer", "count_parameters", "model_summary", ]