BorisTM/loss-guided-static-multi

1"""Hashed token n-gram rows for a static encoder.2 3Fifteen representation-side interventions in this repo produced nothing, and the4per-language numbers say why: quality tracks *tokenisation*, not capacity.5Spearman between tokens-per-word and held-out Recall@1 is -0.607 over 576space-separated languages; the third that tokenise worst average 0.191 against70.295 for the third that tokenise best. Georgian needs 5.7 tokens per word and8scores 0.050.9 10Those languages do not lack a good vector for a word — they lack any row that11*means* a word. A Transformer reassembles the pieces with attention; a12mean-pooled static encoder cannot. The only repair available to it is a bigger13lookup unit.14 15So each adjacent token n-gram inside a sentence gets a row, hashed into a fixed16table, pooled alongside the unigram rows:17 18    h(x) = ( sum_w E[w] + sum_g F[hash(g)] ) / (|U| + |G|)19 20Allocation across languages needs no rule: a fragmented language simply emits21more distinct n-grams and takes more of the table, while a language that already22tokenises into words gains little. The budget follows the deficit by itself.23 24Rows start at zero, so the pooled direction — and therefore every cosine — is25identical to the unconditioned model at step zero, while the rows still receive26gradient from the first batch. Unlike every conditioning head tried here, an27n-gram row is a *new input feature* rather than a second route to an existing28one, so the shared table cannot absorb it: the feature does not otherwise exist.29"""30 31from __future__ import annotations32 33import torch34from torch import nn35 36# Odd 32-bit constants; the product mixes the low bits that adjacent token ids37# share, which matters because ids are correlated within a language.38_MIX_A = 265443576139_MIX_B = 4050340 41 42class NgramTable(nn.Module):43    def __init__(self, buckets: int, dim: int, orders: tuple[int, ...] = (2,)) -> None:44        super().__init__()45        if any(order < 2 for order in orders):46            raise ValueError("n-gram orders must be at least 2")47        self.buckets = int(buckets)48        self.orders = tuple(orders)49        self.rows = nn.Parameter(torch.zeros(self.buckets, dim))50 51    def extra_repr(self) -> str:52        return f"buckets={self.buckets}, orders={self.orders}"53 54    def _hash(self, ids: torch.Tensor, order: int) -> torch.Tensor:55        value = torch.zeros_like(ids[0])56        for position in range(order):57            value = value * _MIX_A + ids[position] * _MIX_B + position58        return value.abs() % self.buckets59 60    def gather(61        self,62        content: torch.Tensor,63        segment: torch.Tensor,64    ) -> tuple[torch.Tensor, torch.Tensor]:65        """Rows and their sentence index for every n-gram inside a sentence."""66        vectors: list[torch.Tensor] = []67        segments: list[torch.Tensor] = []68        for order in self.orders:69            if content.numel() <= order:70                continue71            window = [content[i: content.numel() - order + 1 + i] for i in range(order)]72            starts = segment[: segment.numel() - order + 1]73            # An n-gram must lie inside one sentence, so every position in the74            # window has to carry the same segment id.75            same = torch.ones_like(starts, dtype=torch.bool)76            for i in range(1, order):77                same &= segment[i: segment.numel() - order + 1 + i] == starts78            if not bool(same.any()):79                continue80            index = self._hash([w[same] for w in window], order)81            vectors.append(self.rows[index])82            segments.append(starts[same])83        if not vectors:84            empty = content.new_zeros((0,), dtype=torch.long)85            return self.rows.new_zeros((0, self.rows.shape[1])), empty86        return torch.cat(vectors), torch.cat(segments)87