提高负样本的方法

[Jinkin]

目前负样本的数量还是受限于单卡的显存，即最多为batch size per gpu。
是否有这种操作：把GPU2的样本作为GPU1样本的负样本呢？

[yuxin.wang]

目前我能想到最简单直接的方法就是使用 FSDP，把模型的权重，优化器和梯度的状态分配给不同的机器，从而增加 batch_size。

稍微复杂一点的就是把最终的结果汇聚到一张卡上，计算 loss 后再广播给所有的节点，但我没写过这部分的代码，所以不确定容不容易实现，理论上应该挺容易的。

[Jinkin]

我写了个简单的实现，在我自己的机器上能够跑，后面我也看看这样做的效果。
这里分享下:

import os
from accelerate import Accelerator
class AllGather_multi(torch.autograd.Function):
    @staticmethod
    def forward(ctx, tensor, accelerator: Accelerator):
        ctx.rank = int(os.environ['RANK'])
        ctx.batch_size = tensor.shape[0]
        tensor = accelerator.gather(tensor)
        return tensor

    @staticmethod
    def backward(ctx, grad_output):
        return (
            grad_output[ctx.batch_size * ctx.rank : ctx.batch_size * (ctx.rank + 1)],
            None, None,
        )

class EmbedderForPairAllNegTrain(EmbedderForPairInBatchNegTrain):
    def __init__(
        self,
        accelerator: Accelerator,
        *args, **kwargs,
    ):
        super().__init__(*args, **kwargs)
        self.accelrator = accelerator

    def forward(self, text_ids: torch.Tensor, text_pos_ids: torch.Tensor) -> dict[str, torch.Tensor]:
        text_embeddings = self.embedder(text_ids)
        text_pos_embeddings = self.embedder(text_pos_ids)
        text_embeddings = AllGather_multi.apply(text_embeddings, self.accelrator)
        text_pos_embeddings = AllGather_multi.apply(text_pos_embeddings, self.accelrator)
        loss = self.criterion(text_embeddings, text_pos_embeddings)
        return {'loss': loss}

初始化EmbedderForPairAllNegTrain时候传入一下accelerator就可以了

[Jinkin]

microsoft/unilm#1120 (comment)
但是从E5作者的issue来看，似乎hard negative mining更为重要。
对于M3E的数据集的话，做hard negative mining。
我的理解是先建立一个样本池，对于每一个样本，都把它和样本池里的其他样本用模型作比较，然后筛出top k score的，作为hard negative。
不过这个样本池应该怎样建立比较好呢，首先它们都应该来自同一个task，同一种domain。emmmmm，暂时也没想到更多了.......

[yuxin.wang]

1. 关于如何增加 batch_size？
非常感谢您能提供一个具体实现的 demo，我会测试一下，如果一切正常，我会把这个方法尽快加到库里。
2. hard negative mining 更为重要？
是的，我们发现 m3e 的模型在 reranking 上面的表现并不好，究其原因是缺少 hard negative sample ，导致模型对于文本细节的区分能力较差。
3. 如何构建样本池
我觉得来自同一个 task 就可以 。我能想到比较好的方案是，先通过 m3e 来做召回，在使用下面的模型来做精排。

• openai
• m3e-base
• sgpt 的 cross 方法，使用 chatglm2，百川之类的
  用上面多个模型投票（其他 ensemble 方法也可以）一起来判断 top k

你现在做的方向其实也是 m3e 想要优化的方向，如果有任何进展，欢迎交流~

[Jinkin]

突然想起来，简单地update下:
上面那个demo我这里实际跑起来效果几乎没有提升，而由于需要在forward的时候对tensor进行gather，会增加GPU的显存负载，导致batch size参数还需要设低一些。
我后面的实验也几乎没用这种gather方式了，不过具体不work的原因的话，目前还没深入研究。
以后有时间我会再看看。

[yuxin.wang]

收到