| """A second vocabulary for contested (language, token) pairs. |
| |
| Every other conditioning mechanism in this repo failed for one measured reason: |
| the shared table rewrites about 85% of a row's magnitude in 3,000 steps, so it |
| absorbs any per-language transform a head could learn and leaves the head |
| redundant. |
| |
| This does something different. A selected (language, token) pair gets its **own |
| row**, and a learned gate mixes it with the shared row: |
| |
| E'[w, l] = (1 - g) * E_shared[w] + g * E_split[idx(l, w)] |
| |
| The split row is free capacity, not a correction, so there is nothing for the |
| shared table to absorb. Split rows start as exact copies of their shared row, so |
| the model is numerically identical to the unconditioned one at step zero |
| whatever the gate says, and both paths receive gradient immediately — the |
| failure mode where a zero-initialised head never leaves zero cannot occur. |
| |
| An L1 penalty on the gate makes a split row earn its keep. After training the |
| gate is the answer to the scientific question: which pairs actually needed |
| separating, and does that agree with the interference measured beforehand. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import numpy as np |
| import torch |
| from torch import nn |
|
|
|
|
| class SplitVocabulary(nn.Module): |
| def __init__(self, index_path: str, vocab_size: int, dim: int, n_languages: int, |
| gate_init: float = 0.0) -> None: |
| super().__init__() |
| data = np.load(index_path) |
| language_index = torch.tensor(data["language_index"].astype(np.int64)) |
| token_index = torch.tensor(data["token_index"].astype(np.int64)) |
| self.n_rows = int(language_index.numel()) |
|
|
| |
| |
| lookup = torch.full((n_languages, vocab_size), -1, dtype=torch.long) |
| lookup[language_index, token_index] = torch.arange(self.n_rows) |
| self.register_buffer("lookup", lookup, persistent=False) |
| self.register_buffer("row_language", language_index, persistent=False) |
| self.register_buffer("row_token", token_index, persistent=False) |
|
|
| self.rows = nn.Parameter(torch.zeros(self.n_rows, dim)) |
| self.gate_logit = nn.Parameter(torch.full((self.n_rows,), float(gate_init))) |
| self._initialised = False |
|
|
| @torch.no_grad() |
| def seed_from(self, table: torch.Tensor) -> None: |
| """Copy each split row from the shared row it specialises.""" |
| self.rows.copy_(table[self.row_token]) |
| self._initialised = True |
|
|
| def gate(self) -> torch.Tensor: |
| return torch.sigmoid(self.gate_logit) |
|
|
| def penalty(self) -> torch.Tensor: |
| """Mean gate opening — an L1 that pushes unused rows back to shared.""" |
| return self.gate().mean() |
|
|
| def fusion_penalty(self) -> torch.Tensor: |
| """Pull every language's row for a token toward the consensus for it. |
| |
| This is the merging mechanism. Training starts with one private row per |
| (language, token) — full capacity, no interference — and this term makes |
| a row stay distinct only if the contrastive loss pays for it. Sweeping |
| its weight walks continuously from one vocabulary per language to a |
| single shared one, so the *effective* number of vocabularies is read off |
| rather than chosen. |
| |
| Fusing toward the per-token mean is the O(L) relaxation of the O(L^2) |
| all-pairs fused lasso; rows that need not differ collapse onto the |
| centroid, and which languages stay off it together is the clustering. |
| """ |
| token = self.row_token |
| uniq, inverse = torch.unique(token, return_inverse=True) |
| total = torch.zeros(uniq.numel(), self.rows.shape[1], |
| dtype=self.rows.dtype, device=self.rows.device) |
| total.index_add_(0, inverse, self.rows) |
| count = torch.zeros(uniq.numel(), dtype=self.rows.dtype, device=self.rows.device) |
| count.index_add_(0, inverse, torch.ones_like(inverse, dtype=self.rows.dtype)) |
| centroid = total / count.clamp(min=1.0).unsqueeze(1) |
| |
| return (self.rows - centroid[inverse]).norm(dim=1).mean() |
|
|
| def forward(self, token_ids: torch.Tensor, lang_per_token: torch.Tensor, |
| shared: torch.Tensor, use_gate: bool = True) -> torch.Tensor: |
| row = self.lookup[lang_per_token, token_ids] |
| has = row >= 0 |
| if not bool(has.any()): |
| return shared |
| safe = row.clamp(min=0) |
| if not use_gate: |
| |
| |
| return torch.where(has.unsqueeze(1), self.rows[safe], shared) |
| g = self.gate()[safe].unsqueeze(1) * has.unsqueeze(1) |
| return shared * (1.0 - g) + self.rows[safe] * g |
|
|
| def extra_repr(self) -> str: |
| return f"rows={self.n_rows}" |
|
|
