| |
| |
| """Fixed-shape latent deliberation state for Modilify Mk1 decoding. |
| |
| PyTorch → MLX: |
| nn.LayerNorm mlx.nn.LayerNorm |
| nn.Linear mlx.nn.Linear |
| nn.MultiheadAttention in_proj BiasedMHA query/key/value_proj (split QKV) |
| nn.Sequential token_ff/memory_ff mlx.nn.Sequential (layers.0 / layers.2) |
| """ |
|
|
| from __future__ import annotations |
|
|
| from dataclasses import dataclass |
| import math |
|
|
| import mlx.core as mx |
| import mlx.nn as nn |
|
|
|
|
| @dataclass |
| class LatentDeliberationState: |
| """Persistent, fixed-size state for one or more canvas episodes.""" |
|
|
| token_latents: mx.array |
| memory_slots: mx.array |
| confidence: mx.array |
| entropy: mx.array |
| age: mx.array |
| token_changed: mx.array |
| confidence_delta: mx.array |
| entropy_delta: mx.array |
| ponder_steps: mx.array |
| stagnation_steps: mx.array |
|
|
| @classmethod |
| def empty( |
| cls, |
| *, |
| batch_size: int, |
| canvas_length: int, |
| latent_dim: int, |
| memory_slots: int, |
| dtype: mx.Dtype, |
| ) -> "LatentDeliberationState": |
| """Create a zero-initialized recurrent state.""" |
|
|
| return cls( |
| token_latents=mx.zeros((batch_size, canvas_length, latent_dim), dtype=dtype), |
| memory_slots=mx.zeros((batch_size, memory_slots, latent_dim), dtype=dtype), |
| confidence=mx.zeros((batch_size, canvas_length), dtype=mx.float32), |
| entropy=mx.zeros((batch_size, canvas_length), dtype=mx.float32), |
| age=mx.zeros((batch_size, canvas_length), dtype=mx.int32), |
| token_changed=mx.zeros((batch_size, canvas_length), dtype=mx.float32), |
| confidence_delta=mx.zeros((batch_size, canvas_length), dtype=mx.float32), |
| entropy_delta=mx.zeros((batch_size, canvas_length), dtype=mx.float32), |
| ponder_steps=mx.zeros((batch_size,), dtype=mx.int32), |
| stagnation_steps=mx.zeros((batch_size,), dtype=mx.int32), |
| ) |
|
|
| def shift( |
| self, committed: int, *, entropy_fill_value: float = 0.0 |
| ) -> "LatentDeliberationState": |
| """Drop committed canvas positions without changing long-term memory.""" |
|
|
| 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: mx.array, fill_value: float | int = 0) -> mx.array: |
| result = mx.full(tensor.shape, fill_value, dtype=tensor.dtype) |
| if committed < canvas_length: |
| kept = tensor[:, committed:] |
| result = mx.concatenate( |
| [kept, mx.full( |
| (tensor.shape[0], committed, *tensor.shape[2:]), |
| fill_value, |
| dtype=tensor.dtype, |
| )], |
| axis=1, |
| ) if tensor.ndim > 2 else mx.concatenate( |
| [ |
| kept, |
| mx.full( |
| (tensor.shape[0], committed), |
| fill_value, |
| dtype=tensor.dtype, |
| ), |
| ], |
| axis=1, |
| ) |
| return result |
|
|
| return LatentDeliberationState( |
| token_latents=shifted(self.token_latents), |
| memory_slots=self.memory_slots, |
| 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=mx.zeros_like(self.ponder_steps), |
| stagnation_steps=mx.zeros_like(self.stagnation_steps), |
| ) |
|
|
|
|
| def advance_trajectory_clocks( |
| ponder_steps: mx.array, |
| stagnation_steps: mx.array, |
| *, |
| commit_lengths: mx.array, |
| active_rows: mx.array, |
| progress_scores: mx.array, |
| min_progress: float, |
| ) -> tuple[mx.array, mx.array]: |
| """Advance useful-ponder and true-stagnation clocks for each row.""" |
|
|
| 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 > 0 |
| waiting = active_rows & (~committed) |
| improving = progress_scores >= min_progress |
| next_ponder = mx.where( |
| committed, |
| mx.zeros_like(ponder_steps), |
| ponder_steps + waiting.astype(mx.int32), |
| ) |
| next_stagnation = mx.where( |
| committed, |
| mx.zeros_like(stagnation_steps), |
| mx.where( |
| waiting & improving, |
| mx.zeros_like(stagnation_steps), |
| stagnation_steps + waiting.astype(mx.int32), |
| ), |
| ) |
| return next_ponder.astype(mx.int32), next_stagnation.astype(mx.int32) |
|
|
|
|
| def should_force_trajectory_jump( |
| ponder_steps: mx.array, |
| stagnation_steps: mx.array, |
| *, |
| max_ponder_steps: int, |
| stagnation_threshold: int, |
| ) -> mx.array: |
| """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 >= max_ponder_steps) | ( |
| stagnation_steps >= stagnation_threshold |
| ) |
|
|
|
|
| class BiasedMHA(nn.Module): |
| """Explicit QKV attention matching PyTorch MultiheadAttention weights.""" |
|
|
| def __init__(self, dims: int, num_heads: int) -> None: |
| super().__init__() |
| if dims % num_heads: |
| raise ValueError("`dims` must be divisible by `num_heads`.") |
| self.num_heads = num_heads |
| self.head_dim = dims // num_heads |
| self.scale = self.head_dim**-0.5 |
| self.query_proj = nn.Linear(dims, dims, bias=True) |
| self.key_proj = nn.Linear(dims, dims, bias=True) |
| self.value_proj = nn.Linear(dims, dims, bias=True) |
| self.out_proj = nn.Linear(dims, dims, bias=True) |
|
|
| def __call__( |
| self, |
| queries: mx.array, |
| keys: mx.array, |
| values: mx.array, |
| mask: mx.array | None = None, |
| ) -> mx.array: |
| batch, query_len, dims = queries.shape |
| key_len = keys.shape[1] |
| queries = self.query_proj(queries).reshape( |
| batch, query_len, self.num_heads, self.head_dim |
| ).transpose(0, 2, 1, 3) |
| keys = self.key_proj(keys).reshape( |
| batch, key_len, self.num_heads, self.head_dim |
| ).transpose(0, 2, 1, 3) |
| values = self.value_proj(values).reshape( |
| batch, key_len, self.num_heads, self.head_dim |
| ).transpose(0, 2, 1, 3) |
| attn_mask = None |
| if mask is not None: |
| attn_mask = mask.reshape(1, 1, mask.shape[-2], mask.shape[-1]) |
| output = mx.fast.scaled_dot_product_attention( |
| queries, keys, values, scale=self.scale, mask=attn_mask |
| ) |
| output = output.transpose(0, 2, 1, 3).reshape(batch, query_len, dims) |
| return self.out_proj(output) |
|
|
|
|
| def _silu(x: mx.array) -> mx.array: |
| return x * mx.sigmoid(x) |
|
|
|
|
| class _TemporalTransformerCell(nn.Module): |
| """One-step recurrent token update with fixed-slot memory attention.""" |
|
|
| def __init__( |
| self, |
| latent_dim: int, |
| num_heads: int, |
| 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 = BiasedMHA(latent_dim, num_heads) |
| self.local_attention_window = local_attention_window |
| self.token_memory_attention = BiasedMHA(latent_dim, num_heads) |
| self.memory_token_attention = BiasedMHA(latent_dim, num_heads) |
| 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 _local_mask(self, length: int, dtype: mx.Dtype) -> mx.array: |
| positions = mx.arange(length) |
| allowed = ( |
| mx.abs(positions[:, None] - positions[None, :]) < self.local_attention_window |
| ) |
| return mx.where( |
| allowed, |
| mx.zeros((length, length), dtype=dtype), |
| mx.full((length, length), mx.finfo(dtype).min, dtype=dtype), |
| ) |
|
|
| def __call__( |
| self, |
| previous_tokens: mx.array, |
| observation: mx.array, |
| memory: mx.array, |
| memory_slot_identity: mx.array, |
| ) -> tuple[mx.array, mx.array]: |
| gate_logits, candidate = mx.split( |
| self.temporal_update( |
| mx.concatenate( |
| ( |
| self.state_norm(previous_tokens), |
| self.observation_norm(observation), |
| ), |
| axis=-1, |
| ) |
| ), |
| 2, |
| axis=-1, |
| ) |
| gate = mx.sigmoid(gate_logits) |
| tokens = gate * previous_tokens + (1.0 - gate) * _silu(candidate) |
| local_mask = self._local_mask(tokens.shape[1], tokens.dtype) |
| normed_tokens = self.state_norm(tokens) |
| tokens = tokens + self.local_attention( |
| normed_tokens, normed_tokens, normed_tokens, mask=local_mask |
| ) |
|
|
| addressed_memory = self.memory_address_norm(memory + memory_slot_identity) |
| memory_values = self.memory_value_norm(memory) |
| normed_tokens = self.state_norm(tokens) |
| tokens = tokens + self.token_memory_attention( |
| normed_tokens, addressed_memory, memory_values |
| ) |
| tokens = tokens + self.token_ff(self.token_ff_norm(tokens)) |
|
|
| normed_tokens = self.state_norm(tokens) |
| memory = memory + self.memory_token_attention( |
| addressed_memory, normed_tokens, normed_tokens |
| ) |
| 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__() |
| del dropout |
| 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 = [ |
| _TemporalTransformerCell(latent_dim, num_heads, 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 = mx.zeros((memory_slots, latent_dim)) |
|
|
| def scaled_memory_slot_identity( |
| self, |
| *, |
| batch_size: int, |
| dtype: mx.Dtype, |
| ) -> mx.array: |
| """Return unit directions scaled to LayerNorm RMS.""" |
|
|
| identity = self.memory_slot_identity.astype(mx.float32) |
| identity = identity / mx.maximum( |
| mx.linalg.norm(identity, axis=-1, keepdims=True), 1.0e-12 |
| ) |
| identity = identity * math.sqrt(self.latent_dim) |
| return mx.broadcast_to( |
| identity.astype(dtype)[None, :, :], |
| (batch_size, self.memory_slots, self.latent_dim), |
| ) |
|
|
| def project_context(self, token_latents: mx.array) -> mx.array: |
| """Translate latent state into a self-conditioning embedding.""" |
|
|
| return self.output_projection(self.output_norm(token_latents)) |
|
|
| def __call__( |
| self, |
| *, |
| heavy_hidden: mx.array, |
| token_embeddings: mx.array, |
| confidence: mx.array, |
| entropy: mx.array, |
| state: LatentDeliberationState, |
| ) -> tuple[mx.array, LatentDeliberationState]: |
| """Advance latent memory and produce decoder self-conditioning.""" |
|
|
| 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.") |
|
|
| dtype = heavy_hidden.dtype |
| position = mx.linspace(-1.0, 1.0, canvas_length, dtype=dtype) |
| position = mx.broadcast_to(position[None, :], (batch_size, canvas_length)) |
| age = mx.minimum(state.age.astype(dtype), 32767).astype(dtype) |
| entropy_delta = state.entropy_delta.astype(dtype) |
| scalars = mx.stack( |
| ( |
| confidence.astype(dtype), |
| mx.log1p(entropy.astype(dtype)), |
| mx.log1p(age), |
| position, |
| state.token_changed.astype(dtype), |
| state.confidence_delta.astype(dtype), |
| mx.sign(entropy_delta) * mx.log1p(mx.abs(entropy_delta)), |
| mx.broadcast_to( |
| mx.log1p(state.ponder_steps.astype(dtype))[:, None], |
| (batch_size, canvas_length), |
| ), |
| mx.broadcast_to( |
| mx.log1p(state.stagnation_steps.astype(dtype))[:, None], |
| (batch_size, canvas_length), |
| ), |
| confidence.astype(dtype) |
| * mx.exp(-mx.maximum(entropy.astype(dtype), 0.0)), |
| mx.maximum(state.confidence_delta.astype(dtype), 0.0) |
| + mx.log1p(mx.maximum(-entropy_delta, 0.0)), |
| ), |
| axis=-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, |
| 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.astype(mx.float32), |
| entropy=entropy.astype(mx.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 |
|
|
