Blame for boxes.py

"""
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()