Upload diffusion_llm/model/mcts.py with huggingface_hub

diffusion_llm/model/mcts.py

@@ -0,0 +1,239 @@
+"""AAM Diffusion LLM — MCTS Reasoning Engine
+
+AlphaZero-style tree search for reasoning about narrative arrangement
+from graph evidence. AAM-specific: each node = a sentence arrangement,
+value = narrative coherence, policy = next arrangement step.
+"""
+
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass, field
+from typing import Optional, List, Dict, Any, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class MCTSConfig:
+    def __init__(
+        self,
+        num_simulations: int = 64,
+        c_puct: float = 1.5,
+        temperature: float = 1.0,
+        max_depth: int = 10,
+        use_value_network: bool = True,
+        use_progressive_widening: bool = True,
+        max_children: int = 8,
+    ) -> None:
+        self.num_simulations = num_simulations
+        self.c_puct = c_puct
+        self.temperature = temperature
+        self.max_depth = max_depth
+        self.use_value_network = use_value_network
+        self.use_progressive_widening = use_progressive_widening
+        self.max_children = max_children
+
+        if num_simulations <= 0:
+            raise ValueError(f"num_simulations must be positive, got {num_simulations}")
+        if c_puct <= 0:
+            raise ValueError(f"c_puct must be positive, got {c_puct}")
+        if max_depth <= 0:
+            raise ValueError(f"max_depth must be positive, got {max_depth}")
+
+
+@dataclass
+class MCTSNode:
+    state: Optional[torch.Tensor] = None
+    parent: Optional["MCTSNode"] = None
+    children: List["MCTSNode"] = field(default_factory=list)
+    visit_count: int = 0
+    total_value: float = 0.0
+    prior: float = 0.0
+    depth: int = 0
+    is_expanded: bool = False
+    action: Optional[int] = None
+    hidden_state: Optional[torch.Tensor] = None
+
+    @property
+    def q_value(self) -> float:
+        if self.visit_count == 0:
+            return 0.0
+        return self.total_value / self.visit_count
+
+    @property
+    def is_leaf(self) -> bool:
+        return not self.is_expanded
+
+    @property
+    def is_root(self) -> bool:
+        """Whether this node is the root."""
+        return self.parent is None
+
+
+class ValueNetwork(nn.Module):
+    """Evaluate narrative coherence of a state."""
+    def __init__(self, d_model: int, hidden_dim: int = 256) -> None:
+        super().__init__()
+        self.network = nn.Sequential(
+            nn.Linear(d_model, hidden_dim, bias=False),
+            nn.SiLU(),
+            nn.Linear(hidden_dim, hidden_dim // 2, bias=False),
+            nn.SiLU(),
+            nn.Linear(hidden_dim // 2, 1, bias=False),
+            nn.Tanh(),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.network(x)
+
+
+class PolicyNetwork(nn.Module):
+    """Suggest next arrangement step."""
+    def __init__(self, d_model: int, num_actions: int = 8) -> None:
+        super().__init__()
+        self.network = nn.Sequential(
+            nn.Linear(d_model, d_model // 2, bias=False),
+            nn.SiLU(),
+            nn.Linear(d_model // 2, num_actions, bias=False),
+        )
+        self.num_actions = num_actions
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.network(x)
+
+
+class MCTSReasoner(nn.Module):
+    """MCTS Reasoning Engine for AAM sentence arrangement."""
+
+    def __init__(
+        self,
+        d_model: int,
+        num_actions: int = 8,
+        config: Optional[MCTSConfig] = None,
+    ) -> None:
+        super().__init__()
+        self.d_model = d_model
+        self.num_actions = num_actions
+        self.config = config or MCTSConfig()
+
+        if self.config.use_value_network:
+            self.value_network = ValueNetwork(d_model)
+        else:
+            self.value_network = None
+
+        self.policy_network = PolicyNetwork(d_model, num_actions)
+
+        self.state_encoder = nn.Sequential(
+            nn.Linear(d_model, d_model, bias=False),
+            nn.SiLU(),
+            nn.Linear(d_model, d_model, bias=False),
+        )
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        num_simulations: Optional[int] = None,
+    ) -> Tuple[torch.Tensor, Dict[str, Any]]:
+        batch_size = x.shape[0]
+        n_sims = num_simulations or self.config.num_simulations
+
+        encoded_state = self.state_encoder(x)
+        # Pool over sequence dimension if 3D input
+        if encoded_state.dim() == 3:
+            pooled_state = encoded_state.mean(dim=1)  # (batch, d_model)
+        else:
+            pooled_state = encoded_state  # (batch, d_model)
+
+        policy_logits = self.policy_network(pooled_state)  # (batch, num_actions)
+        policy_probs = F.softmax(policy_logits / self.config.temperature, dim=-1)  # (batch, num_actions)
+
+        if self.value_network is not None:
+            root_value = self.value_network(pooled_state)  # (batch, 1)
+        else:
+            root_value = torch.zeros(batch_size, 1, device=x.device)
+
+        visit_counts = torch.zeros(batch_size, self.num_actions, device=x.device, dtype=torch.float32)
+        total_values = torch.zeros(batch_size, self.num_actions, device=x.device, dtype=torch.float32)
+
+        for sim_idx in range(n_sims):
+            ucb_scores = self._compute_ucb(visit_counts, total_values, policy_probs, n_sims)
+            selected_actions = ucb_scores.argmax(dim=-1)
+
+            if self.value_network is not None:
+                action_onehot = F.one_hot(selected_actions, self.num_actions).float()
+                action_proj = action_onehot @ self.policy_network.network[-1].weight
+                padding = torch.zeros(batch_size, self.d_model - action_proj.shape[-1], device=x.device)
+                action_embedding = torch.cat([action_proj, padding], dim=-1)
+                state_action = pooled_state + action_embedding
+                sim_values = self.value_network(state_action)
+            else:
+                sim_values = torch.rand(batch_size, 1, device=x.device) * 2 - 1
+
+            visit_counts.scatter_add_(1, selected_actions.unsqueeze(1),
+                torch.ones(batch_size, 1, device=x.device, dtype=visit_counts.dtype))
+            total_values.scatter_add_(1, selected_actions.unsqueeze(1), sim_values)
+
+        if self.config.temperature > 0:
+            action_probs = F.softmax(visit_counts.log() / self.config.temperature, dim=-1)
+            action_probs = torch.where(visit_counts > 0, action_probs, torch.zeros_like(action_probs))
+            row_sums = action_probs.sum(dim=-1, keepdim=True)
+            action_probs = torch.where(
+                row_sums > 1e-6,
+                action_probs / (row_sums + 1e-8),
+                torch.full_like(action_probs, 1.0 / self.num_actions),
+            )
+        else:
+            action_probs = F.one_hot(visit_counts.argmax(dim=-1), self.num_actions).float()
+
+        info = {
+            "total_simulations": n_sims,
+            "root_value": root_value.mean().item(),
+            "max_visit_count": visit_counts.max().item(),
+            "entropy": -(action_probs * (action_probs + 1e-8).log()).sum(-1).mean().item(),
+            "visit_distribution": visit_counts / (visit_counts.sum(-1, keepdim=True) + 1e-8),
+        }
+
+        return action_probs, info
+
+    def _compute_ucb(self, visit_counts, total_values, priors, total_simulations):
+        q_values = torch.where(
+            visit_counts > 0,
+            total_values / (visit_counts + 1e-8),
+            torch.zeros_like(total_values),
+        )
+        parent_visits = visit_counts.sum(dim=-1, keepdim=True)
+        exploration = self.config.c_puct * priors * torch.sqrt(parent_visits + 1) / (1 + visit_counts)
+        return q_values + exploration
+
+    def compute_thinking_budget(self, complexity_score: float, base_simulations: int = 16, max_simulations: int = 256) -> int:
+        """Compute number of MCTS simulations based on complexity.
+
+        Adaptive compute budget: more complex inputs get more simulations.
+
+        Args:
+            complexity_score: Complexity score [0, 1] from ThinkingToggle.
+            base_simulations: Minimum number of simulations.
+            max_simulations: Maximum number of simulations.
+
+        Returns:
+            Recommended number of simulations.
+        """
+        return int(base_simulations + (max_simulations - base_simulations) * (complexity_score ** 2))
+
+    def get_reasoning_summary(self, info: Dict[str, Any]) -> str:
+        """Summary of reasoning for logging.
+
+        Args:
+            info: Dictionary from forward output.
+
+        Returns:
+            Summary string.
+        """
+        return (
+            f"MCTS(sims={info['total_simulations']}, "
+            f"root_val={info['root_value']:.3f}, "
+            f"max_visits={info['max_visit_count']:.0f}, "
+            f"entropy={info['entropy']:.3f})"
+        )