文本匹配开山之作--双塔模型及实战

2021-08-20 17:21:45 浏览数 (1)

作者 | 夜小白 整理 | NewBeeNLP

在前面一篇文章中,总结了Representation-Based文本匹配模型的改进方法,

  • 基于表征(Representation)的文本匹配、信息检索、向量召回的方法总结

其中在一篇论文中提到了使用Pre-train方式来提高效果,论文连接如下:

  • Pre-training Tasks for Embedding-based Large-scale Retrieval[1]
  • 谷歌 ICLR 2020 | 向量化召回也需要『预训练』

论文中提到的预训练数据均为,relevant positive Query-Doc 对:

mathcal{T}=left{left(boldsymbol{q}_{i}, boldsymbol{u}_{i}right)right}_{i=1}^{|mathcal{T}|}

训练的目标为最大化当前Postive Query-Doc的Softmax条件概率:

p_{theta}(boldsymbol{d} mid boldsymbol{q})=frac{exp left(f_{theta}(boldsymbol{q}, boldsymbol{d})right)}{sum_{boldsymbol{d}^{prime} in mathcal{D}} exp left(f_{theta}left(boldsymbol{q}_, boldsymbol{d}^{prime}right) bar{chi}right.},

论文中提到,softxmax分母中的

mathcal D

为所有可能的文档集合,这样的话候选文档集合非常大,所以论文中做了近似,「训练时使用当前batch中文档这个子集来代替全集」

mathcal D

,这种方法称为Sample Softmax

TensorFlow中也有这个方法的API实现,但是我一直不是很能理解代码中到底应该怎么实现,突然这几天读到了文本匹配的开山之作 「DSSM」,我发现「DSSM」的训练方法与上面那篇论文非常类似,于是研究了一下源码,有一种豁然开朗的感觉,所以想分享一下,我对这种训练方式的理解。DSSM论链接如下:

  • Learning deep structured semantic models for web search using clickthrough data.[2]

DSSM论文中的训练数据也是Query-Document对,训练目标也为最大化给定Query下点击Doc的条件概率,公式如下,和上面说的Pre-train任务基本一致:

P(D mid Q)=frac{exp (gamma R(Q, D))}{sum_{D prime in D} exp left(gamma Rleft(Q, D^{prime}right)right)}

极大似然估计的公式基本一样,训练都是Point-wise loss,具体各个符号我在下面仔细介绍。

DSSM框架简要介绍

作为文本匹配方向的开山之作,已经有非常多的博客介绍了这个模型,这里我就简单介绍一下,重点放在后面训练源码的阅读。

模型结构

DSSM也是Representation-Based模型,其中Query端 Encoder 和 Doc端 Encoder都是使用 MLP实现,最后Score计算使用的是cosine similarity,后续模型的改进很多都是使用更好的Encoder结构。

输入

DSSM中输入并不是单纯直接使用 bag-of-word,从上面结构图可以看出,输入的时候做了Word Hashing,在进行bag-of-word映射,目的主要如下:

  • 减少词典的大小,直接使用原始word词典非常大(500K),导致输入向量的维数也非常高,使用Word Hashing做分解后,可以减少词典大小,比如letter-trigram(30K)
  • 一定程度解决OOV问题
  • 对拼写错误也有帮助

Word Hashing的做法类似于fast-text中的子词分解,但是不同点在于

  • fast-text中会取多个不同大小窗口对一个单词进行分解,比如2、3、4、5,词表是这些所有的子词构成的集合
  • Word Hashing只会取一个固定大小窗口对单词进行分解,词表是这个固定大小窗口子词的集合,比如letter-bigram,letter-trigram

比如输入的词为#good#,我们选「tri-gram」,则Word-hashing分解后,#good#的表示则为#go,goo,ood,od#,然后就是输入的每个词都映射为tri-gram bag-of-words 向量,出现了的位置为1,否则为0。假设数据集进行「tri-gram」分解后,构成的词表大小为N,那么Query输入处理方式如下:

  • 首先将每个词进行Word Hashing分解
  • 获得每个词的表示,比如 [0,1,1,0,0,0...,0,1] ,维数为N,其中在词表中出现了的位置为1,否则为0
  • 将Query中所有的词的表示向量相加可以得到一个N维向量,「其实就是bag-of-word表示」(只考虑有没有出现,并不考虑出现顺序位置)

Doc端输入的处理也类似于上面Query端的处理,获得Word-Hashing后的向量表示,作为整个模型的输入。

Encoder层

Query端和Doc端Encoder层处理很简单,就是MLP,计算公式如下:

begin{gathered} l_{1}=W_{1} x \ l_{i}=fleft(W_{i} l_{i-1} b_{i}right), i=2, ldots, N-1 \ y=fleft(W_{N} l_{N-1} b_{N}right) end{gathered}

可以看出就是标准的全连接层运算

相似度Score计算

DSSM中最后的相似度计算用的是 cosine similarity,计算公式如下:

R(Q, D)=operatorname{cosine}left(y_{Q}, y_{D}right)=frac{y_{Q}^{T} y_{D}}{left|y_{Q}right|left|y_{D}right|}

模型训练好之后,给定一个Query我们就可以对其所有Doc按照这个计算出来的cosine similarity进行排序。

训练方式

训练数据

DSSM的训练方式是做Point-wise训练,论文中对于训练数据的描述如下:

The clickthrough logs consist of a list of queries and their clicked documents.

给定的是Query以及对应的点击Document,我们需要进行极大似然估计。

训练目标

DSSM首先通过获得的semantic relevance score计算在给定Query下Doc的后验概率:

P(D mid Q)=frac{exp (gamma R(Q, D))}{sum_{D I in D} exp left(gamma Rleft(Q, D^{prime}right)right)}

其中

gamma

为softmax函数的平滑因子,

D

表示所有的待排序的候选文档集合,可以看出这个目标其实和我们一开始提到的Pre-train那篇论文的目标是一样的。我们的候选文档大小可能会非常大,论文在实际训练中,做法如下:

  • 我们使用
(Q,D^ )

来表示一个(Query,Doc)对,其中

D^

表示这个Doc是被点击过的

  • 使用
D^

和四个随机选取没有被点击过的Doc来近似全部文档集合

D

,其中

{D^-_j;j=1,...,4}

表示负样本

上面就是训练时候的实际做法,对于每个

(Q,D^ )

,我们只需要采样K个负样本(K可以自己定),

(Q,D^-_j)

,这样softxmax操作我们也只需要在

hat D ={D^ ,D^-_1,....D^-_k}

这个集合上计算即可,论文中还提到,采样负样本方式对最终结果没有太大影响

In our pilot study, we do not observe any significant difference when different sampling strategies were used to select the unclicked documents.

最后loss选用的就是交叉熵损失:

L(Lambda)=-log prod_{left(Q, D^{ }right)} Pleft(D^{ } mid Qright)
训练方式总结

通过上面的分析,我的理解是DSSM和之前说的Pre-trian那篇论文,训练的时候只需要采样负样本即可,然后softmax操作只在 当前正样本 采样的负样本 集合上计算,最后用交叉熵损失即可。具体负样本怎么采样,我觉的有两种方法:

  • 输入数据中就已经采样好负样本,输入数据直接是正样本 负样本,这样运算量会大些
  • 输入数据batch均为正样本,负样本通过batch中其他Doc构造

DSSM源码阅读

我看的DSSM实现代码是下面两个,其中的不同点就在于上面说的负样本构造不同

  • 训练数据中输入有负样本:InsaneLife/dssm[3]
  • 使用一个batch中其他Doc构造负样本:LiangHao151941/dssm[4]

训练数据中输入有负样本的情况

  • 这部分代码在https://github.com/InsaneLife/dssm/blob/master/dssm_rnn.py
输入数据
代码语言:javascript复制
with tf.name_scope('input'):
    # 预测时只用输入query即可,将其embedding为向量。
    query_batch = tf.placeholder(tf.int32, shape=[None, None], name='query_batch')
    doc_pos_batch = tf.placeholder(tf.int32, shape=[None, None], name='doc_positive_batch')
    doc_neg_batch = tf.placeholder(tf.int32, shape=[None, None], name='doc_negative_batch')
    query_seq_length = tf.placeholder(tf.int32, shape=[None], name='query_sequence_length')
    pos_seq_length = tf.placeholder(tf.int32, shape=[None], name='pos_seq_length')
    neg_seq_length = tf.placeholder(tf.int32, shape=[None], name='neg_sequence_length')
    on_train = tf.placeholder(tf.bool)
    drop_out_prob = tf.placeholder(tf.float32, name='drop_out_prob')
  • doc_pos_batch , 即是论文中说的 $D^ # ,正样本输入
  • doc_neg_batch,即是论文汇总说的
{D^-_j;j=1,...,K}

,负样本输入集合

代码语言:javascript复制
def pull_batch(data_map, batch_id):
    query_in = data_map['query'][batch_id * query_BS:(batch_id   1) * query_BS]
    query_len = data_map['query_len'][batch_id * query_BS:(batch_id   1) * query_BS]
    doc_positive_in = data_map['doc_pos'][batch_id * query_BS:(batch_id   1) * query_BS]
    doc_positive_len = data_map['doc_pos_len'][batch_id * query_BS:(batch_id   1) * query_BS]
    doc_negative_in = data_map['doc_neg'][batch_id * query_BS * NEG:(batch_id   1) * query_BS * NEG]
    doc_negative_len = data_map['doc_neg_len'][batch_id * query_BS * NEG:(batch_id   1) * query_BS * NEG]

    # query_in, doc_positive_in, doc_negative_in = pull_all(query_in, doc_positive_in, doc_negative_in)
    return query_in, doc_positive_in, doc_negative_in, query_len, doc_positive_len, doc_negative_len

这是准备每个batch数据的代码,其中query_BS为batch_size,NEG为负样本采样个数。

合并正负样本与计算余弦相似度

从论文中可以知道,我们需要对「每个Query」选取

hat D ={D^ ,D^-_1,....D^-_k}

这个集合做softmax操作,所以我们计算出每个Query正负样本的Score之后,需要将同一个Query正负样本其合并到一起,Score即为softmax输入的logits。「由于输入数据中直接有负样本」,所以这里不需要我们构造负样本,直接把负样本输出的Score concat即可。下面代码步骤如下:

  • 先把同一个Query下pos_doc和neg_doc经过Encoder之后的隐层表示concat到一起
  • 计算每个Query与正负样本的similarity

计算出来的cosine similarity Tensor如下,每一行是一个Query下正样本和负样本的sim,这样我们在axis = 1上做softmax操作即可:

代码语言:javascript复制
[[query[1]_pos,query[1]_neg[1],query[1]_neg[2],query[1]_neg[3],...],
 [query[2]_pos,query[2]_neg[1],query[2]_neg[2],query[2]_neg[3],...],
 ......,
 [query[n]_pos,query[n]_neg[1],query[n]_neg[2],query[n]_neg[3],...],
]
代码语言:javascript复制
with tf.name_scope('Merge_Negative_Doc'):
    # 合并负样本,tile可选择是否扩展负样本。
    # doc_y = tf.tile(doc_positive_y, [1, 1])
    # 此时doc_y为单独的pos_doc的hidden representation
    doc_y = tf.tile(doc_pos_rnn_output, [1, 1])

    #下面这段代码就是把同一个Query下的neg_doc合并到pos_doc,后续才能计算score 和 softmax
    for i in range(NEG):
        for j in range(query_BS):
            # slice(input_, begin, size)切片API
            # doc_y = tf.concat([doc_y, tf.slice(doc_negative_y, [j * NEG   i, 0], [1, -1])], 0)
            doc_y = tf.concat([doc_y, tf.slice(doc_neg_rnn_output, [j * NEG   i, 0], [1, -1])], 0)

with tf.name_scope('Cosine_Similarity'):
    # Cosine similarity
    # query_norm = sqrt(sum(each x^2))
    query_norm = tf.tile(tf.sqrt(tf.reduce_sum(tf.square(query_rnn_output), 1, True)), [NEG   1, 1])
    # doc_norm = sqrt(sum(each x^2))
    doc_norm = tf.sqrt(tf.reduce_sum(tf.square(doc_y), 1, True))

    prod = tf.reduce_sum(tf.multiply(tf.tile(query_rnn_output, [NEG   1, 1]), doc_y), 1, True)
    norm_prod = tf.multiply(query_norm, doc_norm)

    # cos_sim_raw = query * doc / (||query|| * ||doc||)
    cos_sim_raw = tf.truediv(prod, norm_prod)
    # gamma = 20
    cos_sim = tf.transpose(tf.reshape(tf.transpose(cos_sim_raw), [NEG   1, query_BS])) * 20
    # cos_sim 作为softmax logits输入
softmax操作与计算交叉熵损失

上一步中已经计算出各个Query对其正负样本的cosine similarity,这个将作为softmax输入的logits,然后计算交叉熵损失即可,「因为只有一个正样本,而且其位置在第一个」,所以我们的标签one-hot编码为:

  • [1,0,0,0,0,0,....,0]

所以我们计算交叉熵损失的时候,「只需要取第一列的概率值即可」

代码语言:javascript复制
with tf.name_scope('Loss'):
    # Train Loss
    # 转化为softmax概率矩阵。
    prob = tf.nn.softmax(cos_sim)
    # 只取第一列,即正样本列概率。相当于one-hot标签为[1,0,0,0,.....,0]
    hit_prob = tf.slice(prob, [0, 0], [-1, 1])
    loss = -tf.reduce_sum(tf.log(hit_prob))
    tf.summary.scalar('loss', loss)

使用一个batch中其他Doc构造负样本

上面的方法是在输入数据中直接有负样本,这样计算的时候需要多计算负样本的representation,在输入数据batch中可以只包含正样本,然后再选择同一个batch中的其他Doc构造负样本,这样可以减少计算量

  • 这部分代码在https://github.com/LiangHao151941/dssm/blob/master/single/dssm_v3.py
输入数据
代码语言:javascript复制
with tf.name_scope('input'):
    # Shape [BS, TRIGRAM_D].
    query_batch = tf.sparse_placeholder(tf.float32, shape=query_in_shape, name='QueryBatch')
    # Shape [BS, TRIGRAM_D]
    doc_batch = tf.sparse_placeholder(tf.float32, shape=doc_in_shape, name='DocBatch')

可以看出这里的输入数据只有

(Q,D^ )

,并没有负样本

构造负样本并计算余弦相似度

由于输入数据中没有负样本,所以使用同一个batch中的其他Doc做为负样本,由于所有输入Doc representation在前面已经计算出来了,所以不需要额外再算一遍了,下面的代码就是通过rotate 输入

(Q,D^ )

,来构造负样本,比如:

  • 输入为
{(Q_1,D^ _1),(Q_2,D^ _2),(Q_3,D^ _3)}

,对于每一个

Q_i

,除了

D^ _i

,这个batch中的其他Doc均为负样本

  • 那么对于
Q_1

D^ _2、D^ _3

均为视为

D^-_1

,可以构造负样本为

{(Q_1,D^ _2、D^ _3)}
代码语言:javascript复制
with tf.name_scope('FD_rotate'):
    # Rotate FD  to produce 50 FD-
    temp = tf.tile(doc_y, [1, 1])

    for i in range(NEG):
        rand = int((random.random()   i) * BS / NEG)
        doc_y = tf.concat(0,
                          [doc_y,
                           tf.slice(temp, [rand, 0], [BS - rand, -1]),
                           tf.slice(temp, [0, 0], [rand, -1])])
with tf.name_scope('Cosine_Similarity'):
    # Cosine similarity
    query_norm = tf.tile(tf.sqrt(tf.reduce_sum(tf.square(query_y), 1, True)), [NEG   1, 1])
    doc_norm = tf.sqrt(tf.reduce_sum(tf.square(doc_y), 1, True))

    prod = tf.reduce_sum(tf.mul(tf.tile(query_y, [NEG   1, 1]), doc_y), 1, True)
    norm_prod = tf.mul(query_norm, doc_norm)

    cos_sim_raw = tf.truediv(prod, norm_prod)
    cos_sim = tf.transpose(tf.reshape(tf.transpose(cos_sim_raw), [NEG   1, BS])) * 20
softmax操作与计算交叉熵损失

这一步和前面说的是一样的

代码语言:javascript复制
with tf.name_scope('Loss'):
    # Train Loss
    prob = tf.nn.softmax((cos_sim))
    hit_prob = tf.slice(prob, [0, 0], [-1, 1])
    loss = -tf.reduce_sum(tf.log(hit_prob)) / BS
    tf.scalar_summary('loss', loss)

总结

之前一直对于sampled softmax不太理解,不知道在实际训练中如何做。但是看了DSSM论文和源码之后,真的有一种拨开云雾见月明的感觉,这种训练方式的核心就在于「构造负样本」,这样一说感觉和Pairwise loss中构造pair又有点类似,不过这里构造的不止一个负样本,训练目标也是pointwise,这种方式应该是不需要用到TensorFlow中的tf.nn.sampled_softmax_loss这个函数。

当然上面都是个人理解,最近越来越觉得真正要弄懂一个算法不单要理解数学原理,而且需要去读懂源码,很多在论文中理解不了的信息,在源码中都会清晰的展现出来,这部分我也一直在探索中,之后有什么心得再分享给大家啦~

本文参考资料

[1]

Pre-training Tasks for Embedding-based Large-scale Retrieval: https://arxiv.org/pdf/2002.03932.pdf

[2]

Learning deep structured semantic models for web search using clickthrough data.: https://dl.acm.org/doi/10.1145/2505515.2505665

[3]

InsaneLife/dssm: https://github.com/InsaneLife/dssm

[4]

LiangHao151941/dssm: https://github.com/LiangHao151941/dssm

0 人点赞