prefix turning

Prefix Tuning: Optimizing Continuous Prompts for Large Language Models

原论文链接:arXiv:2101.00190

项目主页:GitHub - XiangLi1999/PrefixTuning

研究背景

大型语言模型(LLMs)在自然语言处理任务中展现出强大的能力。然而,为了使这些模型适应特定下游任务,通常需要进行微调。传统的全量微调方法涉及更新模型的所有参数,这不仅计算成本高昂,而且对存储资源的需求也很大。提示工程(Prompt Engineering)作为一种无需修改模型参数即可引导模型行为的方法,虽然可以降低成本,但在复杂任务上表现往往不如微调。为了结合两者的优势,同时避免各自的缺点,Prefix Tuning应运而生,它旨在通过优化少量的“前缀”参数来高效地适配大型模型。

核心思想

Prefix Tuning 的核心思想是:在不修改大型预训练模型参数的情况下,通过在输入序列前添加一小段可训练的连续向量(即“前缀”),来引导模型执行特定任务。这些前缀向量可以看作是任务特定的“软提示”或“虚拟token”,它们与原始输入一起输入到模型中,从而在推理过程中影响模型的行为。

通过这种方式,Prefix Tuning 实现了以下目标:

  • 参数高效:只训练前缀参数,模型的主体参数保持冻结,大大减少了可训练参数数量。
  • 性能接近全量微调:在许多任务上,Prefix Tuning 的性能可以与全量微调相媲美,甚至在某些情况下超越提示工程。
  • 易于任务切换:不同的任务只需替换不同的前缀,而无需重新加载整个模型。

方法流程详解

Prefix Tuning 的实现过程如下:

  1. 冻结原始模型参数:预训练语言模型的原始权重(例如 GPT-2 或 T5)保持不变。
  2. 引入可训练前缀:在模型的每一层或仅在输入层,引入一系列可训练的连续向量,这些向量构成了“前缀”。对于自回归模型(如GPT-2),前缀添加到输入序列的开头;对于编码器-解码器模型(如T5),前缀同时添加到编码器和解码器的输入中。
  3. 训练前缀参数:在特定任务的数据集上,仅训练这些前缀向量的参数,以最小化任务相关的损失函数。模型的主体参数在整个训练过程中保持冻结。
  4. 推理阶段:在推理时,将训练好的前缀向量与新的输入拼接在一起,然后将组合后的序列输入到冻结的预训练模型中进行预测。

这种方法的优势在于,它将任务特定的知识编码到一小组连续的前缀向量中,从而实现了参数的高效适配。

实验设置与结果

研究团队在多个生成式任务上对 Prefix Tuning 进行了评估,包括表格到文本生成(E2E NLG)、摘要(XSum)和问答(WebNLG)。实验结果表明,Prefix Tuning 在以下方面表现出色:

  • 参数效率:在 GPT-2 上,Prefix Tuning 仅需训练0.1%的模型参数,即可达到与全量微调相近的性能。
  • 性能表现:
    • 在E2E NLG任务上,Prefix Tuning 在数据量较少时优于全量微调,并且在完整数据集上与全量微调性能相当。
    • 在XSum摘要任务上,Prefix Tuning 表现优于全量微调,达到了新的SOTA。
    • 在WebNLG问答任务上,Prefix Tuning 在低数据量和完整数据集设置下都表现出色。
  • 鲁棒性:Prefix Tuning 在低数据量场景下表现更稳定。

此外,Prefix Tuning 在训练和推理阶段都不会引入显著的额外延迟,并且由于参数量极小,可以轻松切换不同任务的前缀。

总结

Prefix Tuning 提供了一种高效且有效的大型语言模型微调方法,它通过优化少量的连续前缀向量来适配模型,从而显著降低了计算资源的需求,同时保持了模型的性能。其核心思想是利用可训练的软提示来引导冻结的预训练模型。

如何使用

在实践中,使用 Prefix Tuning 对大型语言模型进行高效微调,可以显著减少训练所需的参数和计算资源。以下是一个典型的 Prefix Tuning 微调流程,适用于如 GPT-2、T5 等模型,结合了 Hugging Face Transformers 和 PEFT(Parameter-Efficient Fine-Tuning)库的使用。

1. 安装必要的库

确保安装了以下 Python 库:

1
pip install torch transformers datasets peft

2. 加载预训练模型和分词器

以 GPT-2 模型为例,加载预训练模型和对应的分词器:

1
2
3
4
5
6
7
8
9
from transformers import AutoTokenizer, AutoModelForCausalLM

model_name = "gpt2"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

# GPT-2 does not have a padding token by default, which is often needed for batching.
if tokenizer.pad_token is None:
tokenizer.pad_token = tokenizer.eos_token

3. 应用 Prefix Tuning 配置

使用 PEFT 库配置 Prefix Tuning 参数,并将其应用到模型中:

1
2
3
4
5
6
7
8
9
10
11
12
from peft import get_peft_model, PrefixTuningConfig, TaskType

prefix_config = PrefixTuningConfig(
task_type=TaskType.CAUSAL_LM, # For causal language modeling tasks
num_virtual_tokens=20, # Number of virtual tokens in the prefix
encoder_hidden_size=model.config.hidden_size # Hidden size of the encoder (for GPT-2, it's model.config.n_embd)
)

model = get_peft_model(model, prefix_config)

# Print the number of trainable parameters
model.print_trainable_parameters()

在此配置中:

  • task_type:指定任务类型,例如 TaskType.CAUSAL_LM 用于自回归语言模型。
  • num_virtual_tokens:前缀中虚拟token的数量,这个值通常需要根据任务和模型进行调整。
  • encoder_hidden_size:模型的隐藏层维度。

4. 准备训练数据

使用 Hugging Face 的 datasets 库加载和预处理训练数据。对于文本生成任务,你需要将输入和目标文本拼接起来,并进行tokenization。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
from datasets import load_dataset

# Example: a simple text dataset
data_files = {"train": "train.txt", "validation": "validation.txt"}
dataset = load_dataset("text", data_files=data_files)

def tokenize_function(examples):
# This example assumes each line in your text file is a separate example.
# For actual tasks, you'd structure your data as input-output pairs.
return tokenizer(examples["text"], truncation=True, max_length=128)

tokenized_dataset = dataset.map(tokenize_function, batched=True, remove_columns=["text"])

# Format for PyTorch
tokenized_dataset.set_format("torch")

# Data collator for language modeling, which handles padding
from transformers import DataCollatorForLanguageModeling
data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)

5. 定义训练参数和训练器

设置训练参数,并使用 Trainer 进行模型训练:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
from transformers import Trainer, TrainingArguments

training_args = TrainingArguments(
output_dir="./prefix-tuning-gpt2",
per_device_train_batch_size=4,
num_train_epochs=3,
learning_rate=5e-5, # Typically a lower learning rate than full fine-tuning
logging_dir="./logs",
save_total_limit=2,
save_steps=500,
evaluation_strategy="steps",
eval_steps=500,
logging_steps=100,
load_best_model_at_end=True,
report_to="none" # Disable reporting to services like Weights & Biases if not needed
)

trainer = Trainer(
model=model,
args=training_args,
train_dataset=tokenized_dataset["train"],
eval_dataset=tokenized_dataset["validation"],
data_collator=data_collator,
)

6. 开始训练

启动训练过程:

1
trainer.train()

7. 保存和加载 Prefix Tuning 模型

训练完成后,保存 Prefix Tuning 模块的权重:

1
model.save_pretrained("prefix-tuning-gpt2")

在需要时,加载保存的模型进行推理或进一步训练。注意,加载时需要先加载原始预训练模型,然后加载PEFT模型:

1
2
3
4
5
6
7
8
9
from peft import PeftModel

# Load the base model first
base_model = AutoModelForCausalLM.from_pretrained(model_name)

# Load the PeftModel (Prefix Tuning weights) on top of the base model
peft_model = PeftModel.from_pretrained(base_model, "prefix-tuning-gpt2")

# Now peft_model can be used for inference or further training
Contents
  1. 1. Prefix Tuning: Optimizing Continuous Prompts for Large Language Models
    1. 1.1. 研究背景
    2. 1.2. 核心思想
    3. 1.3. 方法流程详解
    4. 1.4. 实验设置与结果
    5. 1.5. 总结
    6. 1.6. 如何使用
      1. 1.6.1. 1. 安装必要的库
      2. 1.6.2. 2. 加载预训练模型和分词器
      3. 1.6.3. 3. 应用 Prefix Tuning 配置
      4. 1.6.4. 4. 准备训练数据
      5. 1.6.5. 5. 定义训练参数和训练器
      6. 1.6.6. 6. 开始训练
      7. 1.6.7. 7. 保存和加载 Prefix Tuning 模型
|