关于预训练语言模型的句子嵌入

这是对以下论文的 TensorFlow 实现：

关于预训练语言模型的句子嵌入
李博涵、周浩、何俊贤、王明轩、杨一鸣、李磊
EMNLP 2020

模型	斯皮尔曼等级相关系数
BERT-large-NLI	77.80
BERT-large-NLI-last2avg	78.45
BERT-large-NLI-flow（仅目标训练）	80.54
BERT-large-NLI-flow（目标及训练/验证/测试集共同训练）	81.18

如有任何问题，请联系 bohanl1@cs.cmu.edu。

需求

• Python >= 3.6
• TensorFlow >= 1.14

准备工作

预训练 BERT 模型

export BERT_PREMODELS="../bert_premodels"
mkdir ${BERT_PREMODELS}; cd ${BERT_PREMODELS}

# 然后从 https://github.com/google-research/bert 下载预训练的 BERT 模型
curl -O https://storage.googleapis.com/bert_models/2018_10_18/uncased_L-12_H-768_A-12.zip
curl -O https://storage.googleapis.com/bert_models/2018_10_18/uncased_L-24_H-1024_A-16.zip

ls ${BERT_PREMODELS}/uncased_L-12_H-768_A-12 # base
ls ${BERT_PREMODELS}/uncased_L-24_H-1024_A-16 # large

GLUE

export GLUE_DIR="../glue_data"
python download_glue_data.py --data_dir=${GLUE_DIR}

# 然后下载 STS-B 的标注测试集
cd ../glue_data/STS-B
curl -O https://raw.githubusercontent.com/kawine/usif/master/STSBenchmark/sts-test.csv

SentEval

cd ..
git clone https://github.com/facebookresearch/SentEval

使用方法

在 NLI 监督下微调 BERT（可选）

export OUTPUT_PARENT_DIR="../exp"
export CACHED_DIR=${OUTPUT_PARENT_DIR}/cached_data
mkdir ${CACHED_DIR}

export RANDOM_SEED=1234
export CUDA_VISIBLE_DEVICES=0
export BERT_NAME="large"
export TASK_NAME="ALLNLI"
unset INIT_CKPT
bash scripts/train_siamese.sh train \
"--exp_name=exp_${BERT_NAME}_${RANDOM_SEED} \
--num_train_epochs=1.0 \
--learning_rate=2e-5 \
--train_batch_size=16 \
--cached_dir=${CACHED_DIR}"


# 评估
export RANDOM_SEED=1234
export CUDA_VISIBLE_DEVICES=0
export TASK_NAME=STS-B
export BERT_NAME=large
export OUTPUT_PARENT_DIR="../exp"
export INIT_CKPT=${OUTPUT_PARENT_DIR}/exp_${BERT_NAME}_${RANDOM_SEED}/model.ckpt-60108
export CACHED_DIR=${OUTPUT_PARENT_DIR}/cached_data
export EXP_NAME=exp_${BERT_NAME}_${RANDOM_SEED}_eval
bash scripts/train_siamese.sh predict \
"--exp_name=${EXP_NAME} \
 --cached_dir=${CACHED_DIR} \
 --sentence_embedding_type=avg \
 --flow=0 --flow_loss=0 \
 --num_examples=0 \
 --num_train_epochs=1e-10"

注意：您可能希望添加 --use_xla 以加快 BERT 微调的速度。

无监督学习基于流的生成模型

export CUDA_VISIBLE_DEVICES=0
export TASK_NAME=STS-B
export BERT_NAME=large
export OUTPUT_PARENT_DIR="../exp"
export INIT_CKPT=${OUTPUT_PARENT_DIR}/exp_large_1234/model.ckpt-60108
export CACHED_DIR=${OUTPUT_PARENT_DIR}/cached_data
bash scripts/train_siamese.sh train \
"--exp_name_prefix=exp \
 --cached_dir=${CACHED_DIR} \
 --sentence_embedding_type=avg-last-2 \
 --flow=1 --flow_loss=1 \
 --num_examples=0 \
 --num_train_epochs=1.0 \
 --flow_learning_rate=1e-3 \
 --use_full_for_training=1"

# 评估
export CUDA_VISIBLE_DEVICES=0
export TASK_NAME=STS-B
export BERT_NAME=large
export OUTPUT_PARENT_DIR="../exp"
export INIT_CKPT=${OUTPUT_PARENT_DIR}/exp_large_1234/model.ckpt-60108
export CACHED_DIR=${OUTPUT_PARENT_DIR}/cached_data
export EXP_NAME=exp_t_STS-B_ep_1.00_lr_5.00e-05_e_avg-last-2_f_11_1.00e-03_allsplits
bash scripts/train_siamese.sh predict \
"--exp_name=${EXP_NAME} \
 --cached_dir=${CACHED_DIR} \
 --sentence_embedding_type=avg-last-2 \
 --flow=1 --flow_loss=1 \
 --num_examples=0 \
 --num_train_epochs=1.0 \
 --flow_learning_rate=1e-3 \
 --use_full_for_training=1"

仅使用 STS-B 的训练集拟合流模型

export CUDA_VISIBLE_DEVICES=0
export TASK_NAME=STS-B
export BERT_NAME=large
export OUTPUT_PARENT_DIR="../exp"
export INIT_CKPT=${OUTPUT_PARENT_DIR}/exp_large_1234/model.ckpt-60108
export CACHED_DIR=${OUTPUT_PARENT_DIR}/cached_data
bash scripts/train_siamese.sh train \
"--exp_name_prefix=exp \
 --cached_dir=${CACHED_DIR} \
 --sentence_embedding_type=avg-last-2 \
 --flow=1 --flow_loss=1 \
 --num_examples=0 \
 --num_train_epochs=1.0 \
 --flow_learning_rate=1e-3 \
 --use_full_for_training=0"

# 评估
export CUDA_VISIBLE_DEVICES=0
export TASK_NAME=STS-B
export BERT_NAME=large
export OUTPUT_PARENT_DIR="../exp"
export INIT_CKPT=${OUTPUT_PARENT_DIR}/exp_large_1234/model.ckpt-60108
export CACHED_DIR=${OUTPUT_PARENT_DIR}/cached_data
export EXP_NAME=exp_t_STS-B_ep_1.00_lr_5.00e-05_e_avg-last-2_f_11_1.00e-03
bash scripts/train_siamese.sh predict \
"--exp_name=${EXP_NAME} \
 --cached_dir=${CACHED_DIR} \
 --sentence_embedding_type=avg-last-2 \
 --flow=1 --flow_loss=1 \
 --num_examples=0 \
 --num_train_epochs=1.0 \
 --flow_learning rate=1e-3 \
 --use_full for training=1"

下载我们的模型

我们的模型可在以下链接获取：https://drive.google.com/file/d/1-vO47t5SPFfzZPKkkhSe4tXhn8u--KLR/view?usp=sharing

参考文献

@inproceedings{li2020emnlp,
    title = {关于预训练语言模型的句子嵌入},
    author = {李博涵、周浩、何俊贤、王明轩、杨一鸣、李磊},
    booktitle = {自然语言处理中的经验方法会议（EMNLP）},
    month = {十一月},
    year = {2020}
}

致谢

本仓库的大部分内容借鉴自以下项目：

• https://github.com/google-research/bert
• https://github.com/zihangdai/xlnet
• https://github.com/tensorflow/tensor2tensor

BERT-flow 快速上手指南

BERT-flow 是一种基于预训练语言模型（如 BERT）生成高质量句子嵌入的方法。它通过引入流式生成模型（Flow-based Generative Models），将 BERT 输出的各向异性分布映射为各向同性的高斯分布，从而显著提升句子语义相似度任务的表现。

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统: Linux (推荐) 或 macOS
• Python: >= 3.6
• 深度学习框架: TensorFlow >= 1.14 (注意：本项目基于 TF 1.x，非 TF 2.x)
• GPU: 可选，但强烈建议使用 CUDA 加速训练过程

安装与数据准备

1. 克隆项目与依赖

首先获取代码库，并安装必要的 Python 依赖（假设项目根目录下有 requirements.txt，若无则需手动安装 tensorflow 等）。

git clone https://github.com/bohanli/BERT-flow.git
cd BERT-flow
# 根据实际需求安装依赖，例如：
pip install tensorflow-gpu==1.15.0

2. 下载预训练 BERT 模型

从 Google 官方源下载预训练模型。国内用户若下载缓慢，可尝试使用国内镜像站或手动下载后上传至服务器。

export BERT_PREMODELS="../bert_premodels"
mkdir -p ${BERT_PREMODELS}
cd ${BERT_PREMODELS}

# 下载 Base 和 Large 模型
curl -O https://storage.googleapis.com/bert_models/2018_10_18/uncased_L-12_H-768_A-12.zip
curl -O https://storage.googleapis.com/bert_models/2018_10_18/uncased_L-24_H-1024_A-16.zip

# 解压
unzip uncased_L-12_H-768_A-12.zip
unzip uncased_L-24_H-1024_A-16.zip

3. 准备评测数据集 (GLUE & STS-B)

下载 GLUE 基准数据集及 STS-B 的测试集标签。

cd ..
export GLUE_DIR="../glue_data"
python download_glue_data.py --data_dir=${GLUE_DIR}

# 下载 STS-B 标注测试集
cd ../glue_data/STS-B
curl -O https://raw.githubusercontent.com/kawine/usif/master/STSBenchmark/sts-test.csv

4. 准备 SentEval 工具

用于后续的句子嵌入质量评估。

cd ../../..
git clone https://github.com/facebookresearch/SentEval

基本使用

以下是使用 BERT-flow 的核心流程：先微调 BERT（可选），然后训练流模型，最后进行评估。

步骤 1: 微调 BERT (可选)

如果您希望先在 NLI 数据上微调 BERT 以获得更好的初始表示，可运行以下命令。若直接使用原始 BERT，可跳过此步并使用官方提供的预训练检查点。

export OUTPUT_PARENT_DIR="../exp"
export CACHED_DIR=${OUTPUT_PARENT_DIR}/cached_data
mkdir -p ${CACHED_DIR}

export RANDOM_SEED=1234
export CUDA_VISIBLE_DEVICES=0
export BERT_NAME="large"
export TASK_NAME="ALLNLI"
unset INIT_CKPT

bash scripts/train_siamese.sh train \
"--exp_name=exp_${BERT_NAME}_${RANDOM_SEED} \
--num_train_epochs=1.0 \
--learning_rate=2e-5 \
--train_batch_size=16 \
--cached_dir=${CACHED_DIR}"

提示: 添加 --use_xla 参数可以加速 BERT 的微调过程。

步骤 2: 训练流模型 (核心步骤)

这是 BERT-flow 的关键步骤，它将 BERT 的输出分布转换为高斯分布。以下示例使用 STS-B 数据进行无监督学习。

export CUDA_VISIBLE_DEVICES=0
export TASK_NAME=STS-B
export BERT_NAME=large
export OUTPUT_PARENT_DIR="../exp"
# 指向步骤 1 生成的模型，或直接使用下载的预训练 BERT 模型路径
export INIT_CKPT=${OUTPUT_PARENT_DIR}/exp_large_1234/model.ckpt-60108
export CACHED_DIR=${OUTPUT_PARENT_DIR}/cached_data

bash scripts/train_siamese.sh train \
"--exp_name_prefix=exp \
 --cached_dir=${CACHED_DIR} \
 --sentence_embedding_type=avg-last-2 \
 --flow=1 --flow_loss=1 \
 --num_examples=0 \
 --num_train_epochs=1.0 \
 --flow_learning_rate=1e-3 \
 --use_full_for_training=1"

关键参数说明：

• --flow=1 --flow_loss=1: 启用流模型训练。
• --sentence_embedding_type=avg-last-2: 推荐使用最后两层的平均池化。
• --use_full_for_training=1: 使用所有可用数据（train+dev+test）进行流模型拟合，通常能获得最佳效果（见论文结果 81.18）。若仅需训练集，设为 0。

步骤 3: 评估模型

训练完成后，使用以下命令在 STS-B 任务上评估性能（计算 Spearman 相关系数）。

export CUDA_VISIBLE_DEVICES=0
export TASK_NAME=STS-B
export BERT_NAME=large
export OUTPUT_PARENT_DIR="../exp"
export INIT_CKPT=${OUTPUT_PARENT_DIR}/exp_large_1234/model.ckpt-60108
export CACHED_DIR=${OUTPUT_PARENT_DIR}/cached_data
# 替换为您实际训练生成的实验名称
export EXP_NAME=exp_t_STS-B_ep_1.00_lr_5.00e-05_e_avg-last-2_f_11_1.00e-03_allsplits

bash scripts/train_siamese.sh predict \
"--exp_name=${EXP_NAME} \
 --cached_dir=${CACHED_DIR} \
 --sentence_embedding_type=avg-last-2 \
 --flow=1 --flow_loss=1 \
 --num_examples=0 \
 --num_train_epochs=1.0 \
 --flow_learning_rate=1e-3 \
 --use_full_for_training=1"

直接使用预训练模型

如果您不想重新训练，可以直接下载作者提供的已训练好的模型进行检查和测试：

• 下载地址：Google Drive Link

下载后，将 INIT_CKPT 环境变量指向解压后的模型文件路径，然后直接运行步骤 3中的评估命令即可。