mlnomad/yat-mmbert-base-contrastive-6000

YAT mmBERT base — multilingual contrastive checkpoint, step 6,000

This is a research sentence-embedding checkpoint from FlaxChat. It continues the YAT mmBERT MLM model through a 500-step contrastive pilot and a separate 6,000-step multilingual contrastive stage. It is a new set of weights, with a distinct training identity from the MLM release. It is intended for experiments in multilingual sentence retrieval; it has not been qualified as a general-purpose production embedding model or evaluated on official MTEB tasks.

The release includes model.safetensors, the exact YAT encoder config.json, tokenizer.json, and an export receipt in export.json. The weights use Flax NNX state paths and require FlaxChat's custom loader; they are not directly loadable with transformers.AutoModel. The 307.8M-parameter architecture has 22 bidirectional layers, width 768, 12 attention heads, a YAT GLU FFN, and YAT attention. The YAT kernel bias is fixed at 1, epsilon at 0.01, and alpha parameters are trainable. Compute is BF16, with FP32 residual paths and centered attention-score arithmetic.

Use

Install a FlaxChat revision containing flaxchat.public_encoder, plus its JAX/Flax dependencies. This example uses the bundled tokenizer and L2-normalizes the mean-pooled output for cosine retrieval:

from pathlib import Path

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

folder = Path(snapshot_download("mlnomad/yat-mmbert-base-contrastive-6000"))
tokenizer = Tokenizer.from_file(str(folder / "tokenizer.json"))
model = load_public_encoder(folder)
ids = jnp.asarray([tokenizer.encode("Hello world").ids], dtype=jnp.int32)
embeddings = model.pool(ids).astype(jnp.float32)
embeddings = embeddings / jnp.maximum(
    jnp.linalg.norm(embeddings, axis=-1, keepdims=True), 1e-12
)

model.pool averages non-padding token states, including special tokens. For batches of unequal lengths, pad with tokenizer ID 0. Training and reported retrieval evaluation used sequences of at most 128 tokens; the larger configured position limit has not been validated for this use.

Training data and method

The final stage used the official train split of all 28 English–other-language shards of Sentence Transformers' Global Voices parallel-sentence package, pinned to revision 4cc20add371f246bb1559b543f8b0dea178a1803. Source content is by Global Voices contributors under CC-BY-3.0; OPUS provided alignment and Sentence Transformers provided packaging. Please preserve that attribution when using or extending the training data. The checkpoint also inherits the earlier mmBERT and YAT MLM training lineage; see the MLM model card for its data and limitations.

The data builder selected up to 10,000 eligible pairs per language by deterministic hash, required 8–128 tokens on both sides, and removed duplicates. It excluded exact normalized text in the previous pilot's train and dev sets and the pinned 112-shard Tatoeba test. The resulting split has 197,827 train and 4,048 development pairs. Exact-text checks do not rule out paraphrase, substring, or earlier base-pretraining exposure. The training manifest SHA-256 is 9083c1ce5fdd6f6489e545e73e40543caa6439547467bf3c2d2f928dc56d4ab0; the recipe and source locks give the full preparation details.

The 6,000-step stage initialized model weights only from the contrastive pilot at step 500 and started a fresh AdamW optimizer. It used a global batch of 64, sequence length 128, peak learning rate 1e-5, 100 warmup steps, a cosine decay, and approximately 1.94 passes over the selected train pairs. Its objective was symmetric InfoNCE with global in-batch negatives, mean-pooled L2 embeddings, and masking of duplicate text and known positive groups. The pilot itself descended from the step-48,236 MLM checkpoint.

Evaluation

The same 4,048-pair Global Voices development split was scored on a physical TPU in both translation directions, with all dev rows in the candidate pool. Multiple known positives were accepted; cross-side identical text outside known positives was excluded. The earlier pilot was evaluated on this new split as an external-data baseline.

Checkpoint	English → translation recall@1	Translation → English recall@1
Parent contrastive pilot, step 500	0.9074	0.8943
This model, step 6,000	0.9654	0.9503

This is in-domain translation retrieval, not general semantic search. The sealed Tatoeba test has not been run on this step-6,000 checkpoint. There is no official MTEB, BEIR, or broad low-resource evaluation for this release, and the table does not establish superiority to mmBERT or other embedding models. English-pivot translated news may transfer poorly to other domains, asymmetric query–document tasks, or underrepresented languages. Training data can contain biases or harmful content. Evaluate the model for your own use before deployment.

Provenance

The source Orbax checkpoint is step 6,000 of encoder-contrastive-global-voices-28-0928. Its checkpoint metadata SHA-256 is 73b45e275ed625cc6276eb102f3f85ace369dbcb3f32ada60551116de13ceafe; the tokenizer SHA-256 is 197d4cc5406ee12cc50c8b5511f2393cc32d9db321545979ce041c1199178356. export.json records the exact safetensors SHA-256 and tensor count. Every exported tensor was checked against the checkpoint's per-tensor hash manifest. The FlaxChat repository contains the trainer, data builder, and evaluation harness; the training PR records this stage.