code/wal_tat/proxy.py

"""Training-only proxy codes for hard-forward ternary optimization."""
from __future__ import annotations

import torch
import torch.nn as nn
import torch.nn.functional as F


def soft_ternary_proxy(proxy: torch.Tensor, temperature: float) -> torch.Tensor:
    """Smooth three-level staircase with transitions at -0.5 and +0.5."""
    tau = max(float(temperature), 1e-4)
    return torch.sigmoid((proxy - 0.5) / tau) - torch.sigmoid(
        (-proxy - 0.5) / tau
    )


def soft_ternary_proxy_derivative(
    proxy: torch.Tensor, temperature: float
) -> torch.Tensor:
    """Analytic derivative of :func:`soft_ternary_proxy` with respect to proxy."""
    tau = max(float(temperature), 1e-4)
    positive = torch.sigmoid((proxy - 0.5) / tau)
    negative = torch.sigmoid((-proxy - 0.5) / tau)
    return (
        positive * (1.0 - positive) + negative * (1.0 - negative)
    ) / tau


class ProxyTernaryMatrix(nn.Module):
    """Optimize a scalar proxy per weight while always executing hard codes."""

    def __init__(
        self,
        codes: torch.Tensor,
        scales: torch.Tensor,
        *,
        compute_dtype: torch.dtype,
        temperature: float = 0.35,
        committed_mask: torch.Tensor | None = None,
        master_weight: torch.Tensor | None = None,
        fake_fp16_scale: bool = False,
        initial_proxy_magnitude: float = 1.0,
        initial_zero_proxy_boundary: float | None = None,
    ):
        super().__init__()
        if codes.ndim != 3 or scales.shape != codes.shape[:2]:
            raise ValueError("codes must be [out, groups, group_size] with matching scales")
        if not torch.all((codes >= -1) & (codes <= 1)):
            raise ValueError("codes must be ternary")
        if committed_mask is None:
            committed_mask = torch.ones(codes.shape[:2], dtype=torch.bool, device=codes.device)
        if committed_mask.shape != codes.shape[:2] or committed_mask.dtype != torch.bool:
            raise ValueError("committed_mask must match the code group shape")
        if not 0.5 <= initial_proxy_magnitude <= 1.5:
            raise ValueError("initial_proxy_magnitude must be in [0.5, 1.5]")
        if initial_zero_proxy_boundary is not None:
            if not 0.0 < initial_zero_proxy_boundary < 0.5:
                raise ValueError(
                    "initial_zero_proxy_boundary must be strictly inside (0, 0.5)"
                )
            if master_weight is None:
                raise ValueError(
                    "initial_zero_proxy_boundary requires master_weight for "
                    "deterministic activation directions"
                )
        full_in_features = codes.shape[1] * codes.shape[2]
        if master_weight is None:
            if not committed_mask.all():
                raise ValueError("partial proxy matrices require master_weight")
            base_weight = torch.zeros_like(codes, dtype=torch.float32)
            in_features = full_in_features
        else:
            if master_weight.ndim != 2 or master_weight.shape[0] != codes.shape[0]:
                raise ValueError("master_weight must match the code output dimension")
            if not 0 < master_weight.shape[1] <= full_in_features:
                raise ValueError("master_weight has an invalid input dimension")
            if full_in_features - master_weight.shape[1] >= codes.shape[2]:
                raise ValueError("master_weight padding must be smaller than one group")
            in_features = int(master_weight.shape[1])
            padded = F.pad(
                master_weight.detach().float(), (0, full_in_features - in_features)
            )
            base_weight = padded.view_as(codes)
        initial_proxy = codes.detach().float().clone() * float(
            initial_proxy_magnitude
        )
        if initial_zero_proxy_boundary is not None:
            # Keep the deployed code exactly zero while placing its continuous
            # training proxy close to the nearest hard boundary.  The sign of
            # the original high-precision weight supplies a deterministic
            # direction for possible 0 -> +/-1 transitions.  Exact zeros use
            # +1 so that initialization never introduces an ambiguous sign.
            source_direction = torch.where(
                base_weight >= 0,
                torch.ones_like(base_weight),
                -torch.ones_like(base_weight),
            )
            eligible_zero = (codes == 0) & committed_mask.unsqueeze(-1)
            boundary_proxy = source_direction * float(initial_zero_proxy_boundary)
            initial_proxy = torch.where(
                eligible_zero, boundary_proxy, initial_proxy
            )
        self.proxy_code = nn.Parameter(initial_proxy)
        self.group_scale = nn.Parameter(scales.detach().float().clone())
        self.register_buffer("initial_codes", codes.detach().to(torch.int8).clone())
        self.register_buffer("committed_mask", committed_mask.detach().clone())
        self.register_buffer("base_weight", base_weight.detach().clone())
        self._in_features = in_features
        self.compute_dtype = compute_dtype
        self.temperature = float(temperature)
        self.fake_fp16_scale = bool(fake_fp16_scale)

    @property
    def out_features(self) -> int:
        return self.proxy_code.shape[0]

    @property
    def in_features(self) -> int:
        return self._in_features

    @property
    def group_size(self) -> int:
        return self.proxy_code.shape[2]

    def hard_codes(self) -> torch.Tensor:
        return self.proxy_code.detach().round().clamp(-1, 1).to(torch.int8)

    def code_churn(self) -> float:
        changed = self.hard_codes() != self.initial_codes
        return float(changed[self.committed_mask].float().mean().item())

    @torch.no_grad()
    def deployment_statistics(self, boundary_epsilon: float = 0.05) -> dict:
        """Summarize hard-code stability and scale health on deployed groups."""
        if boundary_epsilon < 0:
            raise ValueError("boundary_epsilon must be non-negative")
        active_proxy = self.proxy_code[self.committed_mask]
        active_codes = self.hard_codes()[self.committed_mask]
        active_initial = self.initial_codes[self.committed_mask]
        active_scales = self.group_scale.detach().abs()[self.committed_mask]
        if active_proxy.numel() == 0:
            raise ValueError("deployment statistics require committed groups")
        boundary_distance = torch.minimum(
            (active_proxy - 0.5).abs(), (active_proxy + 0.5).abs()
        )
        counts = {
            str(code): int((active_codes == code).sum().item())
            for code in (-1, 0, 1)
        }
        probabilities = torch.tensor(
            list(counts.values()), dtype=torch.float64, device=active_proxy.device
        )
        probabilities /= probabilities.sum().clamp_min(1)
        nonzero = probabilities > 0
        entropy = -(probabilities[nonzero] * probabilities[nonzero].log2()).sum()
        scale_quantiles = torch.quantile(
            active_scales.float(),
            torch.tensor([0.0, 0.5, 0.95, 1.0], device=active_scales.device),
        )
        return {
            "code_counts": counts,
            "zero_fraction": float((active_codes == 0).float().mean().item()),
            "code_entropy_bits": float(entropy.item()),
            "code_churn": float((active_codes != active_initial).float().mean().item()),
            "proxy_abs_displacement_mean": float(
                (active_proxy - active_initial.float()).abs().mean().item()
            ),
            "boundary_epsilon": float(boundary_epsilon),
            "near_boundary_fraction": float(
                (boundary_distance <= boundary_epsilon).float().mean().item()
            ),
            "boundary_distance_min": float(boundary_distance.min().item()),
            "boundary_distance_p01": float(
                torch.quantile(boundary_distance.float(), 0.01).item()
            ),
            "scale_min": float(scale_quantiles[0].item()),
            "scale_median": float(scale_quantiles[1].item()),
            "scale_p95": float(scale_quantiles[2].item()),
            "scale_max": float(scale_quantiles[3].item()),
            "scale_at_clamp_fraction": float(
                (active_scales <= 1.00001e-5).float().mean().item()
            ),
        }

    def proxy_anchor_loss(self) -> torch.Tensor:
        """Squared proxy displacement over deployed ternary groups only."""
        delta = (self.proxy_code - self.initial_codes.float()).square()
        return delta[self.committed_mask].mean()

    def effective_weight(self) -> torch.Tensor:
        soft = soft_ternary_proxy(self.proxy_code, self.temperature)
        hard = self.proxy_code.round().clamp(-1, 1)
        # Exact hard forward with the smooth staircase supplying the gradient.
        code = hard.detach() + soft - soft.detach()
        scale = self.group_scale.abs().clamp_min(1e-5)
        if self.fake_fp16_scale:
            rounded = scale.half().float()
            scale = scale + (rounded - scale).detach()
        value = code * scale.unsqueeze(-1)
        mixed = torch.where(self.committed_mask.unsqueeze(-1), value, self.base_weight)
        return mixed.reshape(self.out_features, -1)[:, : self.in_features].to(
            self.compute_dtype
        )

    @torch.no_grad()
    def constrain_(self) -> None:
        self.proxy_code.clamp_(-1.5, 1.5)
        self.group_scale.clamp_(min=1e-5)


class ProxyTernaryLinear(nn.Module):
    def __init__(self, matrix: ProxyTernaryMatrix, bias=None):
        super().__init__()
        self.matrix = matrix
        self.bias = None if bias is None else nn.Parameter(
            bias.detach().clone(), requires_grad=False
        )
        self.in_features = matrix.in_features
        self.out_features = matrix.out_features

    def forward(self, value: torch.Tensor) -> torch.Tensor:
        return F.linear(value, self.matrix.effective_weight().to(value.dtype), self.bias)