ARCHITECTURE.md

# MillID architecture

**MillID** (MultiLingual Language IDentifier) names the language of a passage of text. v1 is a
byte-level, encoder-only transformer of **115M parameters** (`medium`; 193M as `large`), trained
from scratch in PyTorch on
[baobabtech/millid-dataset](https://huggingface.co/datasets/baobabtech/millid-dataset): 1,641
language-script labels (1,551 ISO 639-3 languages, 56 scripts). v2 adds images on the same
encoder. Everything here is in one file, [`millid.py`](millid.py).

```mermaid
flowchart TD
    T["Passage<br/>a table cell to a paragraph"] --> N["Unicode NFKC<br/>whitespace collapsed"]
    N --> E["Per-byte encoding<br/>byte · script · byte-in-character · character n-grams<br/>B × L × 4 ids, L ≤ 256"]
    E --> EMB["Sum of four embeddings<br/>B × L × d"]
    EMB --> STEM["Depthwise conv stem, k = 5<br/>local byte patterns"]
    STEM --> DOWN["Strided conv, 2 bytes → 1 token<br/>B × L/2 × d"]
    DOWN --> ENC["12 × pre-norm block<br/>RMSNorm → attention with RoPE → RMSNorm → SwiGLU"]
    ENC --> POOL["RMSNorm, masked mean pool<br/>z: B × d"]
    POOL --> HEAD["Language head<br/>z · (label + language + script vectors)ᵀ + b"]
    HEAD --> OUT["Softmax over 1,641 labels<br/>ranked (label, probability)"]
```

## 1. The task

Given a passage, from a table cell of a few characters to a paragraph of a thousand, return a
ranked list of `(label, probability)`. A label is an ISO 639-3 language code and an ISO 15924
script code: `eng_Latn`, `amh_Ethi`, `srp_Cyrl`. The same language in two scripts is two labels,
because telling them apart is part of the job.

## 2. One universal encoding

### Labels

The dataset's `labels.json` fixes label *i* to `labels[i]`, so the dataset's `label_id`, the
model's output index and every report agree. It also carries `aliases`, other spellings of a
language mapped to the one used (`swa_Latn` → `swh_Latn`, `jpn_Hira` → `jpn_Jpan`, 22 in all), and
`dropped`, macrolanguage codes naming several languages GlotLID tells apart (Serbo-Croatian, Fula,
Dinka, Kalenjin), whose text is left out rather than mislabelled. A model trained with extra
sources keeps every dataset id and appends its own labels; its `labels.json` lists the labels it
was trained on (`trained`), and no other label can be an answer (§5).

### Text: four ids per byte

Text is normalised (Unicode NFKC, whitespace collapsed) and read as UTF-8 bytes, so every
language in every script is the same 256 symbols: no tokenizer to train, no language split into
worse pieces than another, no vocabulary to outgrow, no unknown token. Bytes alone hide three
things a language identifier needs, so every byte carries four ids:

| Sequence | Value | What it tells the model |
|---|---|---|
| `ids` | the byte + 3 (0 `PAD`, 1 `CLS`, 2 `EOS`) | the text itself |
| `script` | the Unicode script of the byte's character, hashed into 512 buckets | which writing system: Latin, Ethiopic, Devanagari, Han, Adlam... Digits, punctuation and spaces share one "common" id |
| `charpos` | 1 one-byte character, 2 first byte of a longer one, 3 continuation byte | where characters begin and end |
| `ngrams` | the character 2-, 3- and 4-grams ending at the byte's character, lower-cased and hashed into 32,768 buckets | the features fastText and GlotLID identify languages by, learned as embeddings |

```mermaid
flowchart LR
    C["character ሀ<br/>ETHIOPIC SYLLABLE HA"] --> B1["byte E1<br/>charpos 2"]
    C --> B2["byte 88<br/>charpos 3"]
    C --> B3["byte 80<br/>charpos 3"]
    S["script: ETHIOPIC<br/>same id on all three bytes"] -.-> B1 & B2 & B3
    G["n-grams ending at ሀ<br/>2-gram, 3-gram, 4-gram ids<br/>same on all three bytes"] -.-> B1 & B2 & B3
```

A byte's input vector is the sum of four embeddings: byte, script, byte-in-character, and the sum
of its three n-gram embeddings. Python has no Unicode script property; the first word of a
character's Unicode name (`LATIN`, `ETHIOPIC`, `CJK`, `HANGUL`) is the same grouping for letters
and marks, and hashing it keeps ids stable when Unicode adds scripts; a table of every code point
below U+40000 is built once per process. The n-gram hash is a polynomial over the lower-cased code
points followed by splitmix64, fixed rather than Python's per-process `hash`, so ids are the same
in training and at inference. Before the start of a passage every position holds a value no code
point has, so a word that opens a passage has n-grams of its own; a window cut from the middle of
a passage keeps the n-grams the passage had there.

A passage is at most 256 tokens (`CLS` and 255 bytes, which is 85 to 255 characters). A longer one
is cut to a window of whole characters, never inside a multi-byte one: at inference the start, in
training a random place.

Encoding runs a whole batch at once (`encode_batch`): the batch's passages are one string, and
every step (windows, scripts, n-grams, UTF-8 bytes, padding) is one numpy operation over it. Per
passage, those steps would be dozens of tiny array operations whose overhead dominates; batched,
one CPU core encodes about 5,000 passages a second, so a handful of DataLoader workers keep even
an H200 fed.

## 3. The network

### Shapes

| Stage | Output | Notes |
|---|---|---|
| ids | B × L, four sequences | L ≤ 256 bytes |
| embeddings | B × L × d | byte + script + charpos + Σ n-grams; padding is zero |
| conv stem | B × L × d | residual: x + Conv1×1(GELU(DepthwiseConv₅(x))) |
| downsample | B × L/2 × d | Conv1d, kernel 2, stride 2; a token is real if either byte is |
| 12 blocks | B × L/2 × d | pre-norm, RoPE attention, SwiGLU |
| pool | B × d | RMSNorm, mean over real tokens: `embed()` returns this |
| head | B × C | C = 1,641 |

### The encoder block

```mermaid
flowchart TD
    X["x: B × T × d"] --> N1["RMSNorm"]
    N1 --> QKV["Linear d → 3d<br/>q, k, v in h heads of 64"]
    QKV --> ROPE["RoPE on q and k<br/>θ = 10,000"]
    ROPE --> ATT["scaled dot-product attention<br/>padding keys masked"]
    ATT --> O["Linear d → d"]
    O --> R1(("+"))
    X --> R1
    R1 --> N2["RMSNorm"]
    N2 --> FF["SwiGLU<br/>W₂( SiLU(W₁x) ⊙ W₃x ), width 8d/3"]
    FF --> R2(("+"))
    R1 --> R2
    R2 --> Y["to the next block"]
```

### The language head

Each label's output vector is built from its parts, so a language with little data borrows from
what it shares: `srp_Cyrl` and `srp_Latn` share Serbian, and every Ethiopic-script language shares
Ethiopic.

```mermaid
flowchart LR
    L["label vector<br/>srp_Cyrl"] --> SUM(("+"))
    G["language vector<br/>srp, shared with srp_Latn"] --> SUM
    S["script vector<br/>Cyrl, shared with rus_Cyrl, ukr_Cyrl, ..."] --> SUM
    SUM --> W["row of W: C × d"]
    Z["z: passage vector"] --> DOT["z · Wᵀ + b"]
    W --> DOT
    DOT --> LOGIT["logit per label"]
```

`LanguageHead.language_vectors()` returns W: the learned language embeddings, usable on their own
to compare languages.

### Presets

| Preset | d | Blocks | Heads | SwiGLU width | Down-sample | n-gram buckets | Parameters |
|---|---|---|---|---|---|---|---|
| `tiny` | 192 | 3 | 3 | 512 | 1 | 4,096 | 2.9M (tests) |
| `small` | 384 | 6 | 6 | 1,024 | 2 | 16,384 | 18.9M (sweeps) |
| **`medium`** (default) | 768 | 12 | 12 | 2,048 | 2 | 32,768 | **115.0M** |
| **`large`** | 1,024 | 12 | 16 | 2,752 | 2 | 32,768 | **192.6M** |

Where `medium`'s 115.0M parameters are: 85.0M in the blocks (7.08M each: attention 2.36M, SwiGLU
4.72M), 25.2M in the n-gram table, 2.5M in the language head (1,641 label, 1,551 language and 56
script vectors), 1.8M in the stem and downsampling, 0.6M in the byte, script and position
embeddings.

### Why each part

- **Bytes, not subwords.** A subword vocabulary for 1,600 languages either fragments the small
  ones into near-characters or spends most of its size on the large ones; bytes treat them alike.
  CANINE and ByT5 showed character- and byte-level encoders can match subword ones on
  multilingual tasks, and language identification is about the surface of text.
- **Script, character and n-gram embeddings.** What bytes hide, given directly: the writing
  system, where characters begin, and the character n-grams that fastText (GlotLID's architecture)
  is built on. The transformer starts from the features a fastText model ends with.
- **The conv stem and downsampling.** A width-5 depthwise convolution mixes neighbouring bytes;
  merging every two bytes then halves the sequence the transformer reads, which halves its compute,
  as CANINE does. A token covers two bytes: two-thirds of an Ethiopic or Devanagari character, two
  Latin ones, after the embeddings have already summarised each character.
- **Pre-norm RMSNorm, RoPE, SwiGLU.** The block of current encoders: stable at depth, relative
  positions that need no table, and a gated feed-forward that trains better at the same size.
  Attention runs on PyTorch's fused scaled-dot-product kernels.
- **Mean pooling.** Every token votes, which is robust for short passages. `CLS` stays as the
  first token, which gives v2 a slot to attend from.
- **Structured language vectors.** See the head above.

### Initialisation

Linear and embedding weights from N(0, 0.02), padding rows zero, biases zero. The last projection
of every residual branch (attention output, SwiGLU W₂) from N(0, 0.02 / √(2 × blocks)), so a deep
stack starts close to the identity.

## 4. Training

```mermaid
flowchart LR
    D["train split<br/>1.9M+ passages"] --> SAMP["draw a language ∝ n^0.3<br/>then one of its passages"]
    SAMP --> SHORT["25%: cut to a 20–80 character snippet"]
    SHORT --> WIN["random window of whole characters<br/>≤ 255 bytes"]
    WIN --> ENC["four ids per byte"]
    ENC --> W["8 DataLoader workers<br/>batches of 512"]
    W --> GPU["forward, cross-entropy<br/>label smoothing 0.1, bf16"]
    GPU --> OPT["AdamW, clip 1.0<br/>warm-up then cosine"]
    OPT --> VAL{"every 2,000 steps:<br/>validation macro F1"}
    VAL -->|better| SAVE["keep checkpoint"]
    VAL -->|12,000 steps without gain| STOP["stop"]
```

| | |
|---|---|
| Objective | cross-entropy over the labels, label smoothing 0.1, on float32 logits |
| Optimiser | AdamW, β = (0.9, 0.98), ε = 1e-8; weight decay 0.01 on weight matrices (not embeddings, norms, biases) |
| Schedule | linear warm-up for 2,000 steps to 3e-4, cosine to 5% of the peak |
| Batch | 512 passages, padded to the longest in the batch |
| Steps | 60,000: about 31M passages, 16 passes over the training split |
| Precision | bfloat16 autocast; gradients clipped to norm 1.0; optional `torch.compile` |
| Checkpoint | validation every 2,000 steps (20 passages per language), best macro F1 kept |

- **Which language.** In proportion to n^α, α = 0.3, n its passage count: between the data as it
  is (α = 1), where the largest languages fill every batch, and uniform (α = 0), where a
  100-passage language is seen until memorised. XLM-R and mT5 sample their languages the same way.
- **How long.** A quarter of passages are cut to a 20 to 80 character snippet, on word boundaries
  where the script has them: short text is the hard case, and what a document holds in its
  headers, captions and table cells.
- **Why macro F1 chooses.** It weighs every language equally, the score a thousand-language model
  should be judged on. Top-1 accuracy, which the largest languages dominate, is reported beside it.

### Cost

The transformer reads 128 tokens per passage after downsampling. Training costs about 6 ×
parameters × tokens operations for the dense layers, plus attention; the embedding tables are
lookups and cost almost nothing.

| Preset | Per passage | 60,000 steps × 512 | Measured | Training time |
|---|---|---|---|---|
| `tiny` | small; 256 tokens, no downsampling | | 18,200 passages/s on an H200 (2,100 on an L4) | ≈ 0.5 h (4.1 h on an L4) |
| `medium` | ≈ 67 GFLOPs | ≈ 2.1 × 10¹⁸ | 2,500 passages/s on an H200 | ≈ 3.4 h |
| `large` | ≈ 119 GFLOPs | ≈ 3.7 × 10¹⁸ | 1,790 passages/s on an H200 | ≈ 4.8 h |

On the H200 that is about 17% of peak bf16 throughput: the model runs eagerly, and at these
widths the norms, rotary embeddings and SwiGLU gates, each a separate pass over memory, weigh
as much as the matrix multiplications. `--compile` fuses them and is the first thing to try for
a faster run. Every run is GPU-bound; data loading keeps pace with room to spare (§2,
`encode_batch`). A run logs passages per second every hundred steps, so a new configuration's
cost is known within minutes. The final evaluation (GlotLID included) adds a few minutes.

## 5. Inference

```python
from millid import MillID

model = MillID.from_pretrained("baobabtech/millid")   # config.json, labels.json, model.safetensors
model.predict(["Lave men ou ak savon anvan ou manje."], k=3)
# e.g. [[("hat_Latn", 0.97), ("fra_Latn", 0.01), ...]]
model.embed(["..."])                                   # passage vectors, B × d
model.model.head.language_vectors()                    # language embeddings, C × d
```

The text is encoded as in §2 (the first window), the logits of labels the model was never trained
on are set to −∞, and a softmax gives the probabilities; batches of 512, bfloat16 on GPU.
`medium` is 460 MB in float32 and runs on a CPU, slowly.

## 6. Evaluation

| Set | Used for |
|---|---|
| `validation` | choosing the checkpoint; never reported as the result |
| `test` | the result: split by document, so no test passage shares a document with a training one |
| `udhr` | the honest result: the UDHR, labelled by its translators, not by GlotLID; passages quoting it were removed from training. `udhr` covers the languages MillID knows, `udhr-all` all 429 |
| `test-common`, `udhr-common` | the head-to-head: every model on the passages whose language MillID, GlotLID and fastText all know (about 210) |

Per language: support, TP, FP, FN, TN, precision, recall (= sensitivity), F1, specificity, false
positives per million passages of other languages. Averaged: macro, weighted, micro (= top-1
accuracy), median F1, languages above 0.9 and below 0.5. Top-1, top-3 and top-5 accuracy; top-1 by
passage size.

Two existing identifiers are scored on the same sets, each as it is and restricted to MillID's
labels (its best answer among them), the like-for-like comparison:

| | Labels | Model | Flag |
|---|---|---|---|
| [GlotLID](https://huggingface.co/cis-lmu/glotlid) | 2,000+ | fastText, 1.7 GB | `--compare-glotlid` |
| [fastText, Meta](https://huggingface.co/facebook/fasttext-language-identification) (NLLB) | 218 | fastText, 1.2 GB | `--compare-fasttext` |

Their answers are mapped onto MillID's spelling of a language before scoring (Simplified and
Traditional Han as `Hani`, MillID's aliases such as `pes`/`prs` as `fas`, and `est` as `ekk`), so a
right answer spelled differently is not counted wrong. Each reads a set once however many MillIDs
are evaluated (`millid.py evaluate --model A B C`).

## 7. Configuration and tracking

CleanRL-style: each command (`train`, `evaluate`, `sweep-agent`, `predict`) is a typed dataclass,
its fields the flags and their docstrings the help, parsed by tyro; `python millid.py train
--help` lists every setting. The resolved settings are saved with every run, in `config.json` and
in Weights & Biases, which with `WANDB_API_KEY` set receives the configuration, training loss,
learning rate, gradient norm, throughput, the validation curve, and the final per-set metrics with
per-language tables. `sweep.yaml` is a Bayesian sweep with hyperband stopping; its agents run as
Hugging Face Jobs.

## 8. MillID v2: text and images

```mermaid
flowchart LR
    TXT["text"] --> TE["MillIDTextEncoder<br/>v1, text.*"] --> ZT["z_text: B × d"]
    IMG["page or crop"] --> IE["image encoder<br/>ViT or ConvNeXt, image.*"] --> ZI["z_image: B × d"]
    ZT --> F["gated fusion, fuse.*<br/>z = g·z_text + (1 − g)·z_image"]
    ZI --> F
    F --> H["LanguageHead<br/>v1, head.*, same labels"] --> O["language"]
```

- **Image encoder.** A small vision backbone over a page or a crop, projected to d, reading the
  script and the language from the shapes of the writing: a scan with no text layer, a photo of a
  sign, a figure with text in it. Training images come from the same PDFs as the text, labelled by
  the same documents; CAWST's picture-only materials and GlotLID's `zxx` give a "no language" class.
- **Fusion.** A gated sum first, which falls back to either modality alone when the other is
  missing; cross-attention from text tokens to image patches if that is not enough.
- **Shared head, shared labels.** One label space for both modalities.
- **Compatibility.** v1's keys are `text.*` and `head.*`; v2 adds `image.*` and `fuse.*`. v1
  weights load into v2 as they are; the text path can be frozen while the image path trains, then
  everything fine-tuned together.

## 9. Known limits of v1

- **Close languages** (Bemba and Lamba, Zulu and Ndebele, Arabic varieties, Malay and Indonesian)
  are hard by nature, and short passages of them are often genuinely ambiguous.
- **Labels from FinePDFs and FineWeb-2 are GlotLID's,** so on those sources MillID partly learns
  GlotLID's mistakes. UDHR is the check.
- **Genre.** For many small languages the web is mostly religious text, and the storybooks are
  children's books: a language can be learned together with its genre.
- **One language per passage.** A mixed passage gets the language of most of its text.