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: 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.

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

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

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.

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

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

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	2,000+	fastText, 1.7 GB	--compare-glotlid
fastText, Meta (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

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.