processed/covtoken_code/gate/mask_gumbel.py

"""Differentiable retention mask via Gumbel-sigmoid / straight-through (formalization §4).

The gate pi_theta(x) in [0,1]^n produces per-token keep-logits; the relaxed mask
m_tilde in {0,1}^n is sampled with a straight-through Gumbel-sigmoid (hard forward, soft
backward) so the discrete pruning decision is trainable WITHOUT policy gradients
(anti-goal §5: no RL). At inference the mask is thresholded deterministically.
"""
from __future__ import annotations

import torch


def gumbel_sigmoid(logits: torch.Tensor, tau: float = 0.5, hard: bool = True,
                   generator: torch.Generator | None = None) -> torch.Tensor:
    """Straight-through Gumbel-sigmoid. Returns m in {0,1} (hard) with soft gradients."""
    u = torch.rand(logits.shape, device=logits.device, dtype=logits.dtype,
                   generator=generator).clamp_(1e-6, 1 - 1e-6)
    logistic_noise = torch.log(u) - torch.log1p(-u)
    y_soft = torch.sigmoid((logits + logistic_noise) / tau)
    if not hard:
        return y_soft
    y_hard = (y_soft > 0.5).to(logits.dtype)
    return y_hard + (y_soft - y_soft.detach())  # straight-through


def threshold_mask(logits: torch.Tensor) -> torch.Tensor:
    """Deterministic inference-time mask: keep token iff keep-logit > 0."""
    return (logits > 0).to(logits.dtype)