huggingface 仅训练新的令牌嵌入

Question

我想向 Huggingface 模型添加新令牌并仅训练它们的嵌入。我怎样才能做到这一点？

有一些方法可以仅训练部分权重张量（例如 https://discuss.pytorch.org/t/how-can-i-freeze-weights-in-element-wise/117988），但是它们似乎可以访问和编辑模型类，这可能需要为我将要研究的每个特定 LLM 做大量工作，并且我打算对许多不同的模型进行此培训。

另一种方法是训练所有模型，并在每一步之后重新分配除新标记的嵌入之外的所有权重，这听起来效率低下，甚至对于某些优化器来说可能是不正确的（如上面的链接中所述）。

Answer 1

首先使用标记生成器添加新标记，然后调整模型的标记嵌入的大小：

from transformers import AutoTokenizer, AutoModelForCausalLM

model_name = "your-model-name"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

new_tokens = ["new_token1", "new_token2"]
tokenizer.add_tokens(new_tokens)
model.resize_token_embeddings(len(tokenizer))

然后，冻结除新令牌嵌入之外的所有参数

for param in model.parameters():
    param.requires_grad = False

# Get the new embedding indices
new_token_ids = tokenizer.convert_tokens_to_ids(new_tokens)
new_embeddings = model.get_input_embeddings().weight.data[new_token_ids]

# Unfreeze the new token embeddings
for idx in new_token_ids:
    model.get_input_embeddings().weight.data[idx].requires_grad = True

最后，实现训练循环，确保仅更新新的嵌入：

from torch.optim import AdamW
from torch.nn import functional as F

optimizer = AdamW(model.get_input_embeddings().parameters(), lr=1e-4)

for epoch in range(num_epochs):
    for batch in dataloader:
        inputs = batch["input_ids"]
        labels = batch["labels"]

        optimizer.zero_grad()
        outputs = model(inputs, labels=labels)
        loss = outputs.loss
        loss.backward()
        optimizer.step()

这样您就可以避免重新分配权重的低效率，并且可以直接将其应用于不同的模型。