latent_deliberation.py

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