论文阅读“miCSE: Mutual Information C

作者: 掉了西红柿皮_Kee | 来源:发表于2023-07-22 12:53 被阅读0次

Klein T, Nabi M. miCSE: Mutual Information Contrastive Learning for Low-shot Sentence Embeddings[J]. arXiv preprint arXiv:2211.04928, 2022.

摘要导读

本文提出了一个基于互信息的对比学习框架miCSE，显著地提高了在少量句子嵌入方面的先进水平。该方法在对比学习过程中调整了不同视图之间的注意力模式。通过miCSE学习句子嵌入需要加强每个句子的增强视图的结构一致性，使对比自监督学习提高了效率。因此，在小样本领域该方法取得了很好的效果，在全样本的场景下依旧适用。

本文的贡献如下：

通过添加一个attention-level的目标，将结构性信息引入到语言模型中。
引入了注意力互信息（AMI），一种可以提高样本效率的自监督对比学习方法。

方法浅析

该方法旨在在对比学习的方案中利用到句子结构信息。与传统的仅在嵌入空间中的语义相似度水平上进行操作的对比学习相比，该方法在模型中注入了结构信息。这是通过在训练过程中正则化模型的注意力空间来实现的。

符号声明

给定字符串语料库 $\mathcal{X}$ ，其对应的数据集表示为 $\mathcal{D}=\{x_1, x_2, \cdots,x_{|X|}\}$ ， $x_i \in \mathbb{N}^n$ 表示含有 $n$ 个token的序列。在对句子映射时，本文采用的是bi-encoder-- $f_{\theta}$ ，其输入是输入句子的不同类型的增强表示。这里使用 $v \in \{1,2\}$ 作为增强视图表示的索引值。因此，对batch_size为 $\mathcal{D}_b$ 进行编码，会得到嵌入矩阵 $E_v \in \mathbb{R}^{|{\mathbb{D}_b}|\times U}$ ， $U$ 是嵌入表示的维度。使用Transformer的话，对应于 $E_v$ 产生的还有其相关联的注意力矩阵 $W_v$ 。因此，提出的模型联合优化如下损失，以达到对语义和结构的一致性学习：

显然，损失函数中的第一项是语义的对齐，在嵌入表示空间中使用传统的InfoNCE来实现；第二项是在注意力空间中对句法的对齐，不同的是，句法的对齐仅关注正例。

Embedding-level Momentum-Contrastive Learning (InfoNCE)

InfoNCE loss试图在嵌入空间中将正例对拉在一起，同时将负例对分开。具体来说，嵌入的InfoNCE推动每个样本以及相应的增强嵌入表示之间的相似性。对应的损失函数如下：

其中

e_i \in E_1

，

+e_i \in E_2

分别是对应于

x_i

的两种不同的嵌入表示。

d(x, y)=\text{exp}(sim(x, y)/\tau)

，

sim(\cdot)

为余弦相似度。显然，负例的构成包含两种形式：（1）给定batch中除了当前样本之外的样本；（2）存储在

\mathcal{Q}

中的前序batch中的嵌入表示（这是动量编码器中扩充负样本的常用操作）。

Attention-level Mutual Information (AMI)

首先关于Transformer中的注意力机制的详情这里不再进行赘述。只需要知道在每个注意力头中包含三个矩阵 $Q$ , $K$ , $V$ ，其中 $Q$ , $K$ 进行运算会得到注意力权重矩阵 $W=softmax(f(Q,K)) \in \mathbb{R}^{n \times n}$ ，其中 $f(\cdot)$ 表示缩放点积。最终注意力头的输出为 $WV$ 。当然在实际的应用中为了得到不同的嵌入子空间，一般会将该注意力操作重复 $H$ 次，被称为多头注意力机制。在训练编码器的过程中，自注意力张量 $W$ 的值会受到随机确定性过程的影响，这种随机性由dropout操作产生。因此，基于结构信息的对齐，则是要最大化 $W_v=[w_1, w_2, \cdots, w_{|{\mathcal{D}}_b|}]$ 之间的互信息。本文通过四个步骤来正则化注意力空间。

Attention Tensor Slicing
Attention tensor slicing 在输入上实例化一个Transformer栈会产生一个注意张量 $W$ ，其中包含Transformer的层数 $L$ 和自注意力的头数 $H$ 。显然，这里将 $w$ 打平成1维的张量是为了在三维空间中更好的展示tensor的shape。因此，对于 $x_i$ 的某个自注意力头来说，其输出是 $w_i \in \mathbb{R}^{1 \times n \times n}$ 。考虑到多层和多头，其对应的张量就变成了 $W \in \mathbb{R}^{L \times H \times n \times n}$ 。
slicing函数的主要作用是将每个输入样本的注意力张量切分为 $R$ 个元素：即：其中对于 $r \in [1, R]$ ，每个元素 $w_i^r \in \mathbb{R}^{n \times n}=(w_{j,k})_{1<=j,k<=n}$ 。如果token个数小于 $n$ ，则通过填充对不同长度的序列进行补齐，以适应批处理。（这种切分方式的好处是还是保留了句子 $x_i$ 中 $n$ 个token之间的相关关系。只是不知道 $R$ 的设定是否会对性能产生大的影响。）
Attention Sampling
虽然[PAD]的填充会使得在GPU上可以进行有效的批处理编码，但在查看相关关系时，需要放弃token对[PAD]标记产生的注意力得分。为了适应不同长度的标记化序列，对每个网格单元内的注意得分执行采样操作。采样使用的是多项式分布对于有值的 $s$ ,( $1<=s<=n$ )，构造 $s^2$ 的注意力得分池，每个得分被等概率采样，其余的则概率为 $0$ 。对于每个 $w_i^r$ 对应的表示中，采样 $m$ 个注意力得分构成集合注：这里忽略了样本下标 $i$ 。具体来说， $J_r$ 由以下多项式分布产生因此，对于同一个切分元素 $r$ 而言，会使用相同的采样索引 $J_r$ 。
不同增强视图的采样后表示如下：
Attention Mutual Information Estimation
文章提出使用互信息来衡量不同视图下注意力模式的相似性。具体来说，采用对数正态分布对注意力得分分布进行建模。（torch.Tensor.log_normal_()）
两个正态分布的元组向量 $(z_1, z_2)$ 的互信息可以写为其相关性函数： ρ对应于由 $z_1$ 和 $z_2$ 计算出的相关系数。
因此，对于给定样本 $x_i$ 的第 $r$ 的切分元素而言，其对应的互信息可以写成这里的 $log(\cdot)$ 函数，用于实现从 Log-Normal到Normal随机变量的转换。对应的实现细节如下：实现的伪代码中有三点需要注意：
（1）在切分注意力tensor的时候，给出的伪代码缺少了对 $w_i$ 的切分步骤；
（2）在采样的时候就已经是将 $n \times n$ 的 $w_i^r$ 打平成1维的进行操作了；
（3）采样过后的值，应该就是将 $n^2$ 的 $w_i^r$ 变成了 $m$ 维的向量，然后进行后续的运算，最终返回互信息值的大小。
Mutual Information Aggregation
为了计算注意力正则化的损失分量，需要聚合整个张量的分布相似性。聚合是通过对批次中每个切片 $r \in R$ 和每个样本 $x_i$ 的个体相似性进行平均得到的。给定权重缩放因子 $\lambda \in \mathbb{R}$ ，权重对齐损失写成如下形式：