# Copyright 2026 Modilify # SPDX-License-Identifier: LicenseRef-Modilify-Open-Model-1.0 """Fixed-shape latent deliberation state for Modilify Mk1 decoding. The state contains no vocabulary-sized tensors. Per-canvas information stays in a compact latent space so iterative diffusion does not retain one logits allocation per denoise pass. """ from __future__ import annotations from dataclasses import dataclass import math import torch from torch import nn from torch.nn import functional as F @dataclass class LatentDeliberationState: """Persistent, fixed-size state for one or more canvas episodes.""" token_latents: torch.Tensor memory_slots: torch.Tensor confidence: torch.Tensor entropy: torch.Tensor age: torch.Tensor token_changed: torch.Tensor confidence_delta: torch.Tensor entropy_delta: torch.Tensor ponder_steps: torch.Tensor stagnation_steps: torch.Tensor @classmethod def empty( cls, *, batch_size: int, canvas_length: int, latent_dim: int, memory_slots: int, device: torch.device, dtype: torch.dtype, ) -> "LatentDeliberationState": """Create a zero-initialized recurrent state. Args: batch_size: Number of independent sequences. canvas_length: Number of rolling canvas positions. latent_dim: Width of each latent token and memory slot. memory_slots: Number of persistent memory slots. device: Allocation device. dtype: Floating-point dtype for latent tensors. Returns: A zero-initialized state with integer progress clocks. """ return cls( token_latents=torch.zeros( batch_size, canvas_length, latent_dim, device=device, dtype=dtype ), memory_slots=torch.zeros( batch_size, memory_slots, latent_dim, device=device, dtype=dtype ), confidence=torch.zeros( batch_size, canvas_length, device=device, dtype=torch.float32 ), entropy=torch.zeros( batch_size, canvas_length, device=device, dtype=torch.float32 ), age=torch.zeros( batch_size, canvas_length, device=device, dtype=torch.int32 ), token_changed=torch.zeros( batch_size, canvas_length, device=device, dtype=torch.float32 ), confidence_delta=torch.zeros( batch_size, canvas_length, device=device, dtype=torch.float32 ), entropy_delta=torch.zeros( batch_size, canvas_length, device=device, dtype=torch.float32 ), ponder_steps=torch.zeros(batch_size, device=device, dtype=torch.int32), stagnation_steps=torch.zeros(batch_size, device=device, dtype=torch.int32), ) def shift( self, committed: int, *, entropy_fill_value: float = 0.0 ) -> "LatentDeliberationState": """Drop committed canvas positions without changing long-term memory. Args: committed: Number of leading canvas positions to drop. entropy_fill_value: Fill value for newly exposed entropy slots. Returns: A shifted state with unchanged memory slots. """ canvas_length = self.token_latents.shape[1] if not 0 <= committed <= canvas_length: raise ValueError("`committed` must be in [0, canvas_length].") if committed == 0: return self def shifted(tensor: torch.Tensor, fill_value: float | int = 0) -> torch.Tensor: result = torch.full_like(tensor, fill_value) if committed < canvas_length: result[:, : canvas_length - committed] = tensor[:, committed:] return result return LatentDeliberationState( token_latents=shifted(self.token_latents), memory_slots=self.memory_slots.clone(), confidence=shifted(self.confidence), entropy=shifted(self.entropy, entropy_fill_value), age=shifted(self.age), token_changed=shifted(self.token_changed), confidence_delta=shifted(self.confidence_delta), entropy_delta=shifted(self.entropy_delta), ponder_steps=torch.zeros_like(self.ponder_steps), stagnation_steps=torch.zeros_like(self.stagnation_steps), ) def advance_trajectory_clocks( ponder_steps: torch.Tensor, stagnation_steps: torch.Tensor, *, commit_lengths: torch.LongTensor, active_rows: torch.BoolTensor, progress_scores: torch.Tensor, min_progress: float, ) -> tuple[torch.IntTensor, torch.IntTensor]: """Advance useful-ponder and true-stagnation clocks for each row. Args: ponder_steps: Current useful-ponder clocks, shape ``[batch]``. stagnation_steps: Current stagnation clocks, shape ``[batch]``. commit_lengths: Tokens committed this step. active_rows: Rows that are still generating. progress_scores: Signed fused-risk improvement. min_progress: Minimum improvement counted as progress. Returns: Updated ponder and stagnation clocks. """ if min_progress < 0: raise ValueError("`min_progress` must be non-negative.") if not ( ponder_steps.shape == stagnation_steps.shape == commit_lengths.shape == active_rows.shape == progress_scores.shape ): raise ValueError("Trajectory clock inputs must share shape [batch].") committed = commit_lengths.gt(0) waiting = active_rows & ~committed improving = progress_scores.ge(min_progress) next_ponder = torch.where( committed, torch.zeros_like(ponder_steps), ponder_steps + waiting.to(torch.int32) ) next_stagnation = torch.where( committed, torch.zeros_like(stagnation_steps), torch.where( waiting & improving, torch.zeros_like(stagnation_steps), stagnation_steps + waiting.to(torch.int32), ), ) return next_ponder.to(torch.int32), next_stagnation.to(torch.int32) def should_force_trajectory_jump( ponder_steps: torch.Tensor, stagnation_steps: torch.Tensor, *, max_ponder_steps: int, stagnation_threshold: int, ) -> torch.BoolTensor: """Return whether a row has exhausted ponder or stagnation budget.""" if max_ponder_steps <= 0 or stagnation_threshold <= 0: raise ValueError("Trajectory jump limits must be positive.") return ponder_steps.ge(max_ponder_steps) | stagnation_steps.ge(stagnation_threshold) class _TemporalTransformerCell(nn.Module): """One-step recurrent token update with fixed-slot memory attention.""" def __init__( self, latent_dim: int, num_heads: int, dropout: float, local_attention_window: int, ) -> None: super().__init__() self.state_norm = nn.LayerNorm(latent_dim) self.observation_norm = nn.LayerNorm(latent_dim) self.memory_address_norm = nn.LayerNorm(latent_dim) self.memory_value_norm = nn.LayerNorm(latent_dim) self.temporal_update = nn.Linear(2 * latent_dim, 2 * latent_dim) self.local_attention = nn.MultiheadAttention( latent_dim, num_heads, dropout=dropout, batch_first=True ) self.local_attention_window = local_attention_window self.register_buffer("_local_attention_mask", torch.empty(0), persistent=False) self.token_memory_attention = nn.MultiheadAttention( latent_dim, num_heads, dropout=dropout, batch_first=True ) self.memory_token_attention = nn.MultiheadAttention( latent_dim, num_heads, dropout=dropout, batch_first=True ) self.token_ff_norm = nn.LayerNorm(latent_dim) self.memory_ff_norm = nn.LayerNorm(latent_dim) self.stored_token_norm = nn.LayerNorm(latent_dim) self.stored_memory_norm = nn.LayerNorm(latent_dim) expansion = latent_dim * 4 self.token_ff = nn.Sequential( nn.Linear(latent_dim, expansion), nn.SiLU(), nn.Linear(expansion, latent_dim), ) self.memory_ff = nn.Sequential( nn.Linear(latent_dim, expansion), nn.SiLU(), nn.Linear(expansion, latent_dim), ) def forward( self, previous_tokens: torch.Tensor, observation: torch.Tensor, memory: torch.Tensor, memory_slot_identity: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: """Update token latents and persistent memory for one cell. Slot identity is an addressing key only. After the memory residual is normalized, the same identity is added back so persistent content stays slot-specific and invariant to canvas commits. Args: previous_tokens: Previous token latents, shape ``[batch, canvas, dim]``. observation: Current observation, same shape as ``previous_tokens``. memory: Persistent memory slots, shape ``[batch, slots, dim]``. memory_slot_identity: Scaled slot addresses, same shape as ``memory``. Returns: Updated token latents and memory slots. """ gate_logits, candidate = self.temporal_update( torch.cat( (self.state_norm(previous_tokens), self.observation_norm(observation)), dim=-1, ) ).chunk(2, dim=-1) gate = torch.sigmoid(gate_logits) tokens = gate * previous_tokens + (1.0 - gate) * torch.nn.functional.silu(candidate) if ( self._local_attention_mask.shape != (tokens.shape[1], tokens.shape[1]) or self._local_attention_mask.device != tokens.device or self._local_attention_mask.dtype != tokens.dtype ): positions = torch.arange(tokens.shape[1], device=tokens.device) allowed = (positions[:, None] - positions[None, :]).abs() < self.local_attention_window self._local_attention_mask = torch.zeros( tokens.shape[1], tokens.shape[1], device=tokens.device, dtype=tokens.dtype ).masked_fill(~allowed, torch.finfo(tokens.dtype).min) local_update, _ = self.local_attention( self.state_norm(tokens), self.state_norm(tokens), self.state_norm(tokens), attn_mask=self._local_attention_mask, need_weights=False, ) tokens = tokens + local_update addressed_memory = self.memory_address_norm(memory + memory_slot_identity) memory_values = self.memory_value_norm(memory) token_memory_update, _ = self.token_memory_attention( self.state_norm(tokens), addressed_memory, memory_values, need_weights=False ) tokens = tokens + token_memory_update tokens = tokens + self.token_ff(self.token_ff_norm(tokens)) memory_token_update, _ = self.memory_token_attention( addressed_memory, self.state_norm(tokens), self.state_norm(tokens), need_weights=False, ) memory = memory + memory_token_update memory = memory + self.memory_ff(self.memory_ff_norm(memory)) return ( self.stored_token_norm(tokens), self.stored_memory_norm(memory) + memory_slot_identity, ) class LatentDeliberationTransformer(nn.Module): """Small recurrent Transformer that compresses repeated denoise context.""" def __init__( self, *, hidden_size: int, latent_dim: int = 512, memory_slots: int = 16, num_layers: int = 2, num_heads: int = 8, local_attention_window: int = 32, dropout: float = 0.0, ) -> None: super().__init__() if latent_dim % num_heads: raise ValueError("`latent_dim` must be divisible by `num_heads`.") if local_attention_window <= 0: raise ValueError("`local_attention_window` must be positive.") self.hidden_size = hidden_size self.latent_dim = latent_dim self.memory_slots = memory_slots self.heavy_projection = nn.Linear(hidden_size, latent_dim, bias=False) self.embedding_projection = nn.Linear(hidden_size, latent_dim, bias=False) self.scalar_projection = nn.Linear(11, latent_dim, bias=False) self.blocks = nn.ModuleList( [ _TemporalTransformerCell( latent_dim, num_heads, dropout, local_attention_window ) for _ in range(num_layers) ] ) self.output_norm = nn.LayerNorm(latent_dim) self.output_projection = nn.Linear(latent_dim, hidden_size, bias=False) self.memory_slot_identity = nn.Parameter(torch.empty(memory_slots, latent_dim)) self.reset_memory_slot_identity() @torch.no_grad() def reset_memory_slot_identity(self) -> None: """Restore orthonormal slot addresses after generic initialization.""" workspace = torch.empty_like(self.memory_slot_identity, dtype=torch.float32) if self.memory_slots <= self.latent_dim: nn.init.orthogonal_(workspace) else: nn.init.normal_(workspace, mean=0.0, std=1.0) workspace = F.normalize(workspace, dim=-1) self.memory_slot_identity.copy_(workspace.to(dtype=self.memory_slot_identity.dtype)) def scaled_memory_slot_identity( self, *, batch_size: int, device: torch.device, dtype: torch.dtype, ) -> torch.Tensor: """Return unit directions scaled to LayerNorm RMS. Args: batch_size: Number of sequences to broadcast over. device: Output device. dtype: Output dtype. Returns: Slot identities of shape ``[batch, slots, latent_dim]``. """ identity = F.normalize(self.memory_slot_identity.float(), dim=-1) identity = identity * math.sqrt(self.latent_dim) return identity.to(device=device, dtype=dtype).unsqueeze(0).expand( batch_size, -1, -1 ) def project_context(self, token_latents: torch.Tensor) -> torch.Tensor: """Translate latent state into a self-conditioning embedding.""" return self.output_projection(self.output_norm(token_latents)) def forward( self, *, heavy_hidden: torch.Tensor, token_embeddings: torch.Tensor, confidence: torch.Tensor, entropy: torch.Tensor, state: LatentDeliberationState, ) -> tuple[torch.Tensor, LatentDeliberationState]: """Advance latent memory and produce decoder self-conditioning. Args: heavy_hidden: Hidden states from the previous decoder pass. token_embeddings: Embeddings of current noisy canvas tokens. confidence: Proposal confidence for each canvas position. entropy: Proposal entropy for each canvas position. state: Persistent latent state from the preceding pass. Returns: Self-conditioning embeddings and the next compact latent state. """ if heavy_hidden.ndim != 3: raise ValueError("`heavy_hidden` must have shape [batch, canvas, hidden].") if heavy_hidden.shape != token_embeddings.shape: raise ValueError("`heavy_hidden` and `token_embeddings` must have the same shape.") batch_size, canvas_length, hidden_size = heavy_hidden.shape if hidden_size != self.hidden_size: raise ValueError("Unexpected hidden size for latent deliberation.") expected_state = (batch_size, canvas_length, self.latent_dim) if state.token_latents.shape != expected_state: raise ValueError("State token latents do not match the current canvas.") if state.memory_slots.shape != (batch_size, self.memory_slots, self.latent_dim): raise ValueError("State memory slots do not match this module.") if state.age.dtype is not torch.int32: raise TypeError("Latent deliberation ages must use int32.") scalars = torch.stack( ( confidence.to(dtype=heavy_hidden.dtype), entropy.to(dtype=heavy_hidden.dtype).log1p(), state.age.to(dtype=heavy_hidden.dtype).clamp_max(32767).log1p(), torch.linspace( -1.0, 1.0, canvas_length, device=heavy_hidden.device, dtype=heavy_hidden.dtype, ) .unsqueeze(0) .expand(batch_size, -1), state.token_changed.to(dtype=heavy_hidden.dtype), state.confidence_delta.to(dtype=heavy_hidden.dtype), state.entropy_delta.to(dtype=heavy_hidden.dtype).sign() * state.entropy_delta.to(dtype=heavy_hidden.dtype).abs().log1p(), state.ponder_steps.to(dtype=heavy_hidden.dtype).log1p()[:, None].expand( -1, canvas_length ), state.stagnation_steps.to(dtype=heavy_hidden.dtype).log1p()[:, None].expand( -1, canvas_length ), confidence.to(dtype=heavy_hidden.dtype) * torch.exp(-entropy.to(dtype=heavy_hidden.dtype).clamp_min(0.0)), state.confidence_delta.to(dtype=heavy_hidden.dtype).clamp_min(0.0) + (-state.entropy_delta.to(dtype=heavy_hidden.dtype)).clamp_min(0.0).log1p(), ), dim=-1, ) observation = ( self.heavy_projection(heavy_hidden) + self.embedding_projection(token_embeddings) + self.scalar_projection(scalars) ) tokens = state.token_latents memory = state.memory_slots slot_identity = self.scaled_memory_slot_identity( batch_size=batch_size, device=memory.device, dtype=memory.dtype, ) for block in self.blocks: tokens, memory = block(tokens, observation, memory, slot_identity) observation = tokens next_state = LatentDeliberationState( token_latents=tokens, memory_slots=memory, confidence=confidence.to(dtype=torch.float32), entropy=entropy.to(dtype=torch.float32), age=state.age, token_changed=state.token_changed, confidence_delta=state.confidence_delta, entropy_delta=state.entropy_delta, ponder_steps=state.ponder_steps, stagnation_steps=state.stagnation_steps, ) return self.project_context(tokens), next_state __all__ = [ "LatentDeliberationState", "LatentDeliberationTransformer", "advance_trajectory_clocks", "should_force_trajectory_jump", ]