# 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
a table cell to a paragraph"] --> N["Unicode NFKC
whitespace collapsed"] N --> E["Per-byte encoding
byte · script · byte-in-character · character n-grams
B × L × 4 ids, L ≤ 256"] E --> EMB["Sum of four embeddings
B × L × d"] EMB --> STEM["Depthwise conv stem, k = 5
local byte patterns"] STEM --> DOWN["Strided conv, 2 bytes → 1 token
B × L/2 × d"] DOWN --> ENC["12 × pre-norm block
RMSNorm → attention with RoPE → RMSNorm → SwiGLU"] ENC --> POOL["RMSNorm, masked mean pool
z: B × d"] POOL --> HEAD["Language head
z · (label + language + script vectors)ᵀ + b"] HEAD --> OUT["Softmax over 1,641 labels
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 ሀ
ETHIOPIC SYLLABLE HA"] --> B1["byte E1
charpos 2"] C --> B2["byte 88
charpos 3"] C --> B3["byte 80
charpos 3"] S["script: ETHIOPIC
same id on all three bytes"] -.-> B1 & B2 & B3 G["n-grams ending at ሀ
2-gram, 3-gram, 4-gram ids
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
q, k, v in h heads of 64"] QKV --> ROPE["RoPE on q and k
θ = 10,000"] ROPE --> ATT["scaled dot-product attention
padding keys masked"] ATT --> O["Linear d → d"] O --> R1(("+")) X --> R1 R1 --> N2["RMSNorm"] N2 --> FF["SwiGLU
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
srp_Cyrl"] --> SUM(("+")) G["language vector
srp, shared with srp_Latn"] --> SUM S["script vector
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
1.9M+ passages"] --> SAMP["draw a language ∝ n^0.3
then one of its passages"] SAMP --> SHORT["25%: cut to a 20–80 character snippet"] SHORT --> WIN["random window of whole characters
≤ 255 bytes"] WIN --> ENC["four ids per byte"] ENC --> W["8 DataLoader workers
batches of 512"] W --> GPU["forward, cross-entropy
label smoothing 0.1, bf16"] GPU --> OPT["AdamW, clip 1.0
warm-up then cosine"] OPT --> VAL{"every 2,000 steps:
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
v1, text.*"] --> ZT["z_text: B × d"] IMG["page or crop"] --> IE["image encoder
ViT or ConvNeXt, image.*"] --> ZI["z_image: B × d"] ZT --> F["gated fusion, fuse.*
z = g·z_text + (1 − g)·z_image"] ZI --> F F --> H["LanguageHead
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.