OSF (Ajuste Fino em Subespaços Ortogonais)

O Fine-tuning em Subespaço Ortogonal (OSF) é um método PEFT projetado para aprendizado contínuo que restringe as atualizações de parâmetros para que sejam ortogonais às direções importantes anteriores. Essa abordagem permite o fine-tuning completo, evitando o esquecimento catastrófico sem exigir parâmetros adicionais ou o armazenamento de gradientes anteriores.

O resumo do artigo é:

O aprendizado contínuo em grandes modelos de linguagem (LLMs) é propenso ao esquecimento catastrófico, no qual a adaptação a novas tarefas degrada significativamente o desempenho em tarefas previamente aprendidas. Os métodos existentes geralmente dependem de atualizações de baixo posto e eficientes em parâmetros, que limitam a expressividade do modelo e introduzem parâmetros adicionais por tarefa, levando a problemas de escalabilidade. Para abordar essas limitações, propomos uma nova abordagem de ajuste fino completo contínuo que aproveita a decomposição em valores singulares (SVD) adaptativa. Nosso método identifica dinamicamente subespaços de parâmetros de baixo posto específicos da tarefa e restringe as atualizações para que sejam ortogonais às direções críticas associadas a tarefas anteriores, minimizando assim efetivamente a interferência sem sobrecarga adicional de parâmetros ou armazenamento de gradientes de tarefas anteriores. Avaliamos extensivamente nossa abordagem em benchmarks padrão de aprendizado contínuo, utilizando modelos encoder-decoder (T5-Large) e apenas decoder (LLaMA-2 7B), abrangendo diversas tarefas, incluindo classificação, geração e raciocínio. Empiricamente, nosso método alcança resultados state-of-the-art, com até 7% de aumento na precisão média em comparação com baselines recentes como O-LoRA, e notavelmente mantém as capacidades linguísticas gerais do modelo, a precisão no seguimento de instruções e a segurança ao longo do processo de aprendizado contínuo, reduzindo o esquecimento a níveis quase insignificantes. Nosso framework de SVD adaptativa equilibra efetivamente a plasticidade do modelo e a retenção de conhecimento, fornecendo uma solução prática, fundamentada teoricamente e computacionalmente escalável para cenários de aprendizado contínuo em grandes modelos de linguagem.

Como o OSF Funciona

O OSF decompõe cada matriz de pesos em componentes de alto rank (congelados) e baixo rank (treináveis) usando SVD:

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

Onde:

• U_high, S_high, V_high: Preserva as direções importantes das tarefas anteriores (congeladas)
• U_low, S_low, V_low: Permitem adaptação a novas tarefas (treináveis)

Durante o treinamento, os gradientes são projetados para serem ortogonais ao subespaço de alto rank, garantindo que as atualizações não interfiram no conhecimento previamente aprendido.

Uso Básico

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

Opções de Configuração

Módulos Alvo

Você pode especificar os módulos de destino de várias maneiras:

# 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

Configuração de Rank Efetivo

Controle os subespaços preservados/treináveis:

# 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

Nota: o effective_rank do OSF é o rank preservado (congelado), não o rank treinável. O rank treinável é igual a min(weight.shape) - effective_rank. Isso difere do r do LoRA, que especifica diretamente o rank treinável.

Dicas de Treinamento para Aprendizado Contínuo

Aprendizado de Tarefas Sequenciais

O OSF foi projetado especificamente para tarefas de aprendizado sequencial. Entre as tarefas, recalcule o SVD para que o subespaço preservado reflita os pesos mais recentes. Uma maneira simples é reenvolver o modelo base atualizado com o OSF novamente:

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

Alocação de Orçamento para Sequências de Tarefas

Ao treinar em uma sequência conhecida de n tarefas, uma estratégia eficaz é alocar progressivamente a capacidade do modelo para equilibrar o aprendizado de novas tarefas com a preservação do conhecimento anterior:

• Tarefa 1: Use a capacidade total (treine tudo)
• Tarefa 2: Congelar 1/n da capacidade do modelo, treinar a capacidade restante (n-1)/n
• Tarefa 3: Congelar 2/n da capacidade do modelo, treinar a capacidade restante (n-2)/n
• Tarefa n: Congelar (n-1)/n da capacidade do modelo, usar 1/n da capacidade para a tarefa final

Essa abordagem garante que cada tarefa receba capacidade de aprendizado adequada, preservando progressivamente mais conhecimento das tarefas anteriores.

# 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])

Boas Práticas

1. Seleção eficaz do rank: comece com effective_rank=None (o rank é definido automaticamente como 50% da menor dimensão do peso por módulo alvo) e ajuste com base na complexidade da tarefa
2. Taxa de aprendizado: Use taxas de aprendizado menores (1e-5 a 1e-4) em comparação com o ajuste fino padrão
3. Importância da tarefa: Use rank_pattern para alocar mais capacidade para módulos críticos
4. Arquitetura do modelo: OSF funciona melhor com arquiteturas de transformer que possuem separação clara entre atenção e MLP
5. Planejamento de capacidade: Para sequências de tarefas conhecidas, use alocação progressiva de orçamento (congelamento de 1/n, 2/n, …, (n-1)/n) para equilibrar plasticidade e estabilidade

Considerações de memória

O OSF modifica os pesos in-place e não adiciona parâmetros, tornando-o eficiente em memória:

# 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

Uso Avançado

Módulos de destino personalizados

Para modelos com arquiteturas não padrão:

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

Integração com Outros Métodos

O OSF pode ser combinado com outras técnicas:

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

Visão geral do benchmark

API

OSFConfig

OSFModel

Funções Utilitárias

Decomposição de Pesos

Projeção do Gradiente