Hugging Face 发布多语言编码器模型 mmBERT,覆盖 1833 种语言
mmBERT: ModernBERT goes Multilingual
Hugging Face 与 JHU 团队发布多语言编码器模型 mmBERT,基于 ModernBERT 架构,在 3T+ tokens、1833 种语言的数据上训练,官方称其首次在多语言任务上超越 XLM-R。
原文给出了 1833 种语言的训练策略和基准数据,读者可以据此评估它对多语言编码器选型的影响。
TL;DR
这篇博客文章介绍了 mmBERT,这是一个最先进的大规模多语言编码器模型,在超过 1800 种语言的 3T+ token 文本上训练。它相比之前的多语言模型展现出显著的性能和速度提升,是首个超越 XLM-R 的模型,同时还开发了有效学习低资源语言的新策略。mmBERT 基于 ModernBERT 构建,实现了极快的架构,并添加了新颖的组件以实现高效的多语言学习。
如果你有兴趣亲自尝试这些模型,一些示例样板代码可在本博客文章末尾获取!
训练数据
mmBERT 在精心策划的多语言数据集上训练,总计超过 3T token,分为三个不同的训练阶段。我们训练数据的基础由三个主要的开源高质量网络爬取数据组成,它们同时实现了多语言覆盖和数据质量:
DCLM 和 Filtered DCLM 提供了可用的最高质量英文内容,是强大英文性能的支柱(过滤后的数据来自 Dolmino)。该数据集代表了最先进的网络过滤技术,是一个关键组成部分。由于该数据质量很高,我们使用的英文比例显著高于上一代多语言编码器模型(高达 18%)。
FineWeb2 提供了广泛的多语言网络内容,覆盖超过 1,800 种语言。该数据集使我们能够实现广泛的多语言覆盖,同时在多样化的语系和文字系统中保持合理的质量标准。
FineWeb2-HQ 由 FineWeb2 的过滤子集组成,聚焦于 20 种高资源语言。这个过滤版本提供了更高质量的多语言内容,弥合了仅英文过滤数据与广泛多语言覆盖之间的差距。
训练数据还纳入了来自 Dolma、MegaWika v2、ProLong 等的专门语料库:代码仓库(StarCoder、ProLong)、学术内容(ArXiv、PeS2o)、参考资料(Wikipedia、教科书)和社区讨论(StackExchange),以及指令和数学数据集。
我们数据方法的关键创新是图 1 所示的渐进式语言纳入策略。在每个阶段,我们逐步从更平坦的分布(即更接近均匀)中采样,同时添加新语言。这意味着像俄语这样的高资源语言一开始占数据的很高比例(即 9%),然后在训练的最后阶段结束时约为其一半。我们在预训练期间从 60 种高资源语言开始,在中期训练期间扩展到 110 种语言,最后在衰减阶段纳入 FineWeb2 中全部 1,833 种语言。这使我们能够最大化有限低资源语言数据的影响,避免过度重复,同时保持高整体数据质量。
训练配方和新颖组件
mmBERT 基于 ModernBERT 架构构建,但为多语言学习引入了若干关键创新:
架构
我们采用与 ModernBERT-base 相同的核心架构,具有 22 层和 1152 个中间维度,但改用 Gemma 2 分词器以更好地处理多语言文本。基础模型有 110M 非嵌入参数(由于词汇量更大,总计 307M),而小型变体有 42M 非嵌入参数(总计 140M)。
三阶段训练方法
我们的训练遵循精心设计的三阶段计划:
- 预训练(2.3T tokens):预热和稳定学习率阶段,使用 60 种语言,掩码率为 30%
- 中期训练(600B tokens):上下文扩展到 8192 tokens,更高质量的数据,扩展到 110 种语言,掩码率为 15%
- 衰减阶段(100B tokens):逆平方根学习率衰减,包含全部 1,833 种语言,掩码率为 5%
新颖训练技术
逆掩码率计划:我们不是使用固定的掩码率,而是在训练阶段逐步将掩码率从 30% → 15% → 5% 降低。这允许模型在早期使用较高掩码学习基本表示,然后在较低掩码率下专注于更细致的理解。
退火语言学习:我们动态调整多语言数据采样的温度,从 τ=0.7 → 0.5 → 0.3。这创造了从高资源语言偏置向更均匀采样的过渡,使模型能够在学习低资源语言之前建立强大的多语言基础。
渐进式语言添加:我们不是同时训练所有语言,而是在每个阶段策略性地添加语言(60 → 110 → 1,833)。这通过避免在有限的低资源数据上过度训练,同时仍实现强大性能,从而最大化学习效率。
模型合并:我们在衰减阶段训练三个不同的变体(英语聚焦、110 种语言和所有语言),并使用 TIES 合并将它们的长处结合到最终模型中。
结果
自然语言理解(NLU)
英语性能:在英语 GLUE 基准(表 1)上,mmBERT base 实现了强劲性能,大幅优于其他多语言模型如 XLM-R(多语言 RoBERTa)base 和 mGTE base,同时尽管 mmBERT 训练数据中英语不到 25%,仍与仅英语模型保持竞争力。
多语言性能:如表 2所示,与 XLM-R 相比,mmBERT 在 XTREME 基准上显示出显著改进。显著增益包括在 XNLI 分类上的强劲性能、在 TyDiQA 等问答任务上的大幅改进,以及在 PAWS-X 和 XCOPA 上跨语言理解的竞争性结果。
该模型在大多数类别上表现良好,除了 NER 和 POS 标记等一些结构化预测任务,这可能是由于分词器差异影响了词边界检测。在这些类别上,它的表现与上一代大致相同,但可以应用于更多语言。
检索性能
英语检索:尽管 mmBERT 是为大规模多语言场景设计的,但在 MTEB v2 英语基准测试中(表 3),mmBERT 相比之前的多语言模型展现出显著提升,甚至与 ModernBERT 等纯英语模型的能力相当!
多语言检索:与其他模型相比,mmBERT 在 MTEB v2 多语言基准测试中展现出持续改进(表 4)。
代码检索:得益于现代分词器(基于 Gemma 2),mmBERT 也展现出强大的代码性能(表 5),使其适用于任何类型的文本数据。唯一超越它的模型是 EuroBERT,该模型能够使用非公开访问的 Stack v2 数据集。
在衰减阶段学习语言
mmBERT 最重要的新颖特性之一是证明了低资源语言可以在训练的短暂衰减阶段被有效学习。我们通过在最后 100B token 衰减阶段才引入的语言上进行测试,验证了这一方法。
显著的性能提升:在 TiQuaD(提格里尼亚语)和 FoQA(法罗语)上的测试中,当这些语言被纳入衰减阶段时,我们观察到了大幅改进,如图 2所示。结果证明了我们渐进式语言学习方法的有效性。
与大型模型竞争:尽管仅在最终训练阶段才见到这些语言,mmBERT 仍达到了超越许多更大模型的性能水平。在已对 LLM 进行基准测试的法罗语问答任务上,mmBERT 的表现优于 Google Gemini 2.5 Pro 和 OpenAI o3。
快速学习机制:衰减阶段语言学习的成功源于模型能够利用其在早期阶段构建的强大多语言基础。当接触新语言时,模型可以快速调整现有的跨语言表示,而非从头学习。
模型合并的优势:最终的 mmBERT 模型通过 TIES 合并,成功保留了衰减阶段的大部分改进,同时受益于英语聚焦和高资源变体。
效率提升
mmBERT 通过继承自 ModernBERT 的架构改进,相比之前的多语言编码器模型实现了显著的效率提升:
吞吐性能:如图 3所示,mmBERT 在各种序列长度下处理文本的速度都显著快于现有多语言模型。小型和基础模型相比之前的多语言编码器均展现出大幅速度提升。
现代架构的优势:效率提升来自两项主要技术改进:
- Flash Attention 2:优化的注意力计算,带来更好的内存使用和速度
- 去填充技术:消除处理过程中不必要的填充 token
序列长度扩展:与仅限于 512 个 token 的旧模型不同,mmBERT 可高效处理多达 8,192 个 token,同时保持高吞吐量。这使其适用于多语言应用中日益常见的较长文档处理任务。
能效:更好的吞吐量与现代架构相结合,降低了推理的计算成本,使 mmBERT 在需要大规模多语言支持的生产部署中更加实用。
这些效率提升使 mmBERT 不仅比以往的多语言编码器更准确,而且在实际使用中也显著更加实用。
使用示例
只需几行代码即可使用这些模型!
from transformers import AutoTokenizer, AutoModelForMaskedLM
import torch
tokenizer = AutoTokenizer.from_pretrained("jhu-clsp/mmBERT-base")
model = AutoModelForMaskedLM.from_pretrained("jhu-clsp/mmBERT-base")
def predict_masked_token(text):
inputs = tokenizer(text, return_tensors="pt")
with torch.no_grad():
outputs = model(**inputs)
mask_indices = torch.where(inputs["input_ids"] == tokenizer.mask_token_id)
predictions = outputs.logits[mask_indices]
top_tokens, top_indices = torch.topk(predictions, 5, dim=-1)
return [tokenizer.decode(token) for token in top_indices[0]]
# Works across languages
texts = [
"The capital of France is <mask>.",
"La capital de España es <mask>.",
"Die Hauptstadt von Deutschland ist <mask>.",
]
for text in texts:
predictions = predict_masked_token(text)
print(f"Text: {text}")
print(f"Predictions: {predictions}\n")
微调示例
编码器
点击查看如何使用 Sentence Transformers 将其微调为稠密嵌入模型
import argparse
from datasets import load_dataset
from sentence_transformers import (
SentenceTransformer,
SentenceTransformerTrainer,
SentenceTransformerTrainingArguments,
)
from sentence_transformers.evaluation import TripletEvaluator
from sentence_transformers.losses import CachedMultipleNegativesRankingLoss
from sentence_transformers.training_args import BatchSamplers
def main():
# parse the lr & model name
parser = argparse.ArgumentParser()
parser.add_argument("--lr", type=float, default=8e-5)
parser.add_argument("--model_name", type=str, default="jhu-clsp/mmBERT-small")
args = parser.parse_args()
lr = args.lr
model_name = args.model_name
model_shortname = model_name.split("/")[-1]
# 1. Load a model to finetune
model = SentenceTransformer(model_name)
# 2. Load a dataset to finetune on
dataset = load_dataset(
"sentence-transformers/msmarco-co-condenser-margin-mse-sym-mnrl-mean-v1",
"triplet-hard",
split="train",
)
dataset_dict = dataset.train_test_split(test_size=1_000, seed=12)
train_dataset = dataset_dict["train"].select(range(1_250_000))
eval_dataset = dataset_dict["test"]
# 3. Define a loss function
loss = CachedMultipleNegativesRankingLoss(model, mini_batch_size=16) # Increase mini_batch_size if you have enough VRAM
run_name = f"{model_shortname}-DPR-{lr}"
# 4. (Optional) Specify training arguments
args = SentenceTransformerTrainingArguments(
# Required parameter:
output_dir=f"output/{model_shortname}/{run_name}",
# Optional training parameters:
num_train_epochs=1,
per_device_train_batch_size=512,
per_device_eval_batch_size=512,
warmup_ratio=0.05,
fp16=False, # Set to False if GPU can't handle FP16
bf16=True, # Set to True if GPU supports BF16
batch_sampler=BatchSamplers.NO_DUPLICATES, # (Cached)MultipleNegativesRankingLoss benefits from no duplicates
learning_rate=lr,
# Optional tracking/debugging parameters:
save_strategy="steps",
save_steps=500,
save_total_limit=2,
logging_steps=500,
run_name=run_name, # Used in `wandb`, `tensorboard`, `neptune`, etc. if installed
)
# 5. (Optional) Create an evaluator & evaluate the base model
dev_evaluator = TripletEvaluator(
anchors=eval_dataset["query"],
positives=eval_dataset["positive"],
negatives=eval_dataset["negative"],
name="msmarco-co-condenser-dev",
)
dev_evaluator(model)
# 6. Create a trainer & train
trainer = SentenceTransformerTrainer(
model=model,
args=args,
train_dataset=train_dataset,
eval_dataset=eval_dataset,
loss=loss,
evaluator=dev_evaluator,
)
trainer.train()
# 7. (Optional) Evaluate the trained model on the evaluator after training
dev_evaluator(model)
# 8. Save the model
model.save_pretrained(f"output/{model_shortname}/{run_name}/final")
# 9. (Optional) Push it to the Hugging Face Hub
model.push_to_hub(run_name, private=False)
if __name__ == "__main__":
main()
点击查看如何使用 PyLate 将其微调为多向量嵌入模型
from datasets import load_dataset
from pylate import losses, models, utils
from sentence_transformers import (
SentenceTransformerTrainer,
SentenceTransformerTrainingArguments,
)
def main():
# Load the datasets required for knowledge distillation (train, queries, documents)
train = load_dataset(
path="lightonai/ms-marco-en-bge",
name="train",
)
queries = load_dataset(
path="lightonai/ms-marco-en-bge",
name="queries",
)
documents = load_dataset(
path="lightonai/ms-marco-en-bge",
name="documents",
)
# Set the transformation to load the documents/queries texts using the corresponding ids on the fly
train.set_transform(
utils.KDProcessing(queries=queries, documents=documents).transform,
)
# Define the base model, training parameters, and output directory
num_train_epochs = 1
lr = 8e-5
batch_size = 16
accum_steps = 1
model_name = "jhu-clsp/mmBERT-small"
model_shortname = model_name.split("/")[-1]
# Set the run name for logging and output directory
run_name = f"{model_shortname}-colbert-KD-{lr}"
output_dir = f"output/{model_shortname}/{run_name}"
# Initialize the ColBERT model from the base model
model = models.ColBERT(model_name_or_path=model_name)
# Configure the training arguments (e.g., epochs, batch size, learning rate)
args = SentenceTransformerTrainingArguments(
output_dir=output_dir,
num_train_epochs=num_train_epochs,
per_device_train_batch_size=batch_size,
fp16=False, # Set to False if you get an error that your GPU can't run on FP16
bf16=True, # Set to True if you have a GPU that supports BF16
run_name=run_name,
logging_steps=10,
learning_rate=lr,
gradient_accumulation_steps=accum_steps,
warmup_ratio=0.05,
)
# Use the Distillation loss function for training
train_loss = losses.Distillation(model=model)
# Initialize the trainer
trainer = SentenceTransformerTrainer(
model=model,
args=args,
train_dataset=train,
loss=train_loss,
data_collator=utils.ColBERTCollator(tokenize_fn=model.tokenize),
)
# Start the training process
trainer.train()
model.save_pretrained(f"{output_dir}/final")
if __name__ == "__main__":
main()
点击查看如何使用 Sentence Transformers 将其微调为稀疏检索模型
import logging
from datasets import load_dataset
from sentence_transformers import (
SparseEncoder,
SparseEncoderModelCardData,
SparseEncoderTrainer,
SparseEncoderTrainingArguments,
)
from sentence_transformers.sparse_encoder.evaluation import SparseNanoBEIREvaluator
from sentence_transformers.sparse_encoder.losses import SparseMultipleNegativesRankingLoss, SpladeLoss
from sentence_transformers.training_args import BatchSamplers
logging.basicConfig(format="%(asctime)s - %(message)s", datefmt="%Y-%m-%d %H:%M:%S", level=logging.INFO)
# 1. Load a model to finetune with 2. (Optional) model card data
model = SparseEncoder(
"jhu-clsp/mmBERT-small",
model_card_data=SparseEncoderModelCardData(
language="en",
license="apache-2.0",
)
)
# 3. Load a dataset to finetune on
full_dataset = load_dataset("sentence-transformers/natural-questions", split="train").select(range(100_000))
dataset_dict = full_dataset.train_test_split(test_size=1_000, seed=12)
train_dataset = dataset_dict["train"]
eval_dataset = dataset_dict["test"]
# 4. Define a loss function
loss = SpladeLoss(
model=model,
loss=SparseMultipleNegativesRankingLoss(model=model),
query_regularizer_weight=5e-5,
document_regularizer_weight=3e-5,
)
# 5. (Optional) Specify training arguments
run_name = "splade-distilbert-base-uncased-nq"
args = SparseEncoderTrainingArguments(
# Required parameter:
output_dir=f"models/{run_name}",
# Optional training parameters:
num_train_epochs=1,
per_device_train_batch_size=16,
per_device_eval_batch_size=16,
learning_rate=2e-5,
warmup_ratio=0.1,
fp16=True, # Set to False if you get an error that your GPU can't run on FP16
bf16=False, # Set to True if you have a GPU that supports BF16
batch_sampler=BatchSamplers.NO_DUPLICATES, # MultipleNegativesRankingLoss benefits from no duplicate samples in a batch
# Optional tracking/debugging parameters:
eval_strategy="steps",
eval_steps=1000,
save_strategy="steps",
save_steps=1000,
save_total_limit=2,
logging_steps=200,
run_name=run_name, # Will be used in W&B if `wandb` is installed
)
# 6. (Optional) Create an evaluator & evaluate the base model
dev_evaluator = SparseNanoBEIREvaluator(dataset_names=["msmarco", "nfcorpus", "nq"], batch_size=16)
# 7. Create a trainer & train
trainer = SparseEncoderTrainer(
model=model,
args=args,
train_dataset=train_dataset,
eval_dataset=eval_dataset,
loss=loss,
evaluator=dev_evaluator,
)
trainer.train()
# 8. Evaluate the model performance again after training
dev_evaluator(model)
# 9. Save the trained model
model.save_pretrained(f"models/{run_name}/final")
# 10. (Optional) Push it to the Hugging Face Hub
model.push_to_hub(run_name)
点击查看如何使用 Sentence Transformers 将其微调为重排序模型
import logging
import traceback
import torch
from datasets import load_dataset
from sentence_transformers import SentenceTransformer
from sentence_transformers.cross_encoder import (
CrossEncoder,
CrossEncoderModelCardData,
CrossEncoderTrainer,
CrossEncoderTrainingArguments,
)
from sentence_transformers.cross_encoder.evaluation import (
CrossEncoderNanoBEIREvaluator,
CrossEncoderRerankingEvaluator,
)
from sentence_transformers.cross_encoder.losses import BinaryCrossEntropyLoss
from sentence_transformers.evaluation import SequentialEvaluator
from sentence_transformers.util import mine_hard_negatives
# Set the log level to INFO to get more information
logging.basicConfig(format="%(asctime)s - %(message)s", datefmt="%Y-%m-%d %H:%M:%S", level=logging.INFO)
def main():
model_name = "jhu-clsp/mmBERT-small"
train_batch_size = 64
num_epochs = 1
num_hard_negatives = 5 # How many hard negatives should be mined for each question-answer pair
# 1a. Load a model to finetune with 1b. (Optional) model card data
model = CrossEncoder(
model_name,
model_card_data=CrossEncoderModelCardData(
language="en",
license="apache-2.0",
),
)
print("Model max length:", model.max_length)
print("Model num labels:", model.num_labels)
# 2a. Load the GooAQ dataset: https://huggingface.co/datasets/sentence-transformers/gooaq
logging.info("Read the gooaq training dataset")
full_dataset = load_dataset("sentence-transformers/gooaq", split="train").select(range(100_000))
dataset_dict = full_dataset.train_test_split(test_size=1_000, seed=12)
train_dataset = dataset_dict["train"]
eval_dataset = dataset_dict["test"]
logging.info(train_dataset)
logging.info(eval_dataset)
# 2b. Modify our training dataset to include hard negatives using a very efficient embedding model
embedding_model = SentenceTransformer("sentence-transformers/static-retrieval-mrl-en-v1", device="cpu")
hard_train_dataset = mine_hard_negatives(
train_dataset,
embedding_model,
num_negatives=num_hard_negatives, # How many negatives per question-answer pair
margin=0, # Similarity between query and negative samples should be x lower than query-positive similarity
range_min=0, # Skip the x most similar samples
range_max=100, # Consider only the x most similar samples
sampling_strategy="top", # Sample the top negatives from the range
batch_size=4096, # Use a batch size of 4096 for the embedding model
output_format="labeled-pair", # The output format is (query, passage, label), as required by BinaryCrossEntropyLoss
use_faiss=True,
)
logging.info(hard_train_dataset)
# 2c. (Optionally) Save the hard training dataset to disk
# hard_train_dataset.save_to_disk("gooaq-hard-train")
# Load again with:
# hard_train_dataset = load_from_disk("gooaq-hard-train")
# 3. Define our training loss.
# pos_weight is recommended to be set as the ratio between positives to negatives, a.k.a. `num_hard_negatives`
loss = BinaryCrossEntropyLoss(model=model, pos_weight=torch.tensor(num_hard_negatives))
# 4a. Define evaluators. We use the CrossEncoderNanoBEIREvaluator, which is a light-weight evaluator for English reranking
nano_beir_evaluator = CrossEncoderNanoBEIREvaluator(
dataset_names=["msmarco", "nfcorpus", "nq"],
batch_size=train_batch_size,
)
# 4b. Define a reranking evaluator by mining hard negatives given query-answer pairs
# We include the positive answer in the list of negatives, so the evaluator can use the performance of the
# embedding model as a baseline.
hard_eval_dataset = mine_hard_negatives(
eval_dataset,
embedding_model,
corpus=full_dataset["answer"], # Use the full dataset as the corpus
num_negatives=30, # How many documents to rerank
batch_size=4096,
include_positives=True,
output_format="n-tuple",
use_faiss=True,
)
logging.info(hard_eval_dataset)
reranking_evaluator = CrossEncoderRerankingEvaluator(
samples=[
{
"query": sample["question"],
"positive": [sample["answer"]],
"documents": [sample[column_name] for column_name in hard_eval_dataset.column_names[2:]],
}
for sample in hard_eval_dataset
],
batch_size=train_batch_size,
name="gooaq-dev",
# Realistic setting: only rerank the positives that the retriever found
# Set to True to rerank *all* positives
always_rerank_positives=False,
)
# 4c. Combine the evaluators & run the base model on them
evaluator = SequentialEvaluator([reranking_evaluator, nano_beir_evaluator])
evaluator(model)
# 5. Define the training arguments
short_model_name = model_name if "/" not in model_name else model_name.split("/")[-1]
run_name = f"reranker-{short_model_name}-gooaq-bce"
args = CrossEncoderTrainingArguments(
# Required parameter:
output_dir=f"models/{run_name}",
# Optional training parameters:
num_train_epochs=num_epochs,
per_device_train_batch_size=train_batch_size,
per_device_eval_batch_size=train_batch_size,
learning_rate=2e-5,
warmup_ratio=0.1,
fp16=False, # Set to False if you get an error that your GPU can't run on FP16
bf16=True, # Set to True if you have a GPU that supports BF16
dataloader_num_workers=4,
load_best_model_at_end=True,
metric_for_best_model="eval_gooaq-dev_ndcg@10",
# Optional tracking/debugging parameters:
eval_strategy="steps",
eval_steps=1000,
save_strategy="steps",
save_steps=1000,
save_total_limit=2,
logging_steps=200,
logging_first_step=True,
run_name=run_name, # Will be used in W&B if `wandb` is installed
seed=12,
)
# 6. Create the trainer & start training
trainer = CrossEncoderTrainer(
model=model,
args=args,
train_dataset=hard_train_dataset,
loss=loss,
evaluator=evaluator,
)
trainer.train()
# 7. Evaluate the final model, useful to include these in the model card
evaluator(model)
# 8. Save the final model
final_output_dir = f"models/{run_name}/final"
model.save_pretrained(final_output_dir)
# 9. (Optional) save the model to the Hugging Face Hub!
# It is recommended to run `huggingface-cli login` to log into your Hugging Face account first
try:
model.push_to_hub(run_name)
except Exception:
logging.error(
f"Error uploading model to the Hugging Face Hub:\n{traceback.format_exc()}To upload it manually, you can run "
f"`huggingface-cli login`, followed by loading the model using `model = CrossEncoder({final_output_dir!r})` "
f"and saving it using `model.push_to_hub('{run_name}')`."
)
if __name__ == "__main__":
main()
模型系列与链接
标准模型:
- mmBERT-small(总计 140M 参数,42M 非嵌入参数)
- mmBERT-base(总计 307M 参数,110M 非嵌入参数)
研究资源:
- 🤗 mmBERT 模型合集
- 📝 论文
- 🗂️ 训练数据(3T+ token,完全开放)
- 💻 GitHub 仓库
- 📊 训练检查点,用于研究训练或继续预训练
来源:Hugging Face:Blog · huggingface.co