Prefix tuning

Prefix tuning добавляет в начало входной последовательности набор векторов, специфичных для задачи, которые можно обучать, сохраняя предобученную модель замороженной. Параметры префикса вставляются во все слои модели.

Аннотация из статьи:

Дообучение (fine-tuning) является фактическим стандартом для использования крупных предобученных языковых моделей в задачах, следующих за предобучением. Однако этот подход изменяет все параметры языковой модели, что требует хранения полной копии модели для каждой задачи. В данной статье мы предлагаем prefix-tuning — лёгкую альтернативу дообучению для задач генерации естественного языка, которая сохраняет параметры языковой модели замороженными, но оптимизирует небольшой непрерывный вектор, специфичный для задачи (называемый префиксом). Prefix-tuning заимствует идеи из промптинга, позволяя последующим токенам обращать внимание на этот префикс, как если бы он состоял из «виртуальных токенов». Мы применяем prefix-tuning к GPT-2 для генерации текста из таблиц и к BART для суммаризации. Мы обнаруживаем, что, обучая лишь 0,1\% параметров, prefix-tuning обеспечивает сопоставимую производительность в условиях полного набора данных, превосходит дообучение в условиях малого объёма данных и лучше обобщается на примеры с темами, не встречавшимися во время обучения.

Примечание Для моделей с архитектурой «кодировщик-декодировщик» (seq2seq) префикс применяется только к декодеру, что не соответствует спецификации из статьи (см., например, рисунок 2). Метод prefix tuning всё ещё можно дообучать на таких архитектурах, однако качество может быть ниже ожидаемого; для моделей с архитектурой «кодировщик-декодировщик» рекомендуется использовать другие методы PEFT.

Prefix tuning очень похож на prompt tuning. Главное отличие состоит в том, что параметры префикса вставляются во все слои модели, тогда как prompt tuning добавляет параметры промпта только к входным эмбеддингам модели. Параметры префикса также оптимизируются с помощью отдельной прямой нейронной сети (FFN), а не обучаются напрямую на мягких промптах, поскольку это вызывает нестабильность и ухудшает производительность. FFN отбрасывается после обновления мягких промптов.

В результате авторы установили, что prefix tuning демонстрирует сопоставимую с полным дообучением модели производительность, несмотря на то, что содержит в 1000 раз меньше параметров, и показывает ещё лучшие результаты в условиях малого объёма данных.

Базовое использование

Создайте PrefixTuningConfig с типом задачи и количеством виртуальных токенов, которые нужно добавить и обучить.

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()

Возможные инициализации

По умолчанию prefix tuning использует виртуальные токены с случайной инициализацией. Также есть возможность инициализировать векторы так, чтобы они были близки к no-op (при инициализации нулями они всё равно немного смещают распределение вероятностей). Это означает, что префиксы, внедряемые в KV-кэш, оказывают меньшее влияние с самого начала, что снижает дисперсию в показателях производительности обучения.

PEFT также предоставляет утилиты для инициализации адаптера prefix-tuning из существующего префикса KV-кэша (например, из первых p токенов промпта/корпуса). Это поддерживается только при prefix_projection=False (значение по умолчанию), поскольку в этом случае обучаемыми параметрами является сам 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,
)

Убедитесь, что текст достаточно длинный, чтобы сгенерировать как минимум num_virtual_tokens токенов, иначе инициализация завершится ошибкой.

В качестве ориентира:

• начните с нейтральной стартовой последовательности, используя initialize_kv_prefix_from_text, это может быть очень короткая строка, например “Question: ”
• если это не помогает, используйте более длинную последовательность, релевантную задаче (то есть специально сконструированный промпт), что даст вам больше виртуальных токенов для размещения, но и позволит сильнее направлять модель
• если использовать текст инициализации невозможно или вы хотите быстро проверить, применим ли prefix tuning вообще, используйте нулевую инициализацию без проекции

Обзор бенчмарков

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)