import math from typing import List, Tuple from dataclasses import dataclass import torch import torch.nn as nn import torch.nn.functional as F # --------------------------------------------------------------------------- # Utilities # --------------------------------------------------------------------------- def _trunc_normal_init(tensor: torch.Tensor, std: float = 1.0, lower: float = -2.0, upper: float = 2.0) -> torch.Tensor: with torch.no_grad(): if std == 0: tensor.zero_() else: sqrt2 = math.sqrt(2) a = math.erf(lower / sqrt2) b = math.erf(upper / sqrt2) z = (b - a) / 2 c = (2 * math.pi) ** -0.5 pdf_u = c * math.exp(-0.5 * lower ** 2) pdf_l = c * math.exp(-0.5 * upper ** 2) comp_std = std / math.sqrt( 1 - (upper * pdf_u - lower * pdf_l) / z - ((pdf_u - pdf_l) / z) ** 2 ) tensor.uniform_(a, b) tensor.erfinv_() tensor.mul_(sqrt2 * comp_std) tensor.clip_(lower * comp_std, upper * comp_std) return tensor def _find_multiple(a: int, b: int) -> int: return (-(a // -b)) * b def rms_norm(hidden_states: torch.Tensor, variance_epsilon: float) -> torch.Tensor: input_dtype = hidden_states.dtype hidden_states = hidden_states.to(torch.float32) variance = hidden_states.square().mean(-1, keepdim=True) hidden_states = hidden_states * torch.rsqrt(variance + variance_epsilon) return hidden_states.to(input_dtype) class CastedLinear(nn.Module): def __init__(self, in_features: int, out_features: int, bias: bool): super().__init__() self.weight = nn.Parameter( _trunc_normal_init(torch.empty(out_features, in_features), std=1.0 / (in_features ** 0.5)) ) self.bias = nn.Parameter(torch.zeros(out_features)) if bias else None def forward(self, x: torch.Tensor) -> torch.Tensor: return F.linear( x, self.weight.to(x.dtype), bias=self.bias.to(x.dtype) if self.bias is not None else None, ) class SwiGLU(nn.Module): def __init__(self, hidden_size: int, expansion: float): super().__init__() inter = _find_multiple(round(expansion * hidden_size * 2 / 3), 256) self.gate_up_proj = CastedLinear(hidden_size, inter * 2, bias=False) self.down_proj = CastedLinear(inter, hidden_size, bias=False) def forward(self, x): gate, up = self.gate_up_proj(x).chunk(2, dim=-1) return self.down_proj(F.silu(gate) * up) # --------------------------------------------------------------------------- # RTM block: token-mixing MLP + channel-mixing MLP with RMSNorm residuals # --------------------------------------------------------------------------- class RTMBlock(nn.Module): def __init__(self, hidden_size: int, seq_len: int, expansion: float, rms_norm_eps: float = 1e-5): super().__init__() self.mlp_t = SwiGLU(hidden_size=seq_len, expansion=expansion) self.mlp = SwiGLU(hidden_size=hidden_size, expansion=expansion) self.norm_eps = rms_norm_eps def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: hidden_states = hidden_states.transpose(1, 2) hidden_states = rms_norm( hidden_states + self.mlp_t(hidden_states), variance_epsilon=self.norm_eps, ) hidden_states = hidden_states.transpose(1, 2) hidden_states = rms_norm( hidden_states + self.mlp(hidden_states), variance_epsilon=self.norm_eps, ) return hidden_states class RTMReasoningModule(nn.Module): def __init__(self, layers: List[RTMBlock]): super().__init__() self.layers = nn.ModuleList(layers) def forward(self, hidden_states: torch.Tensor, input_injection: torch.Tensor) -> torch.Tensor: hidden_states = hidden_states + input_injection for layer in self.layers: hidden_states = layer(hidden_states) return hidden_states def _make_level(num_layers: int, hidden_size: int, seq_len: int, expansion: float, rms_norm_eps: float) -> RTMReasoningModule: return RTMReasoningModule([ RTMBlock(hidden_size, seq_len, expansion, rms_norm_eps) for _ in range(num_layers) ]) # --------------------------------------------------------------------------- # Recursive H/L cycle loop # --------------------------------------------------------------------------- @dataclass class RTMCarry: z_H: torch.Tensor z_L: torch.Tensor class RTMInner(nn.Module): def __init__(self, hidden_size: int, expansion: float, H_cycles: int, L_cycles: int, H_layers: int, L_layers: int, num_tokens: int, with_grad: bool = False, cycle_noise_std: float = 0.0, rms_norm_eps: float = 1e-5, forward_dtype: str = "float32"): super().__init__() self.hidden_size = hidden_size self.H_cycles = H_cycles self.L_cycles = L_cycles self.with_grad = with_grad self.cycle_noise_std = max(0.0, float(cycle_noise_std)) self.forward_dtype = getattr(torch, forward_dtype) self.total_seq_len = max(1, int(num_tokens)) self.L_level = _make_level(L_layers, hidden_size, self.total_seq_len, expansion, rms_norm_eps) self.H_init = nn.Buffer( _trunc_normal_init(torch.empty(hidden_size, dtype=self.forward_dtype), std=1), persistent=True, ) self.L_init = nn.Parameter( _trunc_normal_init(torch.empty(hidden_size, dtype=self.forward_dtype), std=1) ) def empty_carry(self, batch_size: int, device=None) -> RTMCarry: if device is None: device = self.H_init.device return RTMCarry( z_H=self.H_init.unsqueeze(0).unsqueeze(0).expand( batch_size, self.total_seq_len, -1), z_L=self.L_init.unsqueeze(0).unsqueeze(0).expand( batch_size, self.total_seq_len, -1), ) def forward(self, carry: RTMCarry, z_H_init: torch.Tensor ) -> Tuple[RTMCarry, torch.Tensor, List[torch.Tensor]]: z_H, z_L = carry.z_H, carry.z_L intermediates: List[torch.Tensor] = [] if self.with_grad: for _ in range(self.H_cycles): for _ in range(self.L_cycles): z_L = self.L_level(z_L, z_H + z_H_init) z_H = self.L_level(z_H, z_L) if self.training and self.cycle_noise_std > 0: z_H = z_H + torch.randn_like(z_H) * self.cycle_noise_std intermediates.append(z_H) else: with torch.no_grad(): for _ in range(self.H_cycles - 1): for _ in range(self.L_cycles): z_L = self.L_level(z_L, z_H + z_H_init) z_H = self.L_level(z_H, z_L) if self.training and self.cycle_noise_std > 0: z_H = z_H + torch.randn_like(z_H) * self.cycle_noise_std intermediates.append(z_H) for _ in range(self.L_cycles): z_L = self.L_level(z_L, z_H + z_H_init) z_H = self.L_level(z_H, z_L) if self.training and self.cycle_noise_std > 0: z_H = z_H + torch.randn_like(z_H) * self.cycle_noise_std intermediates.append(z_H) new_carry = RTMCarry(z_H=z_H.detach(), z_L=z_L.detach()) return new_carry, z_H, intermediates class _PixelNorm(nn.Module): def forward(self, x: torch.Tensor) -> torch.Tensor: return x * (x.square().mean(dim=1, keepdim=True) + 1e-8).rsqrt() class _EqualLinear(nn.Module): def __init__(self, in_dim: int, out_dim: int): super().__init__() self.linear = nn.Linear(in_dim, out_dim) self.linear.bias.data.zero_() def forward(self, x: torch.Tensor) -> torch.Tensor: return self.linear(x) def _coerce_noise_2d(input_tensor: torch.Tensor, code_dim: int, mapper_name: str) -> torch.Tensor: """Accept [B, D] or [B, 1, D] and return [B, D].""" if input_tensor.dim() == 2 and input_tensor.shape[1] == code_dim: return input_tensor if input_tensor.dim() == 3 and input_tensor.shape[1] == 1 and input_tensor.shape[2] == code_dim: return input_tensor.squeeze(1) raise ValueError( f"{mapper_name} expected input shape [B, {code_dim}] or [B, 1, {code_dim}], " f"got {tuple(input_tensor.shape)}" ) class RTMMappingNetwork(nn.Module): def __init__(self, code_dim: int, num_tokens: int = 1, H_cycles: int = 1, L_cycles: int = 1, H_layers: int = 2, L_layers: int = 2, hidden_size: int = 256, expansion: float = 4.0, refinement_steps: int = 1, with_grad: bool = False, cycle_noise_std: float = 0.0, rms_norm_eps: float = 1e-5, forward_dtype: str = "float32"): super().__init__() self.code_dim = code_dim self.refinement_steps = max(1, refinement_steps) total_seq_len = max(1, int(num_tokens)) if hidden_size <= 0: assert code_dim % total_seq_len == 0, ( f"code_dim={code_dim} must be divisible by num_tokens={total_seq_len} " f"when hidden_size is auto" ) hidden_size = code_dim // total_seq_len self.hidden_size = hidden_size self.trm = RTMInner( hidden_size=hidden_size, expansion=expansion, H_cycles=H_cycles, L_cycles=L_cycles, H_layers=H_layers, L_layers=L_layers, num_tokens=num_tokens, with_grad=with_grad, cycle_noise_std=cycle_noise_std, rms_norm_eps=rms_norm_eps, forward_dtype=forward_dtype, ) seq_len = self.trm.total_seq_len capacity = seq_len * hidden_size self.pixel_norm = _PixelNorm() self.mapper_direct = (capacity == code_dim) if not self.mapper_direct: self.z_to_seq = _EqualLinear(code_dim, capacity) self.seq_to_w = _EqualLinear(capacity, code_dim) def _z_H_to_w(self, z_H: torch.Tensor) -> torch.Tensor: flat = z_H.flatten(start_dim=1) if self.mapper_direct: return flat return self.seq_to_w(flat) def forward(self, input): if isinstance(input, (list, tuple)): input = input[0] input = _coerce_noise_2d(input, self.code_dim, self.__class__.__name__) B = input.shape[0] seq_len = self.trm.total_seq_len z_norm = self.pixel_norm(input) if self.mapper_direct: z_seq = z_norm.view(B, seq_len, self.hidden_size) else: z_seq = self.z_to_seq(z_norm).view(B, seq_len, self.hidden_size) carry = self.trm.empty_carry(B, device=input.device) for _ in range(self.refinement_steps): carry, z_H_out, _ = self.trm(carry, z_H_init=z_seq) return [self._z_H_to_w(z_H_out)] def forward_w_trajectory(self, input): """Decode every intermediate H state into a style vector. Returns ``[w_proj, w_after_cycle_1, ...]`` so callers can visualize how w is refined across the H/L cycles. """ if isinstance(input, (list, tuple)): input = input[0] input = _coerce_noise_2d(input, self.code_dim, self.__class__.__name__) B = input.shape[0] seq_len = self.trm.total_seq_len z_norm = self.pixel_norm(input) if self.mapper_direct: z_seq = z_norm.view(B, seq_len, self.hidden_size) else: z_seq = self.z_to_seq(z_norm).view(B, seq_len, self.hidden_size) trajectory = [self._z_H_to_w(z_seq)] carry = self.trm.empty_carry(B, device=input.device) for _ in range(self.refinement_steps): carry, _z_H_out, intermediates = self.trm(carry, z_H_init=z_seq) for zh in intermediates: trajectory.append(self._z_H_to_w(zh)) return trajectory