opensearch-neural-sparse-en-jp-ko

Fine-tuned from opensearch-project/opensearch-neural-sparse-encoding-multilingual-v1 for cross-lingual sparse retrieval across English, Japanese, and Korean.

The training objective pulls non-English queries toward English vocabulary positions in the sparse output — so a Japanese sentence about "電池投資" and its English equivalent about "battery investment" activate similar vocab positions, enabling a single English anchor schema to retrieve across all three scripts.

Domain: automotive, supply-chain, and geopolitics.

This repository includes both the trained weights and the full training harness (training/), so you can reproduce, extend, or retarget the model to your own domain.

---

Quick usage

from transformers import AutoTokenizer, AutoModelForMaskedLM
import torch

tok = AutoTokenizer.from_pretrained("cp500/opensearch-neural-sparse-en-jp-ko")
model = AutoModelForMaskedLM.from_pretrained("cp500/opensearch-neural-sparse-en-jp-ko")

text = "トヨタは電池技術に多額の投資をしている。"
enc = tok(text, return_tensors="pt", return_token_type_ids=False)
logits = model(**enc).logits
activated = torch.log1p(torch.relu(logits))
activated = activated * enc["attention_mask"].unsqueeze(-1)
sparse = activated.max(dim=1).values
top = torch.topk(sparse[0], k=15)
for v, i in zip(top.values, top.indices):
    if v > 0:
        print(f"{{tok.convert_ids_to_tokens(int(i)):>20}}  {{float(v):.3f}}")

---

Loss functions

The training objective is a weighted sum of three terms:

loss = task_loss × (1 + (λ_xl − 1) × xl_fraction)
     + λ_q · FLOPS(query_batch)
     + λ_d · (FLOPS(positive_batch) + FLOPS(negative_batch))

1. Task loss — InfoNCE with in-batch negatives

The core retrieval signal. For each query in a batch of size B:

• Compute sparse representations: q_sparse, p_sparse, n_sparse via SPLADE's log(1 + ReLU(logits)) + attention-masked max-pool over token positions.
• Build scores against every positive in the batch plus one explicit hard negative: scores = concat(q @ p_batch.T, q · n) — shape (B, B + 1).
• Optimise multi-class cross-entropy with the correct label = row index (each query's own positive is the target).

Why cross-entropy and not MarginMSE? During our initial runs, MarginMSE with dot-product scores collapsed to near-empty sparse vectors: the model found it cheaper to shrink all scores toward zero than to learn proper margins. InfoNCE's softmax normalisation makes that strategy lose — the loss is invariant to uniform score scaling, so the only way to improve is to differentiate positives from negatives.

2. Cross-lingual pair up-weighting

The primary goal is cross-lingual alignment — pulling non-English queries onto English vocab positions. We scale the task loss by a per-batch factor proportional to how many cross-lingual pairs it contains:

xl_fraction = count(q_lang ≠ p_lang) / batch_size
task_weight = 1 + (λ_xl − 1) × xl_fraction            # λ_xl = 1.5 default

A batch with 100% cross-lingual pairs contributes 1.5× the gradient of a batch with 100% monolingual pairs. The dataset is deliberately skewed toward cross-lingual triplets (see the corpus card for the language-pair distribution).

3. FLOPS regularisation

Without regularisation, a sparse encoder can "cheat" by activating all vocab positions slightly — sparse in name only. FLOPS regularisation (from the SPLADE line of papers) penalises the squared L1 of the mean activation vector across a batch:

FLOPS(batch) = ‖mean_over_batch(|activations|)‖²_{L1-then-L2}
             = ( Σ_v ( (1/B) Σ_b |a_{b,v}| ) )² computed per vocab-col then
                                                  .pow(2).sum() across cols

This is asymmetrically applied: queries are pushed sparser (λ_q = 3e−4 by default) than documents (λ_d = 1e−4), consistent with the SPLADE production convention that short queries should activate fewer tokens than long documents.

4. FLOPS lambda warm-up

Applying full FLOPS regularisation from step 0 kills the model — it happily zeroes the output before the task loss has started pulling any signal through the layers. We ramp λ_q and λ_d quadratically from 0 to their target values over the first flops_warmup_steps steps (default: min(total_steps / 3, 2000)). This lets the model first establish which vocab positions matter for retrieval, then compress.

5. What we did NOT use and why

• Knowledge distillation from a cross-encoder teacher (MarginMSE canonical recipe): we have no multilingual automotive cross-encoder with trustworthy teacher scores. InfoNCE is the honest fallback when teacher scores aren't available.
• Hard negative mining from a first-pass encoder: the five hard negatives per concept were generated directly by the data-gen prompt, not mined post hoc. This keeps the pipeline simple and reproducible.
• Normalisation of sparse vectors before dot product: SPLADE-v3 and prithivida/Splade_PP both use raw dot product. InfoNCE's softmax makes magnitude-robustness moot.

---

Training configuration

• Base model: opensearch-project/opensearch-neural-sparse-encoding-multilingual-v1 (XLM-RoBERTa, 160M params, 105,879 vocab)
• Dataset: cp500/multilingual-automotive-sparse — 67,500 (query, positive, negative) triplets across 9 EN/JA/KO language-pair combinations
• Hardware: 1× NVIDIA L40S (g6e.2xlarge)
• Epochs: 3
• Batch size: 16
• Optimizer: AdamW, lr = 2e-5, weight decay = 0.01
• LR schedule: Linear with 6% warm-up
• Temperature: 1.0 (InfoNCE softmax)
• Cross-lingual weight (λ_xl): 1.5
• FLOPS λ_query (target): 3e-4
• FLOPS λ_doc (target): 1e-4
• FLOPS warm-up: 2000 quadratic steps before reaching target λ

---

Evaluation

Held-out split: 2,241 (query, positive) pairs across 9 language-pair combinations.

Pair	n	MRR@10	Recall@10
en->en	499	0.7752	0.9940
en->ja	499	0.2844	0.9880
en->ko	499	0.2955	0.9940
ja->en	499	0.4278	0.9800
ja->ja	499	0.6164	0.9960
ja->ko	499	0.3091	0.9820
ko->en	499	0.4345	0.9820
ko->ja	499	0.3253	0.9880
ko->ko	499	0.5826	0.9920
overall	4491	0.4501	0.9884

Mean activations / query: 74 Mean activations / passage: 125

EN-vocab activation ratio by query language (top-50 active dims):

• en: 0.578
• ja: 0.555
• ko: 0.565 Non-EN mean: 0.560

Cross-lingual success metric (the point of this model): the mean fraction of the top-50 active vocab positions on JA/KO queries that correspond to English WordPiece tokens. The base XLM-R sparse encoder sits near 0.05–0.15 on this metric — its sparse output stays in the input language's vocab space. Successful cross-lingual fine-tuning should push this toward 0.30+.

---

Reproducing from scratch

The full pipeline — from seed concepts to a trained model — is included in the training/ directory. Four steps:

Step 1. Generate the training corpus

pip install boto3                                  # for Bedrock
aws sso login                                      # whatever auth your account uses
python training/generate_corpus.py --target 2500   # ~1 hour, ~$20-25 in Bedrock
python training/flatten_triplets.py                # 2500 concepts → ~67k triplets

The concept seeds live in training/concept_seeds.py — edit the OEMS/SUPPLIERS/MATERIALS/etc. lists there to retarget the corpus to a different domain.

Step 2. (Optional) Push the corpus to HuggingFace

See cp500/multilingual-automotive-sparse for the reference dataset. If you want to publish your own variant:

huggingface-cli login
python training/push_corpus.py --repo-id <your-user>/<your-dataset>

Step 3. Train on a GPU

huggingface-cli login              # once, caches your write token
pip install -r training/requirements.txt
bash training/run_on_instance.sh   # train + eval + push; ~1 hour on L40S

Or run train.py directly for full control:

python training/train.py \
    --dataset cp500/multilingual-automotive-sparse \
    --base-model opensearch-project/opensearch-neural-sparse-encoding-multilingual-v1 \
    --epochs 3 --batch-size 16 \
    --flops-lambda-q 3e-4 --flops-lambda-d 1e-4 \
    --crosslingual-weight 1.5 \
    --output-dir checkpoints/

Step 4. Evaluate

python training/eval.py --model checkpoints/final --out eval_summary.json

The eval script computes per-language-pair MRR@10 and Recall@10, plus the cross-lingual success metric (English-vocab activation ratio on JA/KO inputs).

---

License

Apache 2.0. The base model, the training corpus, and this fine-tuned model are all Apache 2.0. The corpus was generated via Anthropic's Claude Haiku through AWS Bedrock — review Anthropic's Acceptable Use Policy for downstream applications.