mlnomad/yat-mmbert-base-mlm-48236

YAT mmBERT base — MLM checkpoint at step 48,236

This is a research checkpoint of a multilingual bidirectional encoder trained with masked language modeling (MLM) in FlaxChat. It starts from jhu-clsp/mmBERT-base weights and replaces the feed-forward gate and attention normalization with the project's YAT implementation. This checkpoint has not received contrastive or retrieval training. Its raw mean-pooled embeddings currently trail the released mmBERT baseline in our downstream evaluation; it is not production-qualified.

The release contains the complete model state in model.safetensors, the exact encoder config.json, and tokenizer.json. Optimizer state and training data are not included. model.safetensors uses Flax NNX state paths and is not directly loadable with transformers.AutoModel. The custom architecture and loader live in FlaxChat; use a revision containing flaxchat.public_encoder:

from huggingface_hub import snapshot_download
from flaxchat.public_encoder import load_public_encoder
from tokenizers import Tokenizer
import jax.numpy as jnp
from pathlib import Path

folder = Path(snapshot_download("mlnomad/yat-mmbert-base-mlm-48236"))
model = load_public_encoder(folder)
tokenizer = Tokenizer.from_file(str(folder / "tokenizer.json"))
ids = jnp.asarray([tokenizer.encode("Hello world").ids], dtype=jnp.int32)
embeddings = model.pool(ids)

The example uses the bundled tokenizer and the custom loader from FlaxChat. pool averages non-padding token states, including special tokens. The checkpoint preserves the MLM head for further masked-token training; model.encode(ids) returns token representations.

Architecture and training

Property	Value
Layers / hidden size / attention heads	22 / 768 / 12
FFN intermediate size	1,152
Vocabulary / maximum configured context	256,000 / 8,192 tokens
Attention	Bidirectional local/global YAT softmax
FFN	YAT GLU
YAT bias / epsilon	Fixed 1 / 0.01
YAT alpha	Trainable
Precision	BF16 compute with FP32 residuals and centered FP32 attention-score arithmetic
Objective	MLM only; no contrastive stage
Final checkpoint	Step 48,236

The final training stage resumed model weights from an earlier step-23,467 checkpoint and started a new optimizer and data schedule. It used a 32-example global batch, 15% masking, AdamW with 2e-5 peak learning rate, and a multilingual text mixture. The held-out MLM negative log-likelihood reported by the training harness was 1.28655. That loss is an internal gate, not a measure of retrieval quality. The tokenizer file SHA-256 is 197d4cc5406ee12cc50c8b5511f2393cc32d9db321545979ce041c1199178356; the frozen training source Python SHA-256 was c7c04f9d7b84e2bfe9a2387917f75b096a5dabf3830a16dd50f06bbfa350534b.

The training mixture included multilingual web text from FineWeb2-HQ and other prepared sources. Because this model initialized from mmBERT-base, it inherits that model's original multilingual pretraining. Our additional MLM stage did not replay mmBERT's full published 3T-token data mixture or its staged language and masking schedule; see the mmBERT paper. This model card does not claim that every source has been exhaustively documented for redistribution; the training data itself is not included. FineWeb2-HQ is listed as ODC-BY, and upstream mmBERT as MIT in their respective repository cards. Please review source licenses and suitability for your use.

Evaluation and limitations

All numbers below are from our pinned, custom evaluations on TPU, with identical mean pooling and cosine similarity for the compared encoders. They are not official MTEB or fine-tuned XNLI scores.

Evaluation	This checkpoint	Released mmBERT-base
Native Tatoeba bitext, 112 subsets / 88,877 pairs, subset-macro support-weighted F1	0.1140	0.3016
Native Tatoeba bitext pair accuracy	0.1393	0.3390
Frozen XNLI English-trained ridge probe, 15-language macro accuracy	0.3826	0.4085

The Tatoeba corpus was used for evaluation, not model selection for this release. The XNLI probe trained a shallow classifier on 1,024 English examples and tested 512 examples per language; it did not fine-tune either encoder. The comparison shows that more MLM pretraining did not yet make this YAT checkpoint a competitive sentence embedding model. We intend to train a separate contrastive successor, with a fresh training identity and evaluation against the same baseline.

The model may perform poorly on retrieval, semantic similarity, low-resource languages, and text outside its training distribution. It can reflect biases or harmful content from web training data. Validate it for your application; do not infer production quality from the MLM loss or the configured context limit.

Reproducibility

The exact weight-only export records its source checkpoint step, tensor count, safetensors SHA-256, and tokenizer identity in export.json. Every exported tensor was checked against the original checkpoint's per-tensor SHA-256 manifest before packaging. The FlaxChat source contains the training and benchmark harness. For citation, cite the mmBERT model, FineWeb2-HQ, and this model repository with its commit revision.