BERT-flow
用 Normalizing Flow 把 BERT 各向异性嵌入空间映射到标准高斯分布,提升句子相似
加载项目详情…
本应用为开源项目,仅供学习研究,请遵守其开源协议。
用 Normalizing Flow 把 BERT 各向异性嵌入空间映射到标准高斯分布,提升句子相似
加载项目详情…
本应用为开源项目,仅供学习研究,请遵守其开源协议。
当你问 ChatGPT"苹果和橙子有什么相同点",它能流畅作答。但你有没有想过,这些答案背后的"语义理解"到底是怎么量化的?在 NLP 研究中,把一句话转换成一段数字向量(Sentence Embedding),是文本分类、相似度计算、语义搜索等一切下游任务的地基。BERT 的出现让语言理解大幅提升,但其原始输出的 embedding 空间存在严重的各向异性(anisotropy)——词向量挤在了一个狭窄的锥形区域内,导致直接用余弦相似度衡量句子间的语义距离并不准确。
这就是 BERT-flow 要解决的核心问题。
BERT-flow 由卡内基梅隆大学(CMU)和字节跳动的研究者共同提出,论文《On the Sentence Embeddings from Pre-trained Language Models》发表于自然语言处理顶级会议 EMNLP 2020(Google Scholar 显示引用量已超过 2000 次)。第一作者 Bohan Li 当时在 CMU 读博,导师为 Yiming Yang 教授,合作者还包括周浩(字节跳动)、何俊贤(港中文)、王明轩(字节跳动)、李磊(字节跳动)等 NLP 领域活跃研究者。该团队同年还在 AAAI、ACL 等顶会发表了多篇语义表示方向的论文,属于 sentence embedding 这个细分领域的早期开拓者之一。
该问题的根源在于:BERT 在预训练时使用的是 LM Loss(语言建模),它优化的是"下一个词的预测准确率",并不直接优化"语义空间的几何结构"。因此 BERT 的词嵌入呈现出词频偏差(高频词分布紧凑、低频词分布稀疏)和语义空间各向异性的特征。Flow 模型最早由 OpenAI 在 Glow 论文中引入图像生成领域,BERT-flow 的核心创新是将 normalizing flow 引入 NLP 的句子嵌入优化,用 flow 把 BERT 的各向异性分布映射到标准高斯分布,使语义相似度计算变得可靠。
如果你把 sentence embedding 想象成把一句话"翻译"成地图上的一个坐标点,那么 BERT-flow 做的事情就是:首先用 BERT 把句子编码成一个向量(相当于用投影仪把三维物体投射到二维平面),然后用 normalizing flow 这个"坐标矫正器",把这个二维平面拉伸、旋转、平移,让所有点均匀分布在一个标准圆内。这样一来,两个点的距离(欧氏距离或余弦相似度)就能真实反映它们之间的语义距离了。
Normalizing flow 的核心思想是:通过一系列可逆变换(bijective transformations),把一个复杂分布逐步转换为简单分布。BERT-flow 采用了 1x1 可逆卷积(invertible 1x1 convolution)的 Glow 结构来实现这一目标。Glow 模块由 actnorm(激活归一化)+ affine coupling layer + 1x1 conv 三部分堆叠而成,每一层都可以精确计算对数行列式(log det),从而用最大似然估计进行端到端训练。相比 VAE 和 GAN,flow 模型的优势是精确的对数似然计算,训练目标直接且稳定。
该仓库是论文的 TensorFlow 1.14+ 实现,代码结构清晰,主要包含以下模块:
modeling.py — BERT 模型的核心实现,直接 fork 自 Google 官方 BERT 仓库,包含 BertConfig、BertModel、BertModelPreTrainingModel 等类,定义了 Transformer 编码器和嵌入层。代码从 Google Research 仓库沿袭了大量注释和文档风格,是标准的预训练语言模型实现。
run_siamese.py — Siamese Network(孪生网络)训练入口。对应论文中的两个训练阶段:① NLI 监督微调:用自然语言推理数据集(NLI)微调 BERT;② Flow 训练:冻结 BERT 参数,只训练 flow 模块,用 STS-B 数据集最大化 BERT 嵌入到标准高斯的 log-likelihood。
siamese_utils.py — 数据处理工具集,包含 StsbProcessor、MnliProcessor、SnliTrainProcessor 等多个 GLUE 任务处理器,负责将原始数据转换为 BERT 所需的 [CLS] text_a [SEP] text_b [SEP] 格式,并处理标签映射和特征序列化。
flow/glow_ops_1x1.py — Flow 模型的核心算子实现。从 tensor2tensor 仓库移植,改写了以下内容:① 用原生 TensorFlow 替代 tensorflow_probability(因为 TF 1.14 不兼容 tfp 0.6/0.7);② 移除了视频相关的 conv3d 算子;③ 移除了条件分布支持。这是一个对旧版本 TensorFlow 的工程适配,代码中保留了大量原始 tensor2tensor 的结构。
flow/glow_1x1.py — Glow 模块的顶层封装,定义了 Glow 类和 AttrDict 配置结构。Glow 由多层 flow 组成,每层包含 actnorm → affine coupling → 1x1 invertible conv 三步。
optimization.py / optimization_bert_flow.py — 优化器实现,包含 AdamWeightDecay 等自定义学习率调度器。
download_glue_data.py — 自动下载 GLUE 基准数据集的工具脚本,支持 CoLA、SST-2、MRPC、QQP、STS-B、MNLI、SNLI、QNLI、RTE、WNLI 等多个任务。
scripts/train_siamese.sh — 封装了训练和推理的 Bash 脚本,通过环境变量配置 BERT 模型路径、数据路径、实验名称等。提供了三种使用模式:train(微调 BERT 或训练 flow)、predict(推理 + 调用 eval_stsb.py 计算 Spearman 相关系数)。
实验流程上,团队首先用 NLI 数据集(SNLI + MNLI)微调 BERT,使模型具备基础的语义推理能力;然后固定 BERT 参数,在 STS-B(语义文本相似度基准)训练集上用 MLE 目标训练 normalizing flow,让 BERT 输出的嵌入分布逐步逼近标准高斯。训练完成后,在 STS-B 测试集上评估,flow 训练仅用训练集时达到 80.54 Spearman's rho,训练集+验证集+测试集全部使用时达到 81.18(远超 BERT-base-NLI 的 77.80)。
该仓库面向研究场景,对工程实践者有一定门槛。主要挑战体现在以下几个方面:
依赖复杂:需要 TensorFlow 1.14+(不是 TF2),Python >= 3.6,且必须手动下载 Google 官方 BERT 预训练模型(uncased_L-12_H-768_A-12.zip 和 uncased_L-24_H-1024_A-16.zip,每个约 1.3GB)。此外还要通过 download_glue_data.py 下载 GLUE 数据集,并手动 clone Facebook Research 的 SentEval 评估库。README 提供了完整的准备步骤,但全流程走完至少需要 2-3 小时。
硬件要求高:微调 BERT-large(24层、1024维)需要至少 16GB 显存的 GPU(V100 或以上),flow 模块本身也需要额外显存。仓库中通过 TF_FORCE_GPU_ALLOW_GROWTH=true 启用 GPU 内存动态增长以避免 OOM。纯 CPU 训练几乎不可行。
非生产级代码:这是典型的学术研究复现,代码未经生产环境验证。变量命名、异常处理、单元测试覆盖都较为简陋,glow_ops_1x1.py 中大量使用 tf.contrib(TensorFlow 1.x 的不稳定模块),在 TF 2.x 中完全不可用。此外仓库未提供 Docker / docker-compose 等容器化支持,无法在隔离环境中复现训练环境。
上手路径:如果你只是想体验效果,可以直接下载作者提供的 Google Drive 预训练模型 checkpoint,跳过训练环节。推理流程只需要:准备 GLUE 数据 → 加载 BERT + flow checkpoint → 计算句子向量 → 计算余弦相似度。不过项目未提供独立的推理脚本,需要阅读 run_siamese.py 的 --do_predict 分支来理解推理逻辑。
与后继工作的对比:BERT-flow 发表于 2020 年底,是 sentence embedding 各向异性问题的早期系统性解决方案。2021 年后,SBERT(Sentence-BERT)、SimCSE、E5、BAISE 等后继工作从对比学习、更深层的表示学习等角度进一步提升了语义表示质量,且大多提供了更友好的 PyTorch 实现和预训练模型下载。BERT-flow 的核心贡献在于奠定了"flow 改善 embedding 分布"的理论基础,后继工作大多沿用了这一思想。
尽管 BERT-flow 取得了当时最优的 STS-B 性能,但其局限性也比较明显。首先,flow 模型训练开销较大,需要额外训练一个 flow 模块,而对比学习方法(如 SimCSE)仅需一个无监督目标即可达到相近效果。其次,flow 模型仅优化了嵌入的分布形状(各向同性),并未显式优化语义相似度目标,对某些语义差异较大的句子对区分能力有限。此外,该仓库只支持英文 sentence embedding,未涉及多语言场景,而多语言 BERT(如 mBERT、XLM-R)的大规模多语言 embedding 优化是另一个活跃的研究方向。
BERT-flow 是 NLP sentence embedding 领域的重要里程碑式工作,首次系统性地用 normalizing flow 解决了 BERT 嵌入各向异性的问题,并在 STS-B 基准上取得了 81.18 的 Spearman's rho(2020 年底最优)。虽然其 TensorFlow 1.x 实现对当代研究者的吸引力有限,但论文中提出的"flow + pre-trained LM"的框架启发了大量后续研究。该仓库的代码结构和实验设计较为规范,适合作为学习 sentence embedding 优化技术的参考实现。
GitHub 数据亮点:
快速体验路径: 下载 Google Drive 预训练模型 → 准备 STS-B 数据 → 复用 run_siamese.py --do_predict 分支推理 → 计算句子余弦相似度。整个过程约 1-2 小时可跑通,比训练更实用。