BGE多模态嵌入技术解析：从原理到工程实践

1次阅读

没有评论

共计 2000 个字符，预计需要花费 5 分钟才能阅读完成。

现代推荐系统和搜索引擎需要处理图像、文本、视频等多种模态的数据。传统方法如 CLIP（Contrastive Language-Image Pretraining）通过对比学习实现跨模态对齐，但在处理非对称模态关系时存在局限性：

单向注意力机制难以捕捉模态间的双向交互
固定温度系数导致不同模态间的相似度分布不匹配
对长尾数据（低频特征）的嵌入质量不稳定

BGE（Bidirectional Generative Embedding）的核心是双向交叉注意力模块。其工作流程可分为三个阶段：

模态特征提取
图像分支使用 ResNet-50 作为骨干网络
文本分支采用 BERT-base 架构
特征交互层
图像到文本的注意力路径：$Attention(Q_{text},K_{image},V_{image})$
文本到图像的注意力路径：$Attention(Q_{image},K_{text},V_{text})$
融合预测层
通过门控机制动态加权两种注意力输出
最终嵌入向量维度默认为 768

在 MS-COCO 数据集上的测试结果：

方法	R@1	R@5	R@10
CLIP	58.3	82.1	89.7
BGE-base	63.2	85.4	91.8
BGE-large	65.7	87.3	93.2

温度系数 τ ：控制相似度分布的陡峭程度
过大导致相似度区分度下降
过小引发训练不稳定
推荐初始值 0.07，按 0.01 间隔网格搜索
嵌入维度：
维度增加提升表达能力但增大计算开销
生产环境中推荐 256-1024 范围

import torch
import torch.nn as nn
from transformers import BertModel, BertTokenizer

class BGEModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.text_encoder = BertModel.from_pretrained('bert-base-uncased')
        self.image_encoder = ResNet50(pretrained=True)
        self.cross_attn = nn.MultiheadAttention(embed_dim=768, num_heads=8)

    def forward(self, text_input, image_input):
        # 文本特征提取 [batch, seq_len, 768]
        text_feat = self.text_encoder(**text_input).last_hidden_state

        # 图像特征提取 [batch, 2048, 7, 7]
        image_feat = self.image_encoder(image_input)
        image_feat = image_feat.flatten(2).transpose(1,2)  # [batch, 49, 2048]

        # 双向注意力
        text2img, _ = self.cross_attn(query=text_feat.mean(1).unsqueeze(0),
            key=image_feat.transpose(0,1),
            value=image_feat.transpose(0,1)
        )
        img2text, _ = self.cross_attn(query=image_feat.mean(1).unsqueeze(0),
            key=text_feat.transpose(0,1),
            value=text_feat.transpose(0,1)
        )

        return torch.cat([text2img.squeeze(0), img2text.squeeze(0)], dim=-1)

使用 DDP（Distributed Data Parallel）时需注意：