论文分享|Arxiv2024‘麦吉尔大学|LLM2Vec—将LLM转换为文本编码器

LLM本身的表征直接用于Embedding，比如用于检索/聚类/STS等任务，效果其实不太好。因此才需要将Embedding模型和大模型区分开来。本文介绍一篇将LLM转换为Embedding模型的工作，代码全开源，值得好好学习。

论文题目：LLM2Vec: Large Language Models Are Secretly Powerful Text Encoders

来源：Arxiv2024/麦吉尔大学

方向：文本编码/LLM

开源地址：https://github.com/McGill-NLP/llm2vec

转换LLM为Text Encoder的三步骤

参考作者的Tutorial ，以LLaMA2和FlashAttention为例介绍主要的三个步骤

第一步：实现双向注意力

对于llama这样的transformer模型来说，每一层都含有一个自注意力的子层，而这些自注意力在得到embedding的时候都会mask后面的词的attention，只保留前面的词语，即单向注意力。所以需要先将这个掩蔽关闭，每个词都可以和句子中所有词语进行注意力交互，从而实现双向注意力。如下图所示：

下面介绍下代码。首先需要修改llama_attention子层的实现。llama_attention有三种实现，在此我们需要将LlamaFlashAttention2中的is_causal设置为False，从而改为双向注意力。其实只修改了__init__函数

class ModifiedLlamaFlashAttention2(LlamaFlashAttention2):

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.is_causal = False  # 原Trnasformer实现中是True

LLAMA_ATTENTION_CLASSES = {
    "eager": LlamaAttention,
    "flash_attention_2": ModifiedLlamaFlashAttention2,  # 原是`LlamaFlashAttention2'
    "sdpa": LlamaSdpaAttention,
}

接下来需要将每个Decoder层中的attention改为双向注意力attention实现。同样只需要修改__init__函数。这一段直接从transformers库中复制即可

class ModifiedLlamaDecoderLayer(LlamaDecoderLayer):
    def __init__(self, config: LlamaConfig, layer_idx: int):
        nn.Module.__init__(self) # Initially, super().__init__()
        self.hidden_size = config.hidden_size
        
        # 这一行将注意力类绑定为LLAMA_ATTENTION_CLASSES中自定义的类
        self.self_attn = LLAMA_ATTENTION_CLASSES[config._attn_implementation](config=config, layer_idx=layer_idx) 

        self.mlp = LlamaMLP(config)
        self.input_layernorm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.post_attention_layernorm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)

最后需要将LlamaModel中的每一层都改为我们修改的decoder层，还是只需要修改__init__函数中。将每个layer中的改为修改的layer

class LlamaBiModel(LlamaModel):
    def __init__(self, config):
        LlamaPreTrainedModel.__init__(self, config) # Initially, super().__init__(config)
        self.padding_idx = config.pad_token_id
        self.vocab_size = config.vocab_size

        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
        self.layers = nn.ModuleList(
            [ModifiedLlamaDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]  # Initially, `LlamaDecoderLayer(config, layer_idx)`
        ) # 这一行中原是LlamaDecoderLayer(config, layer_idx)
        self.norm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.gradient_checkpointing = False

        self.post_init()

至此，通过简单修改了三行代码，就完成了对llama双向注意力的实现

第二步：使用掩蔽下一词预测（masked next token prediction ,MNTP）任务训练

原LlamaForCausalLM类使用的是LlamaModel，需要先修改为上一步的LlamaBiModel

class BiLlamaForMNTP(LlamaForCausalLM):

    def __init__(self, config, attention_dropout=0.0):
        LlamaPreTrainedModel.__init__(self, config) # Initially, super().__init__(config)
        self.model = LlamaBiModel(config)  # 原是LlamaModel(config)
        self.vocab_size = config.vocab_size
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)

        self.post_init()

至此模型部分就修改完毕了，接下来实现损失函数，直接复用LlamaForCausalLM中forward函数里的下一词预测任务，即利用第i-1个token的表征的logits来预测第i个token位置mask的词，但注意由于此时模型已经是双向注意力，所以本质上这其实类似简化版本的掩蔽语言模型（masked language modeling）任务。

# LlamaForCausalLM.forward()中的代码片段
loss = None
if labels is not None:
    # 拷贝表征 从而让小于n的tokens表征预测第n个token 
    # contiguous() 其实是强制拷贝一份数据防止和之前的logits对象共享存储空间
    shift_logits = logits[..., :-1, :].contiguous() # (batch_size,n-1,vocab_size) 即每句话从第1个token开始，第n-1个token结束的token表征 
    shift_labels = labels[..., 1:].contiguous() # (batch_size,n-1) 即每句话从第2个token开始，第n个token结束的token序列
    # 展开token表征
    loss_fct = CrossEntropyLoss() #交叉熵损失
    shift_logits = shift_logits.view(-1, self.config.vocab_size) # (batch_size * (n-1), vocab_size)
    shift_labels = shift_labels.view(-1) (batch_size * (n-1))
    # 允许模型并行
    shift_labels = shift_labels.to(shift_logits.device)
    loss = loss_fct(shift_logits, shift_labels)

训练时，仿造 examples/pytorch/language-modeling/run_mlm.py 脚本。本文主要有以下修改：

将模型类别改成了自己实现的语言模型类而不是原脚本的AutoModelForMaskedLM类
使用PEFT lora进行高效微调
使用了下划线_ 作为mask token。下面展示了data collator中的mask操作代码

class DataCollatorForLanguageModelingWithFullMasking(DataCollatorForLanguageModeling):
    def torch_mask_tokens(
        self,
        inputs: Any,
        special_tokens_mask: Optional[Any] = None,
    ) -> Tuple[Any, Any]:
        """
        为掩蔽语言模型准备inputs/labels，只要被选中为需要mask的，则100%都设置为mask_token
        """
        import torch

        labels = inputs.clone()
        # 使用一个概率 self.mlm_probability 来为每个句子选择少量token用于MLM训练
        probability_matrix = torch.full(labels.shape, self.mlm_probability)
        if special_tokens_mask is None:
            special_tokens_mask = [
                self.tokenizer.get_special_tokens_mask(
                    val, already_has_special_tokens=True
                )
                for val in labels.tolist()
            ]
            special_tokens_mask = torch.tensor(special_tokens_mask, dtype=torch.bool)
        else:
            special_tokens_mask = special_tokens_mask.bool()

        probability_matrix.masked_fill_(special_tokens_mask, value=0.0)
        masked_indices = torch.bernoulli(probability_matrix).bool()
        labels[~masked_indices] = -100  # 只计算masked token处的损失

        # 对于需要mask的token，100%使用mask_token替换
        inputs[masked_indices] = self.tokenizer.convert_tokens_to_ids(
            self.tokenizer.mask_token
        )

        return inputs, labels