""" Box-embedding encoding layer and geometric losses. A box B = [c_min, c_max] is an axis-aligned hyperrectangle in R^d. The encoder maps an LLM hidden state at each position to a box, forming a per-prefix lattice element. We parameterise boxes by their min corner and a strictly positive width (via softplus) so that c_max = c_min + width > c_min always holds. This module provides: * BoxEncoder : hidden state -> (c_min, c_max) * BoxBatch : container with volume / log-volume helpers * geometric_signals : leakage residual r_t, contraction rate kappa_t * containment_loss : hinge loss enforcing B_{t+1} subset B_t * volume_loss : prevents collapse / unbounded growth * anchored_contraction_loss : penalises contraction faster than info gain * concretization_residual : Channel 6, prediction-vs-box in vocabulary space Tensor conventions: boxes are (B, T, d); B = batch, T = time (tokens). """ from __future__ import annotations from dataclasses import dataclass import torch from torch import nn, Tensor import torch.nn.functional as F EPS = 1e-6 @dataclass class BoxBatch: """A batch of per-step boxes. Shapes: (B, T, d).""" c_min: Tensor c_max: Tensor @property def width(self) -> Tensor: return (self.c_max - self.c_min).clamp_min(EPS) @property def log_volume(self) -> Tensor: """log V_t = sum_j log(width_j). Shape (B, T).""" return torch.log(self.width).sum(dim=-1) @property def volume(self) -> Tensor: """V_t in (0, inf). Clamped for stability. Shape (B, T).""" return torch.exp(self.log_volume.clamp(-60.0, 20.0)) @property def modulation_scale(self) -> Tensor: """Dimension-normalized box scale used by the dynamics layer. A raw d-dimensional product volume shrinks exponentially with ``d`` and drove the RSNN time constant into its lower clamp for every token. The geometric mean side length preserves the box-size ordering while staying on a dimension-independent scale. """ return torch.exp(self.log_volume / self.width.shape[-1]) class BoxEncoder(nn.Module): """Map a hidden state to box corners. Output corners are squashed into the unit hypercube [0, 1]^d so that the learned lattice is nested inside the root box B_0 = [0, 1]^d. """ def __init__( self, d_model: int, d_box: int = 256, hidden: int = 512, architecture: str = "mlp", ): super().__init__() self.d_box = d_box self.architecture = architecture if architecture == "mlp": self.net = nn.Sequential( nn.Linear(d_model, hidden), nn.GELU(), nn.Linear(hidden, 2 * d_box), ) output_layer = self.net[-1] elif architecture == "linear": self.net = nn.Linear(d_model, 2 * d_box) output_layer = self.net else: raise ValueError("BoxEncoder architecture must be 'mlp' or 'linear'") # Initialise so the first boxes are a slight contraction of the root box. with torch.no_grad(): output_layer.bias[:d_box].fill_(0.05) # c_min near 0 output_layer.bias[d_box:].fill_(1.5) # pre-softplus width -> large def forward(self, h: Tensor) -> BoxBatch: out = self.net(h) base, width_raw = out[..., : self.d_box], out[..., self.d_box:] c_min = torch.sigmoid(base) # in (0, 1) width = F.softplus(width_raw) # > 0 c_max = torch.clamp(c_min + width, max=1.0) c_max = torch.maximum(c_max, c_min + EPS) # guarantee positive width return BoxBatch(c_min=c_min, c_max=c_max) # ---------------------------------------------------------------------------- # Geometric signals (Channels 1 & 2) # ---------------------------------------------------------------------------- def geometric_signals(box: BoxBatch) -> tuple[Tensor, Tensor]: """Return (leakage residual r_t, contraction rate kappa_t). r_t : (B, T, 2d) containment violation of B_t inside B_{t-1}. kappa_t : (B, T) log-volume contraction, >= 0 for a shrinking box. Both signals are aligned to the current causal token position: step zero is zero-padded because it has no predecessor, and step t uses only B_{t-1} and B_t. In particular, r_t must never read B_{t+1}, whose cached hidden state can already reflect token t. """ c_min, c_max = box.c_min, box.c_max # The adjacent-pair residual belongs to the child/current step, not the # parent/previous step. child_min, parent_min = c_min[:, 1:], c_min[:, :-1] child_max, parent_max = c_max[:, 1:], c_max[:, :-1] leak_lo = F.relu(parent_min - child_min) # child_min should be >= parent_min leak_hi = F.relu(child_max - parent_max) # child_max should be <= parent_max r = torch.cat([leak_lo, leak_hi], dim=-1) # (B, T-1, 2d) r = F.pad(r, (0, 0, 1, 0)) # first step has no predecessor logV = box.log_volume # (B, T) kappa = (logV[:, :-1] - logV[:, 1:]) # (B, T-1), >0 when shrinking kappa = F.pad(kappa, (1, 0)) # first step has no predecessor return r, kappa # ---------------------------------------------------------------------------- # Concretization (Channel 6) # ---------------------------------------------------------------------------- def concretization_residual( box: BoxBatch, topk_probs: Tensor, topk_embeddings: Tensor, ) -> Tensor: """Escaped probability mass m_t (Channel 6). Implements m_t = sum_v p_{t+1}(v) * dist(e_v, B_t) over the top-k tokens of the model's next-token distribution, where dist(e, B) = || relu(c_min - e) + relu(e - c_max) ||_1. Args: box : boxes B_t, corners (B, T, d). topk_probs : (B, T, K) probabilities of the top-K next tokens. topk_embeddings : (B, T, K, d) those tokens' embeddings in box space. Returns: m_t : (B, T) escaped mass weighted by L1 distance outside the box. """ c_min = box.c_min.unsqueeze(2) # (B, T, 1, d) c_max = box.c_max.unsqueeze(2) below = F.relu(c_min - topk_embeddings) above = F.relu(topk_embeddings - c_max) dist = (below + above).sum(dim=-1) # (B, T, K) L1 distance m = (topk_probs * dist).sum(dim=-1) # (B, T) return m # ---------------------------------------------------------------------------- # Losses # ---------------------------------------------------------------------------- def containment_loss(box: BoxBatch) -> Tensor: """Hinge loss enforcing B_{t+1} subset B_t (paper Eq. 1). Containment requires, elementwise, child_min >= parent_min and child_max <= parent_max, where child = step t+1, parent = step t. Each violation is penalised linearly by its magnitude. """ c_min, c_max = box.c_min, box.c_max lo = F.relu(c_min[:, :-1] - c_min[:, 1:]) # parent_min - child_min (>0 = violation) hi = F.relu(c_max[:, 1:] - c_max[:, :-1]) # child_max - parent_max (>0 = violation) return (lo.sum(dim=-1) + hi.sum(dim=-1)).mean() def box_intersection(a: BoxBatch, b: BoxBatch) -> BoxBatch: """Lattice meet (intersection) of two aligned box batches. The meet of boxes A and B is the largest box contained in both: meet_min = max(A.min, B.min) (elementwise) meet_max = min(A.max, B.max) When A and B are disjoint along a dimension, meet_max < meet_min there; the returned BoxBatch clamps such widths to EPS (BoxBatch.width), so the meet volume is a lower bound. `a` and `b` must share the same (B, T, d) shape. """ meet_min = torch.maximum(a.c_min, b.c_min) meet_max = torch.minimum(a.c_max, b.c_max) return BoxBatch(c_min=meet_min, c_max=torch.maximum(meet_max, meet_min + EPS)) def join_loss(box: BoxBatch) -> Tensor: """Volumetric containment via the meet of consecutive boxes. A soft, volume-based complement to `containment_loss` (a corner-wise hinge). For each adjacent pair (parent B_t, child B_{t+1}) we form the intersection box and penalise the child log-volume that the meet fails to capture: L_join = mean relu( logV(child) - logV(meet(parent, child)) ). Under perfect nesting B_{t+1} subset B_t the meet equals the child and the term is zero; when the child leaks outside the parent the meet shrinks along the leaking dimensions and the penalty grows. Working in log-volume makes the term scale-free across dimensions and numerically stable. """ parent = BoxBatch(c_min=box.c_min[:, :-1], c_max=box.c_max[:, :-1]) child = BoxBatch(c_min=box.c_min[:, 1:], c_max=box.c_max[:, 1:]) meet = box_intersection(parent, child) leaked = F.relu(child.log_volume - meet.log_volume) # (B, T-1) return leaked.mean() def volume_loss(box: BoxBatch, lam_collapse: float = 1.0, lam_expand: float = 1.0) -> Tensor: """Penalise both collapse (zero volume) and unbounded growth. The paper's first term exp(-sum log width) is a barrier against collapse; we use the numerically-stable mean form exp(-mean log width) and add a linear volume penalty. Both are averaged over (B, T). """ log_width = torch.log(box.width) # (B, T, d) collapse = torch.exp((-log_width.mean(dim=-1)).clamp(max=30.0)) # (B, T) expand = box.volume # (B, T) return lam_collapse * collapse.mean() + lam_expand * expand.mean() def anchored_contraction_loss(kappa: Tensor, kl_mean: Tensor, gamma: Tensor) -> Tensor: """L_anchor = mean relu(kappa_t - gamma * kl_mean_t). Penalises box contraction that is not justified by incoming distributional information (the knowledge-gap / confabulation signature). Args: kappa : (B, T) contraction rate. kl_mean : (B, T) mean layer-wise KL divergence (incoming surprise). gamma : scalar tensor (learned, kept positive by the caller). """ return F.relu(kappa - gamma * kl_mean).mean()