OSF (Orthogonal Subspace Fine-tuning)

Le Fine-tuning dans un sous-espace orthogonal (OSF) est une méthode PEFT conçue pour l'apprentissage continu, qui contraint les mises à jour des paramètres à être orthogonales aux directions importantes précédentes. Cette approche permet un fine-tuning complet tout en évitant l'oubli catastrophique, sans nécessiter de paramètres supplémentaires ni le stockage des gradients précédents.

L'abstract de l'article est :

L’apprentissage continu dans les grands modèles de langage (LLM) est sujet à l’oubli catastrophique, où l’adaptation à de nouvelles tâches dégrade significativement les performances sur les tâches précédemment apprises. Les méthodes existantes s’appuient généralement sur des mises à jour de faible rang et économes en paramètres, ce qui limite l’expressivité du modèle et introduit des paramètres supplémentaires par tâche, entraînant des problèmes d’évolutivité. Pour remédier à ces limites, nous proposons une nouvelle approche de fine-tuning complet continu exploitant la décomposition en valeurs singulières (SVD) adaptative. Notre méthode identifie dynamiquement les sous-espaces de paramètres de faible rang spécifiques à chaque tâche et contraint les mises à jour pour qu’elles soient orthogonales aux directions critiques associées aux tâches antérieures, minimisant ainsi efficacement les interférences sans surcharge de paramètres supplémentaires ni stockage des gradients des tâches précédentes. Nous évaluons largement notre approche sur des benchmarks standard d’apprentissage continu en utilisant à la fois des modèles encodeur-décodeur (T5-Large) et des modèles uniquement décodeurs (LLaMA-2 7B), couvrant diverses tâches incluant la classification, la génération et le raisonnement. Empiriquement, notre méthode atteint des résultats de l’état de l’art, avec une précision moyenne jusqu’à 7 % supérieure à celle de récents points de référence comme O-LoRA, et maintient de manière notable les capacités linguistiques générales du modèle, la précision du suivi des instructions et la sécurité tout au long du processus d’apprentissage continu en réduisant l’oubli à des niveaux quasi négligeables. Notre cadre SVD adaptatif équilibre efficacement la plasticité du modèle et la rétention des connaissances, offrant une solution pratique, théoriquement fondée et évolutivement calculable pour les scénarios d’apprentissage continu dans les grands modèles de langage.

Fonctionnement d'OSF

OSF décompose chaque matrice de poids en composantes de rang élevé (gelées) et de rang faible (entraînables) à l'aide de la SVD :

W = U_high * S_high * V_high^T + U_low * S_low * V_low^T

Où :

• U_high, S_high, V_high : conserve les directions importantes des tâches précédentes (gelées)
• U_low, S_low, V_low : Permettent l'adaptation à de nouvelles tâches (entraînables)

Lors de l’entraînement, les gradients sont projetés pour être orthogonaux au sous-espace de rang élevé, garantissant que les mises à jour n’interfèrent pas avec les connaissances acquises précédemment.

Utilisation de base

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import OSFConfig, get_peft_model

# Load base model
model = AutoModelForCausalLM.from_pretrained("facebook/opt-125m")

# Configure OSF
config = OSFConfig(
    target_modules=["q_proj", "v_proj"],  # Target attention layers
    effective_rank=8,                     # Default rank for decomposition
    rank_pattern={"q_proj": 16}           # Override rank for specific modules
)

# Apply OSF
model = get_peft_model(model, config)

# Train as usual
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)

tokenizer = AutoTokenizer.from_pretrained("facebook/opt-125m")

inputs = tokenizer("Hello world", return_tensors="pt", padding=True)
loss = model(**inputs, labels=inputs.input_ids).loss
loss.backward()
optimizer.step()
optimizer.zero_grad()

Options de configuration

Modules cibles

Vous pouvez spécifier les modules cibles de plusieurs manières :

# Specific module names
config = OSFConfig(target_modules=["q_proj", "k_proj", "v_proj", "o_proj"])

# All linear layers
config = OSFConfig(target_modules="all-linear")

# Model-specific defaults (automatically detected)
config = OSFConfig()  # Uses model-appropriate defaults

Configuration du rang effectif

Contrôler les sous-espaces préservés/entraînables :

# Global preserved rank (applies to all target modules)
config = OSFConfig(effective_rank=16)  # preserves top-16 singular directions; trains the rest

# Automatic preserved rank (50% of the smaller matrix dimension per target)
config = OSFConfig(effective_rank=None)

# Per-module preserved-rank overrides
config = OSFConfig(
    effective_rank=8,
    rank_pattern={
        "q_proj": 16,      # Higher rank for query projection
        "gate_proj": 4     # Lower rank for gate projection
    }
)

# Fractional preserved rank is supported (interpreted per-target as fraction * min_dim)
config = OSFConfig(effective_rank=0.8)  # preserve 80% of min_dim; train remaining 20%
config = OSFConfig(rank_pattern={"q_proj": 0.5})  # preserve 50% on q_proj, others use global/default

Note : effective_rank d’OSF est le rang préservé (gelé), et non le rang entraînable. Le rang entraînable est égal à min(weight.shape) - effective_rank. Cela diffère de r de LoRA, qui spécifie directement le rang entraînable.

Conseils d'entraînement pour l'apprentissage continu

Apprentissage séquentiel de tâches

OSF est spécifiquement conçu pour les tâches d’apprentissage séquentielles. Entre les tâches, recomputez la SVD afin que le sous-espace préservé reflète les poids les plus récents. Une méthode simple consiste à réappliquer OSF au modèle de base mis à jour :

# Task 1: train on domain A with initial preserved subspace
r = 8  # initial effective rank to preserve
model = get_peft_model(base_model, OSFConfig(effective_rank=r))
train_task(model, task_1_data)

# Task 2: recompute SVD on updated weights and increase preserved subspace
base_model = model.unload()  # unwrap base model without assuming internals
r += 4  # grow preserved subspace to include Task 1 knowledge
model = get_peft_model(base_model, OSFConfig(effective_rank=r))
train_task(model, task_2_data)

# Task 3: recompute again and expand preserved subspace further
base_model = model.unload()
r += 4
model = get_peft_model(base_model, OSFConfig(effective_rank=r))
train_task(model, task_3_data)

Répartition du budget pour les séquences de tâches

Lors de l’entraînement sur une séquence connue de n tâches, une stratégie efficace consiste à allouer progressivement la capacité du modèle afin d’équilibrer l’apprentissage de nouvelles tâches tout en préservant les connaissances antérieures :

• Tâche 1 : Utiliser la pleine capacité (entraîner tout)
• Tâche 2 : Geler 1/n de la capacité du modèle, entraîner la capacité restante (n-1)/n
• Tâche 3 : Geler 2/n de la capacité du modèle, entraîner la capacité restante (n-2)/n
• Tâche n : Geler (n-1)/n de la capacité du modèle, utiliser 1/n de la capacité pour la tâche finale

Cette approche garantit que chaque tâche dispose d'une capacité d'apprentissage adéquate tout en préservant progressivement davantage de connaissances issues des tâches précédentes.

# Example: 4-task sequence with progressive budget allocation
n_tasks = 4
max_preserved_rank = 512  # Upper bound for preserved rank per target (heuristic)

for task_id in range(n_tasks):
    # Freeze increases over time; trainable capacity shrinks
    preserved_fraction = (task_id + 1) / n_tasks
    preserved_rank = int(max_preserved_rank * preserved_fraction)

    config = OSFConfig(
        target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
        effective_rank=preserved_rank,
    )

    print(
        f"Task {task_id + 1}: Preserving rank {preserved_rank} "
        f"({preserved_fraction:.1%} of max_preserved_rank - {max_preserved_rank} frozen); trainable rank = min_dim - preserved_rank"
    )

    model = get_peft_model(base_model, config)
    train_task(model, task_data[task_id])

Bonnes pratiques

1. Sélection du rang effectif : commencez avec effective_rank=None (le rang est automatiquement fixé à 50 % de la plus petite dimension de poids par module cible) et ajustez-le en fonction de la complexité de la tâche
2. Taux d'apprentissage : Utilisez des taux d'apprentissage plus faibles (1e-5 à 1e-4) par rapport à l'ajustement fin standard
3. Importance de la tâche : utilisez rank_pattern pour allouer davantage de capacité aux modules critiques
4. Architecture du modèle : OSF fonctionne le mieux avec les architectures de type transformer présentant une séparation claire entre l'attention et les MLP
5. Planification de la capacité : Pour les séquences de tâches connues, utilisez une allocation budgétaire progressive (gel de 1/n, 2/n, …, (n-1)/n) pour équilibrer plasticité et stabilité

Considérations relatives à la mémoire

OSF modifie les poids en place et n’ajoute pas de paramètres, ce qui le rend efficace en mémoire :

# Memory usage remains close to base model
print(f"Base model parameters: {base_model.num_parameters():,}")
print(f"OSF model parameters: {osf_model.num_parameters():,}")  # Similar count

Utilisation avancée

Modules cibles personnalisés

Pour les modèles aux architectures non standard :

config = OSFConfig(
    target_modules=["dense", "intermediate.dense"],  # Custom layer names
    effective_rank=12,
    rank_pattern={"dense": 8, "intermediate.dense": 16}
)

Intégration avec d'autres méthodes

OSF peut être combiné avec d'autres techniques :

# Use with gradient checkpointing for memory efficiency
model.gradient_checkpointing_enable()

# Apply weight decay selectively (regularizes low-rank factors to limit drift/overfitting in continual updates; keep small)
optimizer = torch.optim.AdamW([
    {"params": [p for n, p in model.named_parameters() if "U_low" in n], "weight_decay": 0.01},
    {"params": [p for n, p in model.named_parameters() if "S_low" in n], "weight_decay": 0.001},
    {"params": [p for n, p in model.named_parameters() if "V_low" in n], "weight_decay": 0.01},
], lr=1e-4)

Aperçu des benchmarks

API

OSFConfig

OSFModel

Fonctions utilitaires

Décomposition des poids

Projection du gradient