Md Saiful Islam Sajol

Upload Lattice_Spike project with full LFS weights

e2c9d42 verified 2 months ago

10.3 kB

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

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