| |
| |
| """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", |
| ] |
|
|
