diff options
Diffstat (limited to 'models')
| -rw-r--r-- | models/srm/hrm_orth_v1.py | 371 |
1 files changed, 371 insertions, 0 deletions
diff --git a/models/srm/hrm_orth_v1.py b/models/srm/hrm_orth_v1.py new file mode 100644 index 0000000..70ec561 --- /dev/null +++ b/models/srm/hrm_orth_v1.py @@ -0,0 +1,371 @@ +"""HRM-Orth v1 — orthogonal patch of HRM per codex round 2 recommendation. + +CORE IDEA (codex Q6 pivot, after pure-orthogonal retract Q1): +Keep HRM's H_level/L_level/ACT structure, just patch the inner Block: + - Attention → cosine-normalized attention (≈ Lipschitz-bounded) + - SwiGLU MLP → CayleyOrth linear + MaxMin + CayleyOrth linear + - rms_norm + add → weighted residual: h_new = (1-σ(w)) · h + σ(w) · f(h) + - "Weak orthogonality": diag(s) scaling with most s≈1, some s∈[0.90, 0.97] for compression + +Per codex Q5 decomp: target +5~+7pp over SRM v1 (0.39 → 0.43-0.46). +Per codex Q3: Cayley used (we have it from srm_aol_v1); Householder would be faster but more impl. +""" +from typing import Tuple, List, Dict, Optional +from dataclasses import dataclass +import math + +import torch +import torch.nn.functional as F +from torch import nn +from pydantic import BaseModel + +from models.common import trunc_normal_init_ +from models.layers import rms_norm, SwiGLU, Attention, RotaryEmbedding, CosSin, CastedEmbedding, CastedLinear +from models.sparse_embedding import CastedSparseEmbedding +from models.srm.srm_aol_v1 import CayleyOrthogonal + + +def maxmin(x: torch.Tensor, group: int = 2) -> torch.Tensor: + """1-Lipschitz norm-preserving activation (Anil et al. 2019 GroupSort). + + Pairs adjacent dims; outputs (min, max) per pair. Permutation a.e. → ||∇|| = 1. + Strictly better than ReLU under norm constraints (no rank-kill). + """ + *prefix, d = x.shape + if d % group != 0: + pad = group - (d % group) + x = F.pad(x, (0, pad)) + d = d + pad + xg = x.reshape(*prefix, d // group, group) + sorted_vals, _ = xg.sort(dim=-1) + return sorted_vals.reshape(*prefix, d) + + +def cosine_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, + tau: float = 8.0) -> torch.Tensor: + """Cosine-normalized softmax attention. Approximately Lipschitz-bounded + (exact bound depends on tau and value norms — see LipsFormer Qi 2023).""" + q = F.normalize(q, dim=-1) + k = F.normalize(k, dim=-1) + attn = (q @ k.transpose(-2, -1)) * tau + attn = attn.softmax(dim=-1) + return attn @ v + + +class OrthLinear(nn.Module): + """Orthogonal linear layer via Cayley. Allows optional row-scaling diag(s) + where s_i ∈ [s_min, 1] to introduce 'weak orthogonality' (codex Q1 fix). + + If s_min < 1, the operator is contractive in some directions: + Lip = max(s) ≤ 1, det = prod(s) ≤ 1 (weak contraction in compressing modes) + """ + def __init__(self, dim: int, s_min: float = 0.95, learn_scale: bool = True): + super().__init__() + self.Q = CayleyOrthogonal(dim) + self.s_min = s_min + # diag scale: sigmoid -> [s_min, 1] + if learn_scale and s_min < 1.0: + self.log_s_raw = nn.Parameter(torch.zeros(dim)) # init sigmoid(0)=0.5 → scale=(s_min+1)/2 + else: + self.register_buffer("log_s_raw", torch.zeros(dim)) + self.learn_scale = learn_scale + + def scale_diag(self) -> torch.Tensor: + if self.s_min >= 1.0 or not self.learn_scale: + return torch.ones_like(self.log_s_raw) + # Affine map sigmoid → [s_min, 1] + return self.s_min + (1.0 - self.s_min) * torch.sigmoid(self.log_s_raw) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + Q = self.Q() # (d, d) orthogonal + s = self.scale_diag().to(Q.dtype) # (d,) in [s_min, 1] + Qs = Q * s.unsqueeze(0) # rescale columns + return F.linear(x, Qs) + + +@dataclass +class HierarchicalReasoningModel_ACTV1InnerCarry: + z_H: torch.Tensor + z_L: torch.Tensor + + +@dataclass +class HierarchicalReasoningModel_ACTV1Carry: + inner_carry: HierarchicalReasoningModel_ACTV1InnerCarry + + steps: torch.Tensor + halted: torch.Tensor + + current_data: Dict[str, torch.Tensor] + + +class HierarchicalReasoningModel_ACTV1Config(BaseModel): + batch_size: int + seq_len: int + puzzle_emb_ndim: int = 0 + num_puzzle_identifiers: int + vocab_size: int + + H_cycles: int + L_cycles: int + + H_layers: int + L_layers: int + + # Transformer config + hidden_size: int + expansion: float + num_heads: int + pos_encodings: str + + rms_norm_eps: float = 1e-5 + rope_theta: float = 10000.0 + + # Halting Q-learning config + halt_max_steps: int + halt_exploration_prob: float + + forward_dtype: str = "bfloat16" + + +class HierarchicalReasoningModel_ACTV1Block(nn.Module): + """Orthogonal-patched HRM Block. + + Replaces (attn + SwiGLU + rms_norm) with (cosine attn + Orth-MLP + weighted residual). + The original class name preserved so the ReasoningModule wrapper is unchanged. + """ + def __init__(self, config: HierarchicalReasoningModel_ACTV1Config) -> None: + super().__init__() + d = config.hidden_size + s_min = getattr(config, "orth_s_min", 0.95) + cosine_tau = getattr(config, "cosine_attn_tau", 8.0) + + # Lipschitz-bounded cosine attention: orthogonal Q/K/V/O projections + self.q_proj = OrthLinear(d, s_min=1.0, learn_scale=False) # strict orth for q/k + self.k_proj = OrthLinear(d, s_min=1.0, learn_scale=False) + self.v_proj = OrthLinear(d, s_min=s_min, learn_scale=True) # weak orth on values + self.o_proj = OrthLinear(d, s_min=s_min, learn_scale=True) + self.cosine_tau = cosine_tau + + # Orth-MLP: OrthLinear -> MaxMin -> OrthLinear (no expansion; uses original d) + self.mlp_in = OrthLinear(d, s_min=s_min, learn_scale=True) + self.mlp_out = OrthLinear(d, s_min=s_min, learn_scale=True) + + # Weighted residual gates (init sigmoid(0)=0.5 → balanced residual) + self.w_attn_logit = nn.Parameter(torch.zeros(())) + self.w_mlp_logit = nn.Parameter(torch.zeros(())) + + def forward(self, cos_sin: CosSin, hidden_states: torch.Tensor) -> torch.Tensor: + # Cosine attention + q = self.q_proj(hidden_states) + k = self.k_proj(hidden_states) + v = self.v_proj(hidden_states) + attn_out = self.o_proj(cosine_attention(q, k, v, tau=self.cosine_tau)) + w_attn = torch.sigmoid(self.w_attn_logit) + hidden_states = (1.0 - w_attn) * hidden_states + w_attn * attn_out + + # Orth-MLP with MaxMin + mlp_out = self.mlp_out(maxmin(self.mlp_in(hidden_states), group=2)) + w_mlp = torch.sigmoid(self.w_mlp_logit) + hidden_states = (1.0 - w_mlp) * hidden_states + w_mlp * mlp_out + return hidden_states + + +class HierarchicalReasoningModel_ACTV1ReasoningModule(nn.Module): + def __init__(self, layers: List[HierarchicalReasoningModel_ACTV1Block]): + super().__init__() + + self.layers = torch.nn.ModuleList(layers) + + def forward(self, hidden_states: torch.Tensor, input_injection: torch.Tensor, **kwargs) -> torch.Tensor: + # Input injection (add) + hidden_states = hidden_states + input_injection + # Layers + for layer in self.layers: + hidden_states = layer(hidden_states=hidden_states, **kwargs) + + return hidden_states + + +class HierarchicalReasoningModel_ACTV1_Inner(nn.Module): + def __init__(self, config: HierarchicalReasoningModel_ACTV1Config) -> None: + super().__init__() + self.config = config + self.forward_dtype = getattr(torch, self.config.forward_dtype) + + # I/O + self.embed_scale = math.sqrt(self.config.hidden_size) + embed_init_std = 1.0 / self.embed_scale + + self.embed_tokens = CastedEmbedding(self.config.vocab_size, self.config.hidden_size, init_std=embed_init_std, cast_to=self.forward_dtype) + self.lm_head = CastedLinear(self.config.hidden_size, self.config.vocab_size, bias=False) + self.q_head = CastedLinear(self.config.hidden_size, 2, bias=True) + + self.puzzle_emb_len = -(self.config.puzzle_emb_ndim // -self.config.hidden_size) # ceil div + if self.config.puzzle_emb_ndim > 0: + # Zero init puzzle embeddings + self.puzzle_emb = CastedSparseEmbedding(self.config.num_puzzle_identifiers, self.config.puzzle_emb_ndim, + batch_size=self.config.batch_size, init_std=0, cast_to=self.forward_dtype) + + # LM Blocks + if self.config.pos_encodings == "rope": + self.rotary_emb = RotaryEmbedding(dim=self.config.hidden_size // self.config.num_heads, + max_position_embeddings=self.config.seq_len + self.puzzle_emb_len, + base=self.config.rope_theta) + elif self.config.pos_encodings == "learned": + self.embed_pos = CastedEmbedding(self.config.seq_len + self.puzzle_emb_len, self.config.hidden_size, init_std=embed_init_std, cast_to=self.forward_dtype) + else: + raise NotImplementedError() + + # Reasoning Layers + self.H_level = HierarchicalReasoningModel_ACTV1ReasoningModule(layers=[HierarchicalReasoningModel_ACTV1Block(self.config) for _i in range(self.config.H_layers)]) + self.L_level = HierarchicalReasoningModel_ACTV1ReasoningModule(layers=[HierarchicalReasoningModel_ACTV1Block(self.config) for _i in range(self.config.L_layers)]) + + # Initial states + self.H_init = nn.Buffer(trunc_normal_init_(torch.empty(self.config.hidden_size, dtype=self.forward_dtype), std=1), persistent=True) + self.L_init = nn.Buffer(trunc_normal_init_(torch.empty(self.config.hidden_size, dtype=self.forward_dtype), std=1), persistent=True) + + # Q head special init + # Init Q to (almost) zero for faster learning during bootstrapping + with torch.no_grad(): + self.q_head.weight.zero_() + self.q_head.bias.fill_(-5) # type: ignore + + def _input_embeddings(self, input: torch.Tensor, puzzle_identifiers: torch.Tensor): + # Token embedding + embedding = self.embed_tokens(input.to(torch.int32)) + + # Puzzle embeddings + if self.config.puzzle_emb_ndim > 0: + puzzle_embedding = self.puzzle_emb(puzzle_identifiers) + + pad_count = self.puzzle_emb_len * self.config.hidden_size - puzzle_embedding.shape[-1] + if pad_count > 0: + puzzle_embedding = F.pad(puzzle_embedding, (0, pad_count)) + + embedding = torch.cat((puzzle_embedding.view(-1, self.puzzle_emb_len, self.config.hidden_size), embedding), dim=-2) + + # Position embeddings + if self.config.pos_encodings == "learned": + # scale by 1/sqrt(2) to maintain forward variance + embedding = 0.707106781 * (embedding + self.embed_pos.embedding_weight.to(self.forward_dtype)) + + # Scale + return self.embed_scale * embedding + + def empty_carry(self, batch_size: int): + return HierarchicalReasoningModel_ACTV1InnerCarry( + z_H=torch.empty(batch_size, self.config.seq_len + self.puzzle_emb_len, self.config.hidden_size, dtype=self.forward_dtype), + z_L=torch.empty(batch_size, self.config.seq_len + self.puzzle_emb_len, self.config.hidden_size, dtype=self.forward_dtype), + ) + + def reset_carry(self, reset_flag: torch.Tensor, carry: HierarchicalReasoningModel_ACTV1InnerCarry): + return HierarchicalReasoningModel_ACTV1InnerCarry( + z_H=torch.where(reset_flag.view(-1, 1, 1), self.H_init, carry.z_H), + z_L=torch.where(reset_flag.view(-1, 1, 1), self.L_init, carry.z_L), + ) + + def forward(self, carry: HierarchicalReasoningModel_ACTV1InnerCarry, batch: Dict[str, torch.Tensor]) -> Tuple[HierarchicalReasoningModel_ACTV1InnerCarry, torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: + seq_info = dict( + cos_sin=self.rotary_emb() if hasattr(self, "rotary_emb") else None, + ) + + # Input encoding + input_embeddings = self._input_embeddings(batch["inputs"], batch["puzzle_identifiers"]) + + # Forward iterations + with torch.no_grad(): + z_H, z_L = carry.z_H, carry.z_L + + for _H_step in range(self.config.H_cycles): + for _L_step in range(self.config.L_cycles): + if not ((_H_step == self.config.H_cycles - 1) and (_L_step == self.config.L_cycles - 1)): + z_L = self.L_level(z_L, z_H + input_embeddings, **seq_info) + + if not (_H_step == self.config.H_cycles - 1): + z_H = self.H_level(z_H, z_L, **seq_info) + + assert not z_H.requires_grad and not z_L.requires_grad + + # 1-step grad + z_L = self.L_level(z_L, z_H + input_embeddings, **seq_info) + z_H = self.H_level(z_H, z_L, **seq_info) + + # LM Outputs + new_carry = HierarchicalReasoningModel_ACTV1InnerCarry(z_H=z_H.detach(), z_L=z_L.detach()) # New carry no grad + output = self.lm_head(z_H)[:, self.puzzle_emb_len:] + + # Q head + q_logits = self.q_head(z_H[:, 0]).to(torch.float32) + + return new_carry, output, (q_logits[..., 0], q_logits[..., 1]) + + +class HierarchicalReasoningModel_ACTV1(nn.Module): + """ACT wrapper.""" + + def __init__(self, config_dict: dict): + super().__init__() + self.config = HierarchicalReasoningModel_ACTV1Config(**config_dict) + self.inner = HierarchicalReasoningModel_ACTV1_Inner(self.config) + + @property + def puzzle_emb(self): + return self.inner.puzzle_emb + + def initial_carry(self, batch: Dict[str, torch.Tensor]): + batch_size = batch["inputs"].shape[0] + + return HierarchicalReasoningModel_ACTV1Carry( + inner_carry=self.inner.empty_carry(batch_size), # Empty is expected, it will be reseted in first pass as all sequences are halted. + + steps=torch.zeros((batch_size, ), dtype=torch.int32), + halted=torch.ones((batch_size, ), dtype=torch.bool), # Default to halted + + current_data={k: torch.empty_like(v) for k, v in batch.items()} + ) + + def forward(self, carry: HierarchicalReasoningModel_ACTV1Carry, batch: Dict[str, torch.Tensor]) -> Tuple[HierarchicalReasoningModel_ACTV1Carry, Dict[str, torch.Tensor]]: + # Update data, carry (removing halted sequences) + new_inner_carry = self.inner.reset_carry(carry.halted, carry.inner_carry) + + new_steps = torch.where(carry.halted, 0, carry.steps) + + new_current_data = {k: torch.where(carry.halted.view((-1, ) + (1, ) * (batch[k].ndim - 1)), batch[k], v) for k, v in carry.current_data.items()} + + # Forward inner model + new_inner_carry, logits, (q_halt_logits, q_continue_logits) = self.inner(new_inner_carry, new_current_data) + + outputs = { + "logits": logits, + "q_halt_logits": q_halt_logits, + "q_continue_logits": q_continue_logits + } + + with torch.no_grad(): + # Step + new_steps = new_steps + 1 + is_last_step = new_steps >= self.config.halt_max_steps + + halted = is_last_step + + # if training, and ACT is enabled + if self.training and (self.config.halt_max_steps > 1): + # Halt signal + # NOTE: During evaluation, always use max steps, this is to guarantee the same halting steps inside a batch for batching purposes + halted = halted | (q_halt_logits > q_continue_logits) + + # Exploration + min_halt_steps = (torch.rand_like(q_halt_logits) < self.config.halt_exploration_prob) * torch.randint_like(new_steps, low=2, high=self.config.halt_max_steps + 1) + + halted = halted & (new_steps >= min_halt_steps) + + # Compute target Q + # NOTE: No replay buffer and target networks for computing target Q-value. + # As batch_size is large, there're many parallel envs. + # Similar concept as PQN https://arxiv.org/abs/2407.04811 + next_q_halt_logits, next_q_continue_logits = self.inner(new_inner_carry, new_current_data)[-1] + + outputs["target_q_continue"] = torch.sigmoid(torch.where(is_last_step, next_q_halt_logits, torch.maximum(next_q_halt_logits, next_q_continue_logits))) + + return HierarchicalReasoningModel_ACTV1Carry(new_inner_carry, new_steps, halted, new_current_data), outputs |