BorisTM/loss-guided-static-multi

1"""A second vocabulary for contested (language, token) pairs.2 3Every other conditioning mechanism in this repo failed for one measured reason:4the shared table rewrites about 85% of a row's magnitude in 3,000 steps, so it5absorbs any per-language transform a head could learn and leaves the head6redundant.7 8This does something different. A selected (language, token) pair gets its **own9row**, and a learned gate mixes it with the shared row:10 11    E'[w, l] = (1 - g) * E_shared[w] + g * E_split[idx(l, w)]12 13The split row is free capacity, not a correction, so there is nothing for the14shared table to absorb. Split rows start as exact copies of their shared row, so15the model is numerically identical to the unconditioned one at step zero16whatever the gate says, and both paths receive gradient immediately — the17failure mode where a zero-initialised head never leaves zero cannot occur.18 19An L1 penalty on the gate makes a split row earn its keep. After training the20gate is the answer to the scientific question: which pairs actually needed21separating, and does that agree with the interference measured beforehand.22"""23 24from __future__ import annotations25 26import numpy as np27import torch28from torch import nn29 30 31class SplitVocabulary(nn.Module):32    def __init__(self, index_path: str, vocab_size: int, dim: int, n_languages: int,33                 gate_init: float = 0.0) -> None:34        super().__init__()35        data = np.load(index_path)36        language_index = torch.tensor(data["language_index"].astype(np.int64))37        token_index = torch.tensor(data["token_index"].astype(np.int64))38        self.n_rows = int(language_index.numel())39 40        # (language, token) -> split row, as one flat lookup. -1 means "shared41        # only", which is the vast majority of pairs.42        lookup = torch.full((n_languages, vocab_size), -1, dtype=torch.long)43        lookup[language_index, token_index] = torch.arange(self.n_rows)44        self.register_buffer("lookup", lookup, persistent=False)45        self.register_buffer("row_language", language_index, persistent=False)46        self.register_buffer("row_token", token_index, persistent=False)47 48        self.rows = nn.Parameter(torch.zeros(self.n_rows, dim))49        self.gate_logit = nn.Parameter(torch.full((self.n_rows,), float(gate_init)))50        self._initialised = False51 52    @torch.no_grad()53    def seed_from(self, table: torch.Tensor) -> None:54        """Copy each split row from the shared row it specialises."""55        self.rows.copy_(table[self.row_token])56        self._initialised = True57 58    def gate(self) -> torch.Tensor:59        return torch.sigmoid(self.gate_logit)60 61    def penalty(self) -> torch.Tensor:62        """Mean gate opening — an L1 that pushes unused rows back to shared."""63        return self.gate().mean()64 65    def fusion_penalty(self) -> torch.Tensor:66        """Pull every language's row for a token toward the consensus for it.67 68        This is the merging mechanism. Training starts with one private row per69        (language, token) — full capacity, no interference — and this term makes70        a row stay distinct only if the contrastive loss pays for it. Sweeping71        its weight walks continuously from one vocabulary per language to a72        single shared one, so the *effective* number of vocabularies is read off73        rather than chosen.74 75        Fusing toward the per-token mean is the O(L) relaxation of the O(L^2)76        all-pairs fused lasso; rows that need not differ collapse onto the77        centroid, and which languages stay off it together is the clustering.78        """79        token = self.row_token80        uniq, inverse = torch.unique(token, return_inverse=True)81        total = torch.zeros(uniq.numel(), self.rows.shape[1],82                            dtype=self.rows.dtype, device=self.rows.device)83        total.index_add_(0, inverse, self.rows)84        count = torch.zeros(uniq.numel(), dtype=self.rows.dtype, device=self.rows.device)85        count.index_add_(0, inverse, torch.ones_like(inverse, dtype=self.rows.dtype))86        centroid = total / count.clamp(min=1.0).unsqueeze(1)87        # Group lasso over languages: the L2 norm of each deviation, averaged.88        return (self.rows - centroid[inverse]).norm(dim=1).mean()89 90    def forward(self, token_ids: torch.Tensor, lang_per_token: torch.Tensor,91                shared: torch.Tensor, use_gate: bool = True) -> torch.Tensor:92        row = self.lookup[lang_per_token, token_ids]93        has = row >= 094        if not bool(has.any()):95            return shared96        safe = row.clamp(min=0)97        if not use_gate:98            # Fusion mode: the private row replaces the shared one outright, and99            # merging is driven by the penalty rather than by a mixing gate.100            return torch.where(has.unsqueeze(1), self.rows[safe], shared)101        g = self.gate()[safe].unsqueeze(1) * has.unsqueeze(1)102        return shared * (1.0 - g) + self.rows[safe] * g103 104    def extra_repr(self) -> str:105        return f"rows={self.n_rows}"106