详解大模型微调方法Prompt Tuning(内附实现代码)

MLTalks
15 min readJun 22, 2023

--

Prompt Tuning是现在大模型微调方法中的一种常用方法,本文通过解读5篇论文来了解Prompt Tuning方法演进的过程。分别是Prefix-Tuning、P-Tuning v1、Parameter-Efficient Prompt Tuning、P-Tuning v2。

1. Prefix-Tuning:Optimizing Continuous Prompts for Generation

Finetuning之前是使用大模型进行下游任务重训的方法,但由于大模型参数量过大,Finetuning需要大量的数据,以及更多的算力去更新学习参数,不够实用。在2021年提出的prefix-tuning算法,并在自然语言生成任务(NLG, Nature Language Generation)上做了验证。这里注意区分下另一个NLP的概念,在NLP中还一类任务叫自然语言理解(NLU, Nature Language Understanding)。

在Prompt思想的启发下,在Prefix-Tuning中提出了给每一个input输入增加一个连续的任务相关的embedding向量(continuous task-specific vectors)来进行训练。

这里的连续(continuous)是相对于人工定义的文本prompt token的离散(discrete)来说的,例如一个人工定义的prompt token数组是['The', 'movie', 'is', '[MASK]'],把其中的token The替换成一个embedding向量做为输入的话,其中embedding就是连续的(continuous)表达方式。在下游任务重训的时候固定原有的大模型所有参数,只用重训跟下游任务相关的前缀向量(prefix embedding)即可。

对于自回归的LM模型(例如GPT-2)来说,会在原有prompt之前增加prefix(z = [PREFIX; x; y]);对于encoder+decoder的LM模型(例如BART)来说,会分别在encoder和decoder的输入前加上prefix(z = [PREFIX; x; PREFIX'; y],)。如下图所示,P_idx表示加的前缀序列, h对应的是可学习的参数, 用Pθ=[h1, h2, h3, ...]表示可学习参数矩阵。

直接学习参数效果不好,所以使用MLP网络对进行了reparameter修正,即Pθ[i,:] = MLP(Pθ'[i,:]),重训完只用保存prefix的相关参数。

huggingface peft关于prefix-tuning的核心代码实现在PrefixEncoder

class PrefixEncoder(torch.nn.Module):
......
if self.prefix_projection and not config.inference_mode:
# Use a two-layer MLP to encode the prefix
self.embedding = torch.nn.Embedding(num_virtual_tokens, token_dim)
self.transform = torch.nn.Sequential(
torch.nn.Linear(token_dim, encoder_hidden_size),
torch.nn.Tanh(),
torch.nn.Linear(encoder_hidden_size, num_layers * 2 * token_dim),
)
else:
self.embedding = torch.nn.Embedding(num_virtual_tokens, num_layers * 2 * token_dim)

def forward(self, prefix: torch.Tensor):
if self.prefix_projection:
prefix_tokens = self.embedding(prefix)
past_key_values = self.transform(prefix_tokens)
else:
past_key_values = self.embedding(prefix)
return past_key_values

2. P-Tuning v1:GPT Understands, Too

P-Tuning这个词在这篇文中被第一次提出,为了跟后续P-Tuning v2区分,这里加了个v1。本论文主要是为了解决GPT大模型在自然语言理解任务(NLU, Nature Language Understanding 重训效果不好的问题。在P-Tuning方法中会在连续向量空间中自动搜索合适的prompt,来增强重训练的效果。

对于之前存在的离散prompt搜索方法(discrete prompt search)来说, 比如AUTOPROMPTLPAQA, 其中的Prompt Generator通过接受离散的反馈来选择合适的prompt。

对于Prompt Generator来说,给定一个词库V和语言模型M, P_i表示在prompt模版T中第i个token,会用词库V中的词来填充模版并并生成embedding向量,

例如:[e_template(The), e_template(capital), e_template(of), e_input(Britain), e_template(is), e_output([Mask])], 其中template表示模版的,input表示输入,ouput表示输出。

而在P-Tuning中通过Prompt Encoder来实现prompt的生成,跟之前的区别在于这里使用了伪prompt和反向传播来对encoder进行更新。在embedding的输入上有所不同,模版中的prompt token embedding向量都是从Prompt Encoder生成出来的,没有对应词库中具体的词。

例如:{h0, ..., hi, e_input(Britain), hi+1, ..., hm, e([MASK])}

在网络结构上使用embedding层加上基于双层LSTM和relu激活的MLP来实现。训练过程中使用LSTM,但在推理过程中可以把LSTM给去掉。

最终在效果上实现了bert等大模型重训的提升。

huggingface peft关于p-tuning的核心代码实现在PromptEncoder

class PromptEncoder(torch.nn.Module):
......
def forward(self, indices):
input_embeds = self.embedding(indices)
if self.encoder_type == PromptEncoderReparameterizationType.LSTM:
output_embeds = self.mlp_head(self.lstm_head(input_embeds)[0])
elif self.encoder_type == PromptEncoderReparameterizationType.MLP:
output_embeds = self.mlp_head(input_embeds)
else:
raise ValueError("Prompt encoder type not recognized. Please use one of MLP (recommended) or LSTM.")

return output_embeds

3. Parameter-Efficient Prompt Tuning

本篇论文可以看成是prefix-tuning的简化版,一方面文中实验证明了使用自动生成的soft prompt方法进行tuning的效果跟model tuning差不多,同时超过了人工设计的prompt。

另一方面,文中对model tuning和prompt tuning做了如下图的对比,prompt tuning可以大幅节省参数量。对于T5的XXL的model来说,全量的model tuning每个下游任务需要11B的参数量,用prompt tuning只需要20480个参数。需要注意跟prefix-tuning不同点:这里的prompt-tuning没有包含中间层的prefix,也没有对下游任务的输出网络进行修改。在prefix-tuning中使用了MLP进行prefix的reparameter

huggingface peft关于prompt-tuning的核心代码实现在PromptEmbedding

class PromptEmbedding(torch.nn.Module):
......
def forward(self, indices):
# Just get embeddings
prompt_embeddings = self.embedding(indices)
return prompt_embeddings

4. P-Tuning v2

在模型参数量小于10B的训练中,prompt tuning效果还是不及FT(fine-tuning), P-Tuning v2支持330M~10B规模的多任务tuning。P-Tuning v2可以看成是Deep Prompt Tuning在NLU领域的实现,而Deep Prompt Tuning用于问答任务的训练。

跟之前区别主要有以下几点:

  1. 对于NLU任务没有使用像MLP的Reparameterization。
  2. 在模型的每一层上都加上了layer prompt,不同任务可以共享相同的网络参数,支持多任务学习
  3. 在分类头的verbalizer中使用了一个随机初始化的linear head
  4. Prompt长度对于简单分类任务小于20,对于像序列标注这样的复杂任务需要100左右

以bert分类任务为例,P-Tuning v2相关的核心代码实现在BertPrefixForTokenClassification

在实现的时候有一个get_prompt方法,通过这个函数提前生成各层的prompt的prefix向量

class BertPrefixForTokenClassification(BertPreTrainedModel):
def get_prompt(self, batch_size):
prefix_tokens = self.prefix_tokens.unsqueeze(0).expand(batch_size, -1).to(self.bert.device)
past_key_values = self.prefix_encoder(prefix_tokens)
# bsz, seqlen, _ = past_key_values.shape
past_key_values = past_key_values.view(
batch_size,
self.pre_seq_len,
self.n_layer * 2,
self.n_head,
self.n_embd
)
past_key_values = self.dropout(past_key_values)
past_key_values = past_key_values.permute([2, 0, 3, 1, 4]).split(2)
return past_key_values

然后在forward中通过self.bert方法中的past_key_values方法把prefix向量传入,在前向计算时会把传入的prefix向量进行拼接。

class BertPrefixForTokenClassification(BertPreTrainedModel):    
...
def forward(
self,
input_ids=None,
attention_mask=None,
token_type_ids=None,
position_ids=None,
head_mask=None,
inputs_embeds=None,
labels=None,
output_attentions=None,
output_hidden_states=None,
return_dict=None,
):
return_dict = return_dict if return_dict is not None else self.config.use_return_dict

batch_size = input_ids.shape[0]
past_key_values = self.get_prompt(batch_size=batch_size)
prefix_attention_mask = torch.ones(batch_size, self.pre_seq_len).to(self.bert.device)
attention_mask = torch.cat((prefix_attention_mask, attention_mask), dim=1)

outputs = self.bert(
input_ids,
attention_mask=attention_mask,
token_type_ids=token_type_ids,
position_ids=position_ids,
head_mask=head_mask,
inputs_embeds=inputs_embeds,
output_attentions=output_attentions,
output_hidden_states=output_hidden_states,
return_dict=return_dict,
past_key_values=past_key_values,
)
...

对于past_key_values可参考huggingface transformer中的实现【BertSelfAttention】。在forward中如果设置了past_key_value会通过torch.cat和layer的参数进行拼接。

class BertSelfAttention(nn.Module):
...
def forward(
self,
hidden_states: torch.Tensor,
attention_mask: Optional[torch.FloatTensor] = None,
head_mask: Optional[torch.FloatTensor] = None,
encoder_hidden_states: Optional[torch.FloatTensor] = None,
encoder_attention_mask: Optional[torch.FloatTensor] = None,
past_key_value: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
output_attentions: Optional[bool] = False,
) -> Tuple[torch.Tensor]:
...
elif past_key_value is not None:
key_layer = self.transpose_for_scores(self.key(hidden_states))
value_layer = self.transpose_for_scores(self.value(hidden_states))
key_layer = torch.cat([past_key_value[0], key_layer], dim=2)
value_layer = torch.cat([past_key_value[1], value_layer], dim=2)
...
...

5. 参考

--

--