| """K vocabularies, and a learned assignment of languages to them. |
| |
| The earlier fusion attempt merged *per token*: two languages could share a row |
| for one word and keep separate rows for the next. That is cell-level sharing, |
| not vocabulary sharing, and it is why the readout came out as mush — the merge |
| decision was never coupled across a language's tokens. |
| |
| Here the decision lives at the language level. Each language holds one |
| distribution over K vocabularies, used for every token it emits: |
| |
| E[l, w] = E_shared[w] + sum_k pi[l, k] * Delta_k[w] |
| |
| `pi` is L x K — 48 numbers for twelve languages and four vocabularies — so the |
| grouping is read directly off the parameter instead of being inferred from row |
| distances. An entropy penalty pushes each language toward one vocabulary, and K |
| is the explicit axis: K=1 is today's single shared vocabulary, K=L with a hard |
| assignment is one per language, and where quality stops improving is the answer |
| to how many a multilingual encoder actually needs. |
| |
| Offsets start at zero so the model is numerically identical to the shared one at |
| step zero. The assignment logits carry the small random asymmetry instead: with |
| identical offsets and a uniform assignment every expert would receive the same |
| gradient and the K vocabularies would collapse back into one by symmetry. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import numpy as np |
| import torch |
| from torch import nn |
|
|
|
|
| class VocabAssignment(nn.Module): |
| def __init__(self, index_path: str, vocab_size: int, dim: int, |
| n_languages: int, n_vocabs: int = 4, seed: int = 20260727) -> None: |
| super().__init__() |
| data = np.load(index_path) |
| tokens = np.unique(data["token_index"].astype(np.int64)) |
| self.n_tokens = int(tokens.size) |
| self.n_vocabs = int(n_vocabs) |
|
|
| slot = torch.full((vocab_size,), -1, dtype=torch.long) |
| slot[torch.tensor(tokens)] = torch.arange(self.n_tokens) |
| self.register_buffer("slot", slot, persistent=False) |
|
|
| self.offsets = nn.Parameter(torch.zeros(self.n_vocabs, self.n_tokens, dim)) |
| generator = torch.Generator().manual_seed(seed) |
| self.logits = nn.Parameter( |
| torch.randn(n_languages, self.n_vocabs, generator=generator) * 0.01) |
|
|
| def assignment(self) -> torch.Tensor: |
| return torch.softmax(self.logits, dim=-1) |
|
|
| def entropy(self) -> torch.Tensor: |
| """Mean assignment entropy; penalising it drives each language to one slot.""" |
| pi = self.assignment() |
| return -(pi * torch.log(pi.clamp_min(1e-9))).sum(-1).mean() |
|
|
| def forward(self, token_ids: torch.Tensor, lang_per_token: torch.Tensor, |
| shared: torch.Tensor) -> torch.Tensor: |
| index = self.slot[token_ids] |
| has = index >= 0 |
| if not bool(has.any()): |
| return shared |
| safe = index.clamp(min=0) |
| |
| gathered = self.offsets[:, safe] |
| weights = self.assignment()[lang_per_token] |
| delta = torch.einsum("nk,knd->nd", weights, gathered) |
| return shared + delta * has.unsqueeze(1) |
|
|
| def extra_repr(self) -> str: |
| return f"tokens={self.n_tokens}, vocabs={self.n_vocabs}" |
|
|
