💨

LoRAとNEFTUNEの関係性

2024/12/02に公開

NEFTUNE(Noise Enhanced Fine-Tuning)は、LoRA(Low-Rank Adaptation)の embedding 層以外にも効果を発揮 する可能性があります。特に、自己注意層(Attention)フィードフォワード層(Feed-Forward Layers) のような、LoRA が通常適用される層においても有効です。


NEFTUNE と LoRA の関係

NEFTUNE は、トレーニング中にノイズを注入してモデルのロバスト性を向上させる技術です。LoRA の学習では以下の点で有効に働くと考えられます:

  1. ノイズの効果が期待される層

    • 自己注意層(q_proj, k_proj, v_proj, o_proj:
      • ノイズによって文脈理解の多様性を増やし、モデルがより広範な入力に対応可能になります。
    • フィードフォワード層(gate_proj, up_proj, down_proj:
      • 特定の非線形変換を補正する過程で、ノイズが新しい一般化を促進する可能性があります。
  2. ロバスト性と汎化性能の向上

    • ノイズは過学習を抑制し、少ないパラメータで汎化性能を向上させる効果があります。
    • LoRA はパラメータ効率が高いため、NEFTUNE のノイズ注入と相性が良いです。

LoRA の層別に見る NEFTUNE の効果

1. Embedding 層

  • NEFTUNE のノイズは、特定のトークン埋め込み(embed_tokens)に多様性を加えます。
  • 効果:
    • 語彙やトークンベースのタスクでのモデル性能を向上。
    • 例えば、異なる文脈での同じ単語の扱いが改善される。

2. 自己注意層

  • LoRA が一般的に適用される q_proj, k_proj, v_proj, o_proj にノイズを加えると、文脈や関係性の表現力が向上する可能性があります。
  • 効果:
    • よりロバストな文脈表現が得られ、特に長文の入力やノイズの多いデータに対する性能が向上。

3. フィードフォワード層

  • gate_proj, up_proj, down_proj に適用した場合、特定の非線形変換に多様性を与えることで、特定タスクの適応性が向上します。
  • 効果:
    • 生成タスクや特定の数値的変換が必要なタスクで性能向上が期待されます。

NEFTUNE の適用例: Embedding 層以外へのノイズ注入

以下の例では、LoRA を自己注意層とフィードフォワード層に適用し、NEFTUNE のノイズを注入しています。

コード例

from peft import get_peft_model, LoraConfig, TaskType
from transformers import AutoModelForCausalLM, AutoTokenizer, Trainer, TrainingArguments
import torch

# モデルとトークナイザの準備
model_name = "gpt2"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

# LoRA 設定(注意層と FFN に適用)
lora_config = LoraConfig(
    task_type=TaskType.CAUSAL_LM,
    r=16,
    lora_alpha=32,
    target_modules=["q_proj", "v_proj", "gate_proj", "down_proj"],  # 注意層と FFN
    lora_dropout=0.1
)
model = get_peft_model(model, lora_config)

# NEFTUNE: 入力にノイズを注入するカスタムデータセット
def add_noise(batch, noise_level=0.02):
    noisy_input = batch["input_ids"] + torch.normal(0, noise_level, batch["input_ids"].shape).long()
    return {"input_ids": noisy_input, "labels": batch["labels"]}

# ダミーデータセット
train_data = [
    {"input_ids": tokenizer("What is 2 + 3?", return_tensors="pt").input_ids,
     "labels": tokenizer("The answer is 5.", return_tensors="pt").input_ids}
]
train_dataset = [{"input_ids": add_noise(d)["input_ids"], "labels": d["labels"]} for d in train_data]

# トレーニング設定
training_args = TrainingArguments(
    output_dir="./lora_neftune",
    per_device_train_batch_size=4,
    num_train_epochs=3,
    learning_rate=5e-4,
    optim="adamw_torch_fused",
    fp16=True
)

# トレーナーの準備
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
)
trainer.train()

効果が期待できる場面

  1. 汎化性能が重要な場合

    • NEFTUNE によるノイズ注入が、未知の文脈や新しいデータに対する応答性を向上させます。
  2. 計算や文脈の厳密さが必要なタスク

    • Attention 層や Feed-Forward 層での NEFTUNE 適用は、文脈的な情報をより適切に捉えるため、計算や論理的推論タスクに有効です。
  3. ノイズの多い入力データ

    • 現実のデータがノイズを含む場合、NEFTUNE によりモデルがよりロバストになります。

注意点

  1. ノイズレベルの設定:

    • 適切なノイズレベル(例: 0.01〜0.05)を選ぶ必要があります。ノイズが強すぎると学習が安定しなくなる可能性があります。
  2. トレーニング時間の増加:

    • ノイズ処理により計算コストが若干増加する場合があります。
  3. パフォーマンスの検証:

    • LoRA の特定層における効果がタスクによって異なるため、NEFTUNE を適用する前後で性能を比較する必要があります。

結論

NEFTUNE は LoRA の embedding 層以外(特に Attention 層や Feed-Forward 層)にも効果があります。ノイズ注入によりモデルの汎化性能やロバスト性が向上し、特定のタスクで性能向上が期待できます。ノイズレベルや適用範囲を慎重に設定することで、LoRA モデルの特性を最大限に活かすことが可能です。

Discussion