modilify/Modilify-Mk1-preview
120
1# Copyright 2026 Modilify2# SPDX-License-Identifier: LicenseRef-Modilify-Open-Model-1.03"""Fixed-shape latent deliberation state for Modilify Mk1 decoding.4 5The state deliberately contains no vocabulary-sized tensors. Keeping the6per-canvas information in a small latent space prevents iterative diffusion7rollouts from retaining one logits/probability allocation per denoise pass.8"""9 10from __future__ import annotations11 12from dataclasses import dataclass13 14import torch15from torch import nn16 17 18@dataclass19class LatentDeliberationState:20 """Persistent, fixed-size state for one or more canvas episodes."""21 22 token_latents: torch.Tensor23 memory_slots: torch.Tensor24 confidence: torch.Tensor25 entropy: torch.Tensor26 age: torch.Tensor27 token_changed: torch.Tensor28 confidence_delta: torch.Tensor29 entropy_delta: torch.Tensor30 ponder_steps: torch.Tensor31 stagnation_steps: torch.Tensor32 33 @classmethod34 def empty(35 cls,36 *,37 batch_size: int,38 canvas_length: int,39 latent_dim: int,40 memory_slots: int,41 device: torch.device,42 dtype: torch.dtype,43 ) -> "LatentDeliberationState":44 """Create a zero-initialized recurrent state.45 46 Args:47 batch_size: Number of independent sequences.48 canvas_length: Number of rolling canvas positions.49 latent_dim: Width of each latent token and memory slot.50 memory_slots: Number of persistent memory slots.51 device: Allocation device.52 dtype: Floating-point dtype for latent tensors.53 54 Returns:55 A zero-initialized state with integer progress clocks.56 """57 58 return cls(59 token_latents=torch.zeros(60 batch_size, canvas_length, latent_dim, device=device, dtype=dtype61 ),62 memory_slots=torch.zeros(63 batch_size, memory_slots, latent_dim, device=device, dtype=dtype64 ),65 confidence=torch.zeros(66 batch_size, canvas_length, device=device, dtype=torch.float3267 ),68 entropy=torch.zeros(69 batch_size, canvas_length, device=device, dtype=torch.float3270 ),71 age=torch.zeros(72 batch_size, canvas_length, device=device, dtype=torch.int3273 ),74 token_changed=torch.zeros(75 batch_size, canvas_length, device=device, dtype=torch.float3276 ),77 confidence_delta=torch.zeros(78 batch_size, canvas_length, device=device, dtype=torch.float3279 ),80 entropy_delta=torch.zeros(81 batch_size, canvas_length, device=device, dtype=torch.float3282 ),83 ponder_steps=torch.zeros(batch_size, device=device, dtype=torch.int32),84 stagnation_steps=torch.zeros(batch_size, device=device, dtype=torch.int32),85 )86 87 88def advance_trajectory_clocks(89 ponder_steps: torch.Tensor,90 stagnation_steps: torch.Tensor,91 *,92 commit_lengths: torch.LongTensor,93 active_rows: torch.BoolTensor,94 progress_scores: torch.Tensor,95 min_progress: float,96) -> tuple[torch.IntTensor, torch.IntTensor]:97 """Advance useful-ponder and true-stagnation clocks for each row.98 99 Args:100 ponder_steps: Total waiting steps for each row.101 stagnation_steps: Consecutive non-improving steps for each row.102 commit_lengths: Number of committed tokens for each row.103 active_rows: Rows that are still generating.104 progress_scores: Signed fused-risk improvements.105 min_progress: Smallest improvement that resets stagnation.106 107 Returns:108 Updated ponder and stagnation counters.109 """110 111 if min_progress < 0:112 raise ValueError("`min_progress` must be non-negative.")113 if not (114 ponder_steps.shape == stagnation_steps.shape == commit_lengths.shape115 == active_rows.shape == progress_scores.shape116 ):117 raise ValueError("Trajectory clock inputs must share shape [batch].")118 committed = commit_lengths.gt(0)119 waiting = active_rows & ~committed120 improving = progress_scores.ge(min_progress)121 next_ponder = torch.where(122 committed, torch.zeros_like(ponder_steps), ponder_steps + waiting.to(torch.int32)123 )124 next_stagnation = torch.where(125 committed,126 torch.zeros_like(stagnation_steps),127 torch.where(128 waiting & improving,129 torch.zeros_like(stagnation_steps),130 stagnation_steps + waiting.to(torch.int32),131 ),132 )133 return next_ponder.to(torch.int32), next_stagnation.to(torch.int32)134 135 136def should_force_trajectory_jump(137 ponder_steps: torch.Tensor,138 stagnation_steps: torch.Tensor,139 *,140 max_ponder_steps: int,141 stagnation_threshold: int,142) -> torch.BoolTensor:143 """Return rows that exhausted either inference progress clock.144 145 Args:146 ponder_steps: Total waiting steps for each row.147 stagnation_steps: Consecutive non-improving steps for each row.148 max_ponder_steps: Maximum allowed waiting steps.149 stagnation_threshold: Maximum consecutive stagnation steps.150 151 Returns:152 Boolean mask selecting rows that must use a forced jump.153 """154 155 if max_ponder_steps <= 0 or stagnation_threshold <= 0:156 raise ValueError("Trajectory jump limits must be positive.")157 return ponder_steps.ge(max_ponder_steps) | stagnation_steps.ge(stagnation_threshold)158 159 160class _TemporalTransformerCell(nn.Module):161 """One-step recurrent token update with fixed-slot memory attention."""162 163 def __init__(164 self, latent_dim: int, num_heads: int, dropout: float,165 local_attention_window: int,166 ) -> None:167 super().__init__()168 self.state_norm = nn.LayerNorm(latent_dim)169 self.observation_norm = nn.LayerNorm(latent_dim)170 # A slot's learned identity is only used for attention addressing. The171 # recurrent state itself remains pure memory content so commit shifts172 # cannot accidentally write positional identity into persistent state.173 self.memory_address_norm = nn.LayerNorm(latent_dim)174 self.memory_value_norm = nn.LayerNorm(latent_dim)175 self.temporal_update = nn.Linear(2 * latent_dim, 2 * latent_dim)176 self.local_attention = nn.MultiheadAttention(177 latent_dim, num_heads, dropout=dropout, batch_first=True178 )179 self.local_attention_window = local_attention_window180 self.register_buffer("_local_attention_mask", torch.empty(0), persistent=False)181 self.token_memory_attention = nn.MultiheadAttention(182 latent_dim, num_heads, dropout=dropout, batch_first=True183 )184 self.memory_token_attention = nn.MultiheadAttention(185 latent_dim, num_heads, dropout=dropout, batch_first=True186 )187 self.token_ff_norm = nn.LayerNorm(latent_dim)188 self.memory_ff_norm = nn.LayerNorm(latent_dim)189 self.stored_token_norm = nn.LayerNorm(latent_dim)190 self.stored_memory_norm = nn.LayerNorm(latent_dim)191 expansion = latent_dim * 4192 self.token_ff = nn.Sequential(193 nn.Linear(latent_dim, expansion),194 nn.SiLU(),195 nn.Linear(expansion, latent_dim),196 )197 self.memory_ff = nn.Sequential(198 nn.Linear(latent_dim, expansion),199 nn.SiLU(),200 nn.Linear(expansion, latent_dim),201 )202 203 def forward(204 self,205 previous_tokens: torch.Tensor,206 observation: torch.Tensor,207 memory: torch.Tensor,208 memory_slot_identity: torch.Tensor,209 ) -> tuple[torch.Tensor, torch.Tensor]:210 gate_logits, candidate = self.temporal_update(211 torch.cat(212 (self.state_norm(previous_tokens), self.observation_norm(observation)),213 dim=-1,214 )215 ).chunk(2, dim=-1)216 gate = torch.sigmoid(gate_logits)217 tokens = gate * previous_tokens + (1.0 - gate) * torch.nn.functional.silu(candidate)218 if (219 self._local_attention_mask.shape != (tokens.shape[1], tokens.shape[1])220 or self._local_attention_mask.device != tokens.device221 or self._local_attention_mask.dtype != tokens.dtype222 ):223 positions = torch.arange(tokens.shape[1], device=tokens.device)224 allowed = (225 positions[:, None] - positions[None, :]226 ).abs() < self.local_attention_window227 self._local_attention_mask = torch.zeros(228 tokens.shape[1], tokens.shape[1], device=tokens.device, dtype=tokens.dtype229 ).masked_fill(~allowed, torch.finfo(tokens.dtype).min)230 local_update, _ = self.local_attention(231 self.state_norm(tokens), self.state_norm(tokens), self.state_norm(tokens),232 attn_mask=self._local_attention_mask, need_weights=False,233 )234 tokens = tokens + local_update235 236 addressed_memory = self.memory_address_norm(memory + memory_slot_identity)237 memory_values = self.memory_value_norm(memory)238 token_memory_update, _ = self.token_memory_attention(239 self.state_norm(tokens), addressed_memory, memory_values, need_weights=False240 )241 tokens = tokens + token_memory_update242 tokens = tokens + self.token_ff(self.token_ff_norm(tokens))243 244 memory_token_update, _ = self.memory_token_attention(245 addressed_memory,246 self.state_norm(tokens),247 self.state_norm(tokens),248 need_weights=False,249 )250 memory = memory + memory_token_update251 memory = memory + self.memory_ff(self.memory_ff_norm(memory))252 # This module is a recurrent cell, not a depth-only Transformer block.253 # Persist normalized state so repeated denoise updates cannot accumulate254 # an unbounded residual magnitude across time.255 return self.stored_token_norm(tokens), self.stored_memory_norm(memory)256 257 258class LatentDeliberationTransformer(nn.Module):259 """Small recurrent Transformer that compresses repeated denoise context."""260 261 def __init__(262 self,263 *,264 hidden_size: int,265 latent_dim: int = 512,266 memory_slots: int = 16,267 num_layers: int = 2,268 num_heads: int = 8,269 local_attention_window: int = 32,270 dropout: float = 0.0,271 ) -> None:272 super().__init__()273 if latent_dim % num_heads:274 raise ValueError("`latent_dim` must be divisible by `num_heads`.")275 if local_attention_window <= 0:276 raise ValueError("`local_attention_window` must be positive.")277 self.hidden_size = hidden_size278 self.latent_dim = latent_dim279 self.memory_slots = memory_slots280 self.heavy_projection = nn.Linear(hidden_size, latent_dim, bias=False)281 self.embedding_projection = nn.Linear(hidden_size, latent_dim, bias=False)282 self.scalar_projection = nn.Linear(11, latent_dim, bias=False)283 self.blocks = nn.ModuleList(284 [285 _TemporalTransformerCell(286 latent_dim, num_heads, dropout, local_attention_window287 )288 for _ in range(num_layers)289 ]290 )291 self.output_norm = nn.LayerNorm(latent_dim)292 self.output_projection = nn.Linear(latent_dim, hidden_size, bias=False)293 self.memory_slot_identity = nn.Parameter(torch.empty(memory_slots, latent_dim))294 self.reset_memory_slot_identity()295 296 @torch.no_grad()297 def reset_memory_slot_identity(self) -> None:298 """Restore learned memory addresses after generic initialization."""299 300 nn.init.normal_(self.memory_slot_identity, mean=0.0, std=0.02)301 302 def project_context(self, token_latents: torch.Tensor) -> torch.Tensor:303 """Translate latent state into a self-conditioning embedding."""304 305 normalized_tokens = self.output_norm(token_latents)306 return self.output_projection(normalized_tokens)307 308 def forward(309 self,310 *,311 heavy_hidden: torch.Tensor,312 token_embeddings: torch.Tensor,313 confidence: torch.Tensor,314 entropy: torch.Tensor,315 state: LatentDeliberationState,316 ) -> tuple[torch.Tensor, LatentDeliberationState]:317 """Advance latent memory and produce decoder self-conditioning.318 319 Args:320 heavy_hidden: Hidden states from the previous decoder pass.321 token_embeddings: Embeddings of current noisy canvas tokens.322 confidence: Proposal confidence for each canvas position.323 entropy: Proposal entropy for each canvas position.324 state: Persistent latent state from the preceding pass.325 326 Returns:327 Self-conditioning embeddings and the next compact latent state.328 """329 330 if heavy_hidden.ndim != 3:331 raise ValueError("`heavy_hidden` must have shape [batch, canvas, hidden].")332 if heavy_hidden.shape != token_embeddings.shape:333 raise ValueError("`heavy_hidden` and `token_embeddings` must have the same shape.")334 batch_size, canvas_length, hidden_size = heavy_hidden.shape335 if hidden_size != self.hidden_size:336 raise ValueError("Unexpected hidden size for latent deliberation.")337 expected_state = (batch_size, canvas_length, self.latent_dim)338 if state.token_latents.shape != expected_state:339 raise ValueError("State token latents do not match the current canvas.")340 if state.memory_slots.shape != (batch_size, self.memory_slots, self.latent_dim):341 raise ValueError("State memory slots do not match this module.")342 if state.age.dtype is not torch.int32:343 raise TypeError("Latent deliberation ages must use int32.")344 345 scalars = torch.stack(346 (347 confidence.to(dtype=heavy_hidden.dtype),348 entropy.to(dtype=heavy_hidden.dtype).log1p(),349 state.age.to(dtype=heavy_hidden.dtype).clamp_max(32767).log1p(),350 torch.linspace(351 -1.0, 1.0, canvas_length, device=heavy_hidden.device,352 dtype=heavy_hidden.dtype,353 ).unsqueeze(0).expand(batch_size, -1),354 state.token_changed.to(dtype=heavy_hidden.dtype),355 state.confidence_delta.to(dtype=heavy_hidden.dtype),356 state.entropy_delta.to(dtype=heavy_hidden.dtype).sign()357 * state.entropy_delta.to(dtype=heavy_hidden.dtype).abs().log1p(),358 state.ponder_steps.to(dtype=heavy_hidden.dtype).log1p()[:, None]359 .expand(-1, canvas_length),360 state.stagnation_steps.to(dtype=heavy_hidden.dtype).log1p()[:, None]361 .expand(-1, canvas_length),362 confidence.to(dtype=heavy_hidden.dtype)363 * torch.exp(-entropy.to(dtype=heavy_hidden.dtype).clamp_min(0.0)),364 state.confidence_delta.to(dtype=heavy_hidden.dtype).clamp_min(0.0)365 + (-state.entropy_delta.to(dtype=heavy_hidden.dtype)).clamp_min(0.0).log1p(),366 ),367 dim=-1,368 )369 observation = (370 self.heavy_projection(heavy_hidden)371 + self.embedding_projection(token_embeddings)372 + self.scalar_projection(scalars)373 )374 tokens = state.token_latents375 memory = state.memory_slots376 slot_identity = self.memory_slot_identity.to(device=memory.device, dtype=memory.dtype)377 slot_identity = slot_identity.unsqueeze(0).expand(batch_size, -1, -1)378 for block in self.blocks:379 tokens, memory = block(tokens, observation, memory, slot_identity)380 observation = tokens381 382 context = self.project_context(tokens)383 next_state = LatentDeliberationState(384 token_latents=tokens,385 memory_slots=memory,386 confidence=confidence.to(dtype=torch.float32),387 entropy=entropy.to(dtype=torch.float32),388 age=state.age,389 token_changed=state.token_changed,390 confidence_delta=state.confidence_delta,391 entropy_delta=state.entropy_delta,392 ponder_steps=state.ponder_steps,393 stagnation_steps=state.stagnation_steps,394 )395 return context, next_state396 397 398__all__ = [399 "LatentDeliberationState", "LatentDeliberationTransformer",400 "advance_trajectory_clocks", "should_force_trajectory_jump",401]402