Implementation:FlagOpen FlagEmbedding BGE Finetune Modeling
| Knowledge Sources | |
|---|---|
| Domains | Machine_Learning, Natural_Language_Processing, Information_Retrieval |
| Last Updated | 2026-02-09 00:00 GMT |
Overview
Bi-encoder model architecture for training BGE embedding models with contrastive learning.
Description
The BiEncoderModel implements a dual-encoder architecture for learning dense text embeddings optimized for retrieval tasks. Key features include:
Architecture: Uses AutoModel as the base encoder that processes queries and passages independently. Embeddings are extracted from either the [CLS] token or mean pooling of the last hidden state. Optional L2 normalization produces unit vectors for cosine similarity computation.
Training: Employs in-batch negatives contrastive learning with cross-entropy loss. For each query, the model learns to rank its positive passage higher than all negative passages in the batch. Supports cross-device negative sharing in distributed training to increase the effective batch size. Temperature scaling controls the sharpness of the similarity distribution.
Optimization: Supports gradient checkpointing for memory efficiency and handles both inner product (unnormalized) and cosine similarity (normalized) scoring. When not normalized, temperature is fixed to 1.0 for stability.
Usage
Use this for fine-tuning pre-trained language models on retrieval tasks with contrastive learning, particularly for building general-purpose embedding models.
Code Reference
Source Location
- Repository: FlagOpen_FlagEmbedding
- File: research/baai_general_embedding/finetune/modeling.py
- Lines: 1-141
Signature
class BiEncoderModel(nn.Module):
def __init__(self, model_name: str = None, normlized: bool = False,
sentence_pooling_method: str = 'cls',
negatives_cross_device: bool = False,
temperature: float = 1.0,
use_inbatch_neg: bool = True)
def encode(self, features)
def forward(self, query: Dict[str, Tensor] = None,
passage: Dict[str, Tensor] = None,
teacher_score: Tensor = None)
Import
from research.baai_general_embedding.finetune.modeling import BiEncoderModel
I/O Contract
Inputs
| Name | Type | Required | Description |
|---|---|---|---|
| model_name | str | Yes | HuggingFace model name or path |
| normlized | bool | No | Whether to L2-normalize embeddings for cosine similarity |
| sentence_pooling_method | str | No | Pooling method: 'cls' or 'mean' |
| negatives_cross_device | bool | No | Share negatives across GPUs in distributed training |
| temperature | float | No | Temperature for scaling similarity scores |
| query | Dict[str, Tensor] | Yes | Tokenized query inputs |
| passage | Dict[str, Tensor] | Yes | Tokenized passage inputs (1 pos + N neg per query) |
Outputs
| Name | Type | Description |
|---|---|---|
| loss | Tensor | Contrastive learning loss (training only) |
| scores | Tensor | Similarity scores between queries and passages |
| q_reps | Tensor | Query embeddings |
| p_reps | Tensor | Passage embeddings |
Usage Examples
from transformers import AutoModel
from research.baai_general_embedding.finetune.modeling import BiEncoderModel
# Initialize model
model = BiEncoderModel(
model_name="bert-base-uncased",
normlized=True, # Use cosine similarity
sentence_pooling_method='cls', # Use [CLS] token
negatives_cross_device=True, # Share negatives across GPUs
temperature=0.02, # Temperature scaling
use_inbatch_neg=True # Use in-batch negatives
)
# Training forward pass
query_inputs = {
"input_ids": query_ids, # [batch_size, seq_len]
"attention_mask": query_mask
}
passage_inputs = {
"input_ids": passage_ids, # [batch_size * group_size, seq_len]
"attention_mask": passage_mask
}
outputs = model(query=query_inputs, passage=passage_inputs)
loss = outputs.loss # Contrastive loss
loss.backward()
# Inference
with torch.no_grad():
q_reps = model.encode(query_inputs) # [batch_size, hidden_dim]
p_reps = model.encode(passage_inputs)
scores = torch.matmul(q_reps, p_reps.T) # Similarity scores