Pearl-Chat: Engineering a native Luganda language model from first principles in JAX and Flax NNX

Wesley Kambale; Roland Sankara

cs.CLOct 15, 2026

Pearl-Chat: Engineering a native Luganda language model from first principles in JAX and Flax NNX

Authors: Wesley Kambale, Roland Sankara

Organizations: AI for Uganda Initiative · ISBAT University

Abstract

Autoregressive foundation language models predominantly target high-resource Indo-European languages, imposing computational and morphological penalties on under-represented African languages. Generic multilingual tokenizers routinely exceed fertility rates of 3.2 tokens per word on agglutinative Bantu languages, degrading context efficiency and inflating latency. We present Pearl-Chat, an 8,473,122-parameter ( $\sim$ 8.5M) causal decoder-only transformer engineered for Luganda (Oluganda, ISO 639-2: lg), a low-resource Niger-Congo Bantu language spoken by over 10 million people in East Africa. Built from first mathematical principles in JAX and Flax NNX, Pearl-Chat uses explicit functional state decomposition (nnx.split and nnx.merge), deterministic resumable data streaming via Google Grain, and SPMD tensor parallelism using jax.sharding.Mesh. We introduce a byte-level byte-pair encoding (BPE) subword tokenizer trained from scratch on curated Luganda corpora ( $V = 8,192$ ), achieving 4.88 characters per token compression and 1.60 tokens per word fertility, a two-fold compression improvement over LLaMA-3 (128K) and Gemma (256K). The model is trained with a two-stage curriculum: Stage 1 pre-training on 59,501 sequences across parallel, biblical, encyclopedic, and journalistic corpora, and Stage 2 instruction fine-tuning on 510 dialogues with assistant-only loss masking. Across 4,000 optimization steps, validation cross-entropy loss drops monotonically from 9.4695 (perplexity: 12,958.98) to 0.1818 (perplexity: 1.20), a decrease of 9.2877 nats. On a three-tier held-out diagnostic benchmark, Pearl-Chat achieves 0.9792 ROUGE-L and 0.9405 BLEU while maintaining concordial agreement across Bantu noun classes. We provide a Python Server-Sent Events (SSE) streaming runtime and open-source all weights, tokenizers, and code.

Figures & tables

Class	Prefix	Concord	Example gloss
Cl 1 / 2	omu- / aba-	a- / ba-	omuntu / abantu (person / people)
Cl 3 / 4	omu- / emi-	gu- / gi-	omuti / emiti (tree / trees)
Cl 5 / 6	eri- / ama-	li- / ga-	erinnya / amannya (name / names)
Cl 7 / 8	eki- / ebi-	ki- / bi-	ekitabo / ebitabo (book / books)
Cl 9 / 10	en- / ezin-	e- / zi-	ente / ente (cow / cows)
Cl 11	olu-	lu-	olulimi (language / tongue)

Table 1: Luganda noun class taxonomy and concordial agreement template.

Figure 1: Subword tokenizer efficiency benchmark on native Luganda prose. Pearl-Chat BPE ( $V=8,192$ ) achieves a 4.88 characters per token compression ratio and 1.60 tokens per word fertility, outperforming 100K+ multilingual vocabularies.

Tokenizer	Vocabulary size	Compression (chars/tok)	Fertility (tok/word)	Total tokens
Pearl-Chat BPE (Ours)	8,192	4.88x	1.60 tok/w	16
Gemma BPE [ 10 ]	256,000	2.52x	3.10 tok/w	31
GPT-4o BPE [ 9 ]	200,000	2.44x	3.20 tok/w	32
LLaMA-3 BPE [ 4 ]	128,256	2.29x	3.40 tok/w	34
mBERT WordPiece [ 16 ]	119,547	2.00x	3.90 tok/w	39
Character baseline	256	1.00x	7.80 tok/w	78

Table 2: Cross-tokenizer compression benchmarks on a standard Luganda validation sentence (78 characters, 10 words).

Figure 2: Curriculum token volume and sequence distribution across multi-source Luganda corpora.

Corpus name	Records	Est. tokens	Domain
Luganda Bible corpus	30,844 verses	$\sim$ 1,620,000	Classical Bantu prose
Parallel corpus	24,973 pairs	$\sim$ 1,240,000	Contemporary bilingual
Luganda Wikipedia	350 articles	$\sim$ 380,000	Encyclopedic prose
Sunbird SALT MT	300 pairs	$\sim$ 18,000	Bilingual parallel
Curated SFT dialogues	510 dialogues	$\sim$ 35,000	Conversational chat
Total curated pool	60,017 records	$\sim$ 3,293,000	Multi-domain balanced

Table 3: Composition of training corpora and token allocations.

Figure 3: Pearl-Chat transformer architecture: (a) macroscopic end-to-end pipeline with weight-tied output embeddings ( $W_{\text{out}}=W_{E}^{T}$ ), and (b) internal Pre-Layer Normalization (Pre-LN) decoder block ( $l\in\{1,\dots,8\}$ ) detailing causal multi-head self-attention, GELU feed-forward network, and dual residual bypass streams.

Component	Tensor dimension	Trainable parameters
Token embedding $W_{E}$	$8,192\times 256$	2,097,152
Position embedding $W_{P}$	$256\times 256$	65,536
Transformer blocks ( $\times 8$ )	8 layers	6,306,816
– Attention ( $Q,K,V,O$ )	$4\times(256\times 256)$	2,097,152 (total across 8 layers)
– LayerNorms ( $\times 2$ )	$2\times 256$ (scale and bias)	32,768 (total across 8 layers)
– MLP ( $W_{1},W_{2}$ )	$(256\times 1024)+(1024\times 256)$	4,194,304 (total across 8 layers)

Table 4: Pearl-Chat model hyperparameters and parameter budget breakdown.

Figure 4: Explicit state graph decomposition in Flax NNX. The function nnx.split separates mutable module graphs into immutable static topology ( GraphDef ) and array parameter containers ( State ), recombined after XLA compilation using nnx.merge .

Dimension	PyTorch ( torch.nn )	Flax NNX ( flax.nnx )
State tracking	Implicit object mutation	Explicit state pytrees
Graph splitting	Deep copies and hooks	Functional nnx.split / merge
JIT compilation	TorchScript / Dynamo	Native nnx.jit via XLA
Autodiff	In-place .grad buffers	Functional nnx.value_and_grad
Sharding	Imperative DDP hooks	Declarative NamedSharding

Table 5: System paradigm comparison: PyTorch versus Flax NNX.

Figure 5: Logical 2D device mesh sharding topology. Activations are partitioned along the data axis using NamedSharding $P(\text{``data''},\text{``None''})$ , while linear weight kernels are sharded along the model axis using $P(\text{``None''},\text{``model''})$ .

Figure 6: Training trajectory across two-stage curriculum: (a) Cross-entropy loss dropping 9.2877 nats; (b) Validation perplexity decreasing from 12,959 to 1.20; (c) Optax warmup cosine decay schedule.

Step	Curriculum stage	Training loss	Validation loss	Validation PPL	Learning rate ( $\times 10^{-4}$ )
0	Uninitialized baseline	–	9.4695	12,958.98	0.00
1	Stage 1 (Pre-training)	9.4470	9.4695	12,958.98	0.06
250	Stage 1 (Pre-training)	6.1544	6.0726	433.80	5.98
500	Stage 1 (Pre-training)	5.7434	5.3179	203.96	5.85
1000	Stage 1 (Pre-training)	5.5017	4.9537	141.69	5.26
1500	Stage 1 (Pre-training)	5.4477	4.2903	72.99	4.31

Table 6: Empirical optimization metrics across curriculum steps.

Figure 7: Quantitative evaluation benchmark across Tiers 1, 2, and 3.

Evaluation tier	Exact match rate	ROUGE-L F1	BLEU-4 score
Tier 1: Memorization and syntactic integrity	100.0%	1.0000	1.0000
Tier 2: In-domain validation split	92.0%	0.9650	0.9420
Tier 3: Out-of-domain held-out generalization	83.3%	0.9792	0.9405

Table 7: Quantitative benchmark metrics across diagnostic tiers.

Configuration	Validation loss	Validation PPL	Concord accuracy	ROUGE-L F1
Pearl-Chat (Full curriculum)	0.1818	1.20	98.5%	0.9792
– No Stage 1 pre-training	2.9135	18.42	48.2%	0.3410
– No assistant loss masking	0.8420	2.32	82.1%	0.5215
– Small vocabulary ( $V=2,048$ )	0.6120	1.84	89.4%	0.6830
– Large vocabulary ( $V=16,384$ )	0.7240	2.06	84.7%	0.6310
– Post-LN baseline	1.4502	4.26	71.0%	0.4850

Table 8: Ablation experiments demonstrating impact of curriculum, masking, and vocabulary capacity.

Figure 8: Pearl-Chat serving and streaming inference systems architecture.

Appendix figures & tables2 assets

Supplementary material from the paper’s appendix.

Appendix

Parameter	Value	Description
Model type	Decoder-only	Causal autoregressive transformer
Vocab size ( $V$ )	8,192	Byte-level BPE tokens
Context length ( $T$ )	256	Maximum sequence tokens
Embed dim ( $d$ )	256	Hidden state dimension
Num layers ( $L$ )	8	Transformer decoder blocks
Num heads ( $h$ )	8	Attention heads ( $d_{k}=32$ )

Appendix

Table 9: Comprehensive model and optimization hyperparameters.

Target scale	Parameters	Optimal tokens	Hardware mesh	Estimated time
Pearl-Chat (Ours)	8.5M	3.3M (curated)	1x Tesla T4 GPU	227.2 sec
Pearl-Medium	35M	700M	4x TPU v5e	4.2 hours
Pearl-Large	125M	2.5B	16x TPU v5e	14.8 hours
Pearl-Bantu Multi	500M	10.0B	64x TPU v4	36.0 hours

Appendix

Table 10: Compute-optimal scaling projections for future African language foundation models.

Language Modeling

Language modeling develops systems that learn patterns in text and other sequences to understand, predict, and generate language.

See all 498 papers

AfriqueLLM: How Data Mixing and Model Architecture Impact Continued Pre-training for African Languages
Jan 10, 2026Hao Yu, Tianyi Xu, Michael A. Hedderich +3Language Model Pretraining Multilingual Language Models
McGill University, Canada · Mila-Quebec AI Institute, Canada · LMU Munich & Munich Center for Machine Learning, Germany +2
MzansiText and MzansiLM: An Open Corpus and Decoder-Only Language Model for South African Languages
Mar 21, 2026Anri Lombard, Simbarashe Mawere, Temi Aina +5African Language Processing Multilingual Language Models
University of Cape Town, Cape Town, South Africa
Mining Large Language Models for Low-Resource Language Data: Comparing Elicitation Strategies for Hausa and Fongbe
Apr 14, 2026Mahounan Pericles Adjovi, Roald Eiselen, Prasenjit MitraLow-Resource Language Processing LLM Prompting
Carnegie Mellon University Africa, Kigali, Rwanda · Centre for Text Technology, North-West University, Potchefstroom, South Africa

Jan 10, 2026J/K move · Enter open · S save

AfriqueLLM: How Data Mixing and Model Architecture Impact Continued Pre-training for African Languages

Hao Yu, Tianyi Xu, Michael A. Hedderich, Wassim Hamidouche+2

McGill University, Canada · Mila-Quebec AI Institute, Canada · LMU Munich & Munich Center for Machine Learning, Germany+2

Language Model Pretraining Multilingual Language Models Continued Pretraining Low-Resource Language Processing African Language Processing

Open paper AskPDF arXiv

Large language models (LLMs) are increasingly multilingual, yet open models continue to underperform relative to proprietary systems, with the gap most pronounced for African languages. Continued pre-training (CPT) offers a practical route to language adaptation, but improvements on demanding capabilities such as mathematical reasoning often remain limited. This limitation is driven in part by the uneven domain coverage and missing task-relevant knowledge that characterize many low-resource language corpora. We present AfriqueLLM, a suite of open LLMs adapted to 20 African languages through CPT on 26B tokens. We perform a comprehensive empirical study across five base models spanning sizes and architectures, including Llama 3.1, Gemma 3, and Qwen 3, and systematically analyze how CPT data composition shapes downstream performance. In particular, we vary mixtures that include math, code, and synthetic translated data, and evaluate the resulting models on a range of multilingual benchmarks. Our results identify data composition as the primary driver of CPT gains. Adding math, code, and synthetic translated data yields consistent improvements, including on reasoning-oriented evaluations. Within a fixed architecture, larger models typically improve performance, but architectural choices dominate scale when comparing across model families. Moreover, strong multilingual performance in the base model does not reliably predict post-CPT outcomes; robust architectures coupled with task-aligned data provide a more dependable recipe. Finally, our best models improve long-context performance, including document-level translation. Models and code have been released on Huggingface and Github.

Abstract

Language Modeling

Explore similar work

AfriqueLLM: How Data Mixing and Model Architecture Impact Continued Pre-training for African Languages

MzansiText and MzansiLM: An Open Corpus and Decoder-Only Language Model for South African Languages

Mining Large Language Models for Low-Resource Language Data: Comparing Elicitation Strategies for Hausa and Fongbe