Prefix tuning

Prefix tuning adiciona uma série de vetores específicos da tarefa ao início da sequência de entrada, que podem ser aprendidos mantendo o modelo pré-treinado congelado. Os parâmetros de prefixo são inseridos em todas as camadas do modelo.

O resumo do artigo é:

O fine-tuning é a maneira de fato de aproveitar grandes modelos de linguagem pré-treinados para executar tarefas de nível inferior. No entanto, ele modifica todos os parâmetros do modelo de linguagem e, portanto, exige armazenar uma cópia completa para cada tarefa. Neste artigo, propomos o prefix-tuning, uma alternativa leve ao fine-tuning para tarefas de geração de linguagem natural, que mantém os parâmetros do modelo de linguagem congelados, mas otimiza um pequeno vetor contínuo específico da tarefa (chamado de prefixo). O prefix-tuning se inspira no prompting, permitindo que os tokens subsequentes prestem atenção a esse prefixo como se fossem “tokens virtuais”. Aplicamos o prefix-tuning ao GPT-2 para geração de texto a partir de tabelas e ao BART para sumarização. Descobrimos que, ao aprender apenas 0,1\% dos parâmetros, o prefix-tuning obtém desempenho comparável na configuração de dados completos, supera o fine-tuning em configurações com poucos dados e se extrapola melhor para exemplos com tópicos não vistos durante o treinamento.

Nota Para modelos encoder-decoder (seq2seq), o prefixo é aplicado apenas ao decoder, o que não corresponde à especificação do artigo (veja, por exemplo, a Figura 2). O prefix tuning ainda pode ser ajustado nesses modelos, mas o desempenho pode ser inferior; considere usar outros métodos PEFT para modelos encoder-decoder.

O prefix tuning é muito semelhante ao prompt tuning. A principal diferença é que os parâmetros de prefixo são inseridos em todas as camadas do modelo, enquanto o prompt tuning adiciona os parâmetros de prompt apenas aos embeddings de entrada do modelo. Os parâmetros de prefixo também são otimizados por uma rede feed-forward (FFN) separada, em vez de treinar diretamente nos soft prompts, pois isso causa instabilidade e prejudica o desempenho. A FFN é descartada após a atualização dos soft prompts.

Como resultado, os autores constataram que o prefix tuning apresenta desempenho comparável ao do finetuning completo de um modelo, apesar de ter 1000x menos parâmetros, e seu desempenho é ainda melhor em cenários com poucos dados.

Uso Básico

Crie um PrefixTuningConfig com o tipo de tarefa e o número de tokens virtuais a serem adicionados e aprendidos.

from peft import PrefixTuningConfig, get_peft_model

peft_config = PrefixTuningConfig(task_type="CAUSAL_LM", num_virtual_tokens=20)
model = get_peft_model(model, peft_config)
model.print_healthcheck()

Inicializações Possíveis

Por padrão, o prefix tuning usa tokens virtuais inicializados aleatoriamente. Há também a opção de inicializar os vetores para ficarem próximos de uma operação nula (inicializados com zero, ainda assim deslocarão um pouco a massa de probabilidade). Isso significa que os prefixos injetados no KV-cache têm menos impacto desde o início e reduzem a variância no desempenho do treinamento.

O PEFT também fornece utilitários para inicializar um adaptador de prefix-tuning a partir de um prefixo de cache KV existente (por exemplo, a partir dos primeiros p tokens de um prompt/corpus). Isso é suportado apenas quando prefix_projection=False (o padrão), porque, nesse caso, os parâmetros aprendidos são o próprio prefixo KV.

from transformers import AutoModelForCausalLM, AutoTokenizer

from peft import PrefixTuningConfig, get_peft_model, initialize_kv_prefix_from_text

base = AutoModelForCausalLM.from_pretrained("gpt2")
tok = AutoTokenizer.from_pretrained("gpt2")

peft_cfg = PrefixTuningConfig(task_type="CAUSAL_LM", num_virtual_tokens=20, prefix_projection=False)
model = get_peft_model(base, peft_cfg)

initialize_kv_prefix_from_text(
    model,
    tok,
    text="...a long context with at least num_virtual_tokens tokens...",
    use_chat_template=False,
)m peft import PrefixTuningConfig, get_peft_model, initialize_kv_prefix_from_text

base = AutoModelForCausalLM.from_pretrained("gpt2")
tok = AutoTokenizer.from_pretrained("gpt2")

peft_cfg = PrefixTuningConfig(task_type="CAUSAL_LM", num_virtual_tokens=20, prefix_projection=False)
model = get_peft_model(base, peft_cfg)

initialize_kv_prefix_from_text(
    model,
    tok,
    text="...a long context with at least num_virtual_tokens tokens...",
    use_chat_template=False,
)

Certifique-se de que o texto seja longo o suficiente para produzir pelo menos num_virtual_tokens tokens, caso contrário a inicialização falhará.

Como diretriz:

• comece com uma sequência inicial neutra usando initialize_kv_prefix_from_text, pode ser uma string muito curta como “Question: ”
• se isso não ajudar, use uma sequência mais longa com relevância para a tarefa (ou seja, um prompt projetado), o que lhe dá mais tokens virtuais para encaixar, mas também mais direcionamento do modelo
• se não for possível usar um texto de inicialização ou se você quiser verificar rapidamente se o prefix tuning é viável, use uma inicialização zero sem projeção

Visão geral do benchmark

API

PrefixTuningConfig

PrefixEncoder

>>> from peft import PrefixEncoder, PrefixTuningConfig

>>> config = PrefixTuningConfig(
...     peft_type="PREFIX_TUNING",
...     task_type="SEQ_2_SEQ_LM",
...     num_virtual_tokens=20,
...     token_dim=768,
...     num_transformer_submodules=1,
...     num_attention_heads=12,
...     num_layers=12,
...     encoder_hidden_size=768,
... )
>>> prefix_encoder = PrefixEncoder(config)