models/answer_selector.py

import math

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

from config import DROPOUT


def bounded_phase(x):
    """Bound an unconstrained head into (-pi, pi]."""

    return math.pi * torch.tanh(x)


def _masked_standardize(x, valid, eps=1e-5):
    """
    Zero-mean, unit-variance per sample over the (K, N) axes, ignoring
    padded entries. Padded entries are returned as 0.
    """

    m = valid.to(x.dtype)

    n = m.sum(dim=(1, 2), keepdim=True).clamp(min=1.0)

    mean = (x * m).sum(dim=(1, 2), keepdim=True) / n

    var = (((x - mean) ** 2) * m).sum(dim=(1, 2), keepdim=True) / n

    out = (x - mean) / torch.sqrt(var + eps)

    return out * m


class EnergyAnswerSelector(nn.Module):
    """
    Pairwise (hypothesis x option) energy. Lower energy = better answer.

    Fixed in v45
    ------------
    * **Confidence self-gating removed.** v44 computed
      `energy = energy * (0.5 + confidence)` where `confidence` was a
      head reading the same features as the energy itself. That is a
      multiplicative self-gate: it rescales each (k, n) pair by a
      different learned factor, so energies stop being comparable across
      pairs -- and the marginalization downstream compares them. It has
      no probabilistic interpretation and it is gone. (The earlier form,
      `energy / confidence.clamp(min=0.2)`, was worse: up to 5x
      amplification straight into a +/-6 clamp, which saturated and
      killed gradients -- the documented cause of the reverted warm-start
      experiment.)

    * **The three energies are scale-normalized before fusion.** They
      live on different scales: `learned_energy` is an unbounded MLP
      output, while `hamiltonian_energy` and `cosine_energy` are inner
      products of normalized vectors, bounded in roughly [-1, 1].
      Combining them with per-pair softmax weights meant the learned term
      dominated by magnitude alone, whatever the weights said. Each term
      is now standardized over the (K, N) axes per sample first, so the
      fusion weights actually express a preference.

    * **Alignment and polarity features.** v45 generates a supporting and
      an attacking hypothesis per option. `align[k, n] = 1` when
      hypothesis k is ABOUT option n, and `polarity[k] = +/-1` says which
      direction it argues. Without these the support/attack distinction
      would be a sign bit floating free of the option it refers to, and
      "the objection to option 2" could not be used as evidence against
      option 2 specifically.

    * **Phase head.** Emits theta_kn for models/interference.py, so the
      compatibility amplitude v_kn = t_kn * exp(i theta_kn) is complex
      and the Born rule stops being a reparameterized softmax.

    Kept from v44
    -------------
    Question conditioning (`use_question`) -- the largest measured single
    effect in this project: on cached ARC embeddings a plain MLP scores
    0.3585 from options alone and 0.4849 once the question is available.
    The question text sat in the cache un-encoded for the project's whole
    history.
    """

    def __init__(self, dim, use_question=True):

        super().__init__()

        self.use_question = use_question

        self.hamiltonian = nn.Linear(dim, dim, bias=False)

        # [H, O, H-O, H*O] plus, when available, [Q, Q*O, Q-O];
        # then three scalars: align, polarity, align*polarity.
        n_vec = 7 if use_question else 4

        n_feat = dim * n_vec + 3

        self.energy_net = nn.Sequential(
            nn.Linear(n_feat, dim * 2),
            nn.GELU(),
            nn.LayerNorm(dim * 2),
            nn.Dropout(DROPOUT),
            nn.Linear(dim * 2, dim),
            nn.GELU(),
            nn.Linear(dim, 1),
        )

        self.fusion = nn.Sequential(
            nn.Linear(n_feat, dim),
            nn.GELU(),
            nn.Linear(dim, 3),
        )

        # Deliberately a single linear map rather than an MLP: the phase
        # is one bounded scalar per pair and this keeps the parameter
        # cost of the complex-amplitude extension negligible, so a gain
        # from it cannot be explained as extra capacity.
        self.phase_net = nn.Linear(n_feat, 1)

        self.temperature = nn.Parameter(torch.tensor(1.0))

    def forward(self, H, O, Q=None, align=None, polarity=None,
                H_mask=None, O_mask=None):
        """
        H        : (B, K, D)
        O        : (B, N, D)
        Q        : (B, D)
        align    : (B, K, N) 1.0 where hypothesis k is about option n
        polarity : (B, K)    +1 support / -1 attack
        """

        B, K, D = H.shape
        _, N, _ = O.shape

        H = F.normalize(H, dim=-1)
        O = F.normalize(O, dim=-1)

        O_proj = self.hamiltonian(O)

        H_exp = H.unsqueeze(2).expand(B, K, N, D)
        O_exp = O_proj.unsqueeze(1).expand(B, K, N, D)

        parts = [H_exp, O_exp, H_exp - O_exp, H_exp * O_exp]

        if self.use_question:

            if Q is None:
                raise ValueError(
                    "EnergyAnswerSelector was built with use_question=True "
                    "but forward() got Q=None. Rebuild the cache so it "
                    "carries question embeddings, or construct the model "
                    "with use_question=False."
                )

            Q_exp = F.normalize(Q, dim=-1).view(B, 1, 1, D).expand(B, K, N, D)

            parts += [Q_exp, Q_exp * O_exp, Q_exp - O_exp]

        if align is None:
            align = torch.zeros(B, K, N, device=H.device, dtype=H.dtype)

        if polarity is None:
            polarity = torch.zeros(B, K, device=H.device, dtype=H.dtype)

        a = align.unsqueeze(-1)
        p = polarity.view(B, K, 1, 1).expand(B, K, N, 1)

        parts += [a, p, a * p]

        features = torch.cat(parts, dim=-1)

        # ------------------------------------------------------------
        # Three energies, standardized onto a common scale
        # ------------------------------------------------------------

        valid = torch.ones(B, K, N, dtype=torch.bool, device=H.device)

        if H_mask is not None:
            valid = valid & H_mask.view(B, K, 1)

        if O_mask is not None:
            valid = valid & O_mask.view(B, 1, N)

        learned_energy = self.energy_net(features).squeeze(-1)

        hamiltonian_energy = -torch.einsum("bkd,bnd->bkn", H, O_proj)

        cosine_energy = -F.cosine_similarity(H_exp, O_exp, dim=-1)

        learned_energy = _masked_standardize(learned_energy, valid)
        hamiltonian_energy = _masked_standardize(hamiltonian_energy, valid)
        cosine_energy = _masked_standardize(cosine_energy, valid)

        fusion_weights = F.softmax(self.fusion(features), dim=-1)

        energy = (
            fusion_weights[..., 0] * learned_energy
            + fusion_weights[..., 1] * hamiltonian_energy
            + fusion_weights[..., 2] * cosine_energy
        )

        temperature = 0.5 + F.softplus(self.temperature)

        energy = energy / temperature

        # Wide safety net against runaway energies. Should not bind in
        # normal training -- if it does, something upstream is diverging.
        energy = torch.clamp(energy, -12, 12)

        # ------------------------------------------------------------
        # Compatibility phase for the coherent collapse
        # ------------------------------------------------------------

        phase = bounded_phase(self.phase_net(features).squeeze(-1))

        return {
            "energy": energy,
            "phase": phase,
            "fusion_weights": fusion_weights,
        }