| """ |
| Chimera 5.2 — Parcae Prelude / Loop / Coda controller. |
| |
| Same numerics as the previous draft but cleaner: |
| * Loop count is deterministic during training so gradient checkpointing |
| recompute is consistent. |
| * Backward truncation only retains gradients on the last ``n_loops // 2`` |
| iterations; earlier iterates are detached, mirroring the original |
| intuition while keeping the implementation in pure PyTorch. |
| * Adaptive early-exit during inference based on residual magnitude. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import torch |
| import torch.nn as nn |
|
|
|
|
| class ParcaeInjection(nn.Module): |
| """ZOH-stable diagonal injection: ``h' = exp(-Δ·A)·h + Δ·B·e``.""" |
|
|
| __constants__ = ["hidden_size"] |
|
|
| def __init__(self, hidden_size: int): |
| super().__init__() |
| self.hidden_size = int(hidden_size) |
| self.log_A = nn.Parameter(torch.zeros(self.hidden_size)) |
| self.log_A._no_weight_decay = True |
| self.B_raw = nn.Parameter(torch.randn(self.hidden_size) * 0.02) |
| self.delta = nn.Parameter(torch.full((self.hidden_size,), 0.5)) |
|
|
| def forward(self, h_prev: torch.Tensor, e: torch.Tensor) -> torch.Tensor: |
| A_bar = (-self.delta * self.log_A.exp()).exp() |
| B_bar = self.delta * self.B_raw |
| return A_bar * h_prev + B_bar * e |
|
|
|
|
| class ParcaeLoopController(nn.Module): |
| """Iterative refinement controller used by the looped trunk.""" |
|
|
| __constants__ = ["loop_min", "loop_max", "loop_default"] |
|
|
| def __init__(self, hidden_size: int, |
| loop_range: tuple = (1, 6), loop_default: int = 2, |
| adaptive_exit_threshold: float = 0.01, |
| spectral_radius_bound: float = 1.0): |
| super().__init__() |
| self.injection = ParcaeInjection(hidden_size) |
| self.loop_min, self.loop_max = int(loop_range[0]), int(loop_range[1]) |
| self.loop_default = int(loop_default) |
| self.exit_threshold = float(adaptive_exit_threshold) |
| self.e_norm = nn.LayerNorm(hidden_size) |
|
|
| def forward(self, prelude_output: torch.Tensor, loop_fn, |
| num_loops: int = None) -> torch.Tensor: |
| e = self.e_norm(prelude_output) |
| h = torch.zeros_like(e) |
| n_loops = int(num_loops) if num_loops is not None else self.loop_default |
| n_loops = max(self.loop_min, min(self.loop_max, n_loops)) |
|
|
| n_bwd = max(1, n_loops // 2) if self.training else n_loops |
|
|
| for t in range(n_loops): |
| h_new = loop_fn(self.injection(h, e)) |
| backprop = (not self.training) or (t >= n_loops - n_bwd) |
| h = h_new if backprop else h_new.detach() |
| if not self.training and t > 0: |
| if (h_new - h).abs().mean().item() < self.exit_threshold: |
| break |
| return h |
|
|
|
|
| __all__ = ["ParcaeInjection", "ParcaeLoopController"] |
|
|
