11  O Loop de Treinamento e Otimização

Implementação do ciclo de treinamento, cálculo da função de perda (Cross-Entropy), Backpropagation e configuração do otimizador AdamW.

O “coração” de qualquer sistema de Deep Learning é o loop de treinamento. É neste ciclo iterativo que os parâmetros do modelo (pesos e vieses) são ajustados matematicamente para minimizar o erro entre as previsões do modelo e os dados reais.

Neste capítulo, dissecamos a arquitetura do ciclo de treinamento, a matemática por trás da função de perda Cross-Entropy, a mecânica da Backpropagation e a justificativa para o uso do otimizador AdamW em arquiteturas baseadas em Transformers.

11.1 1. Visão Geral do Ciclo de Treinamento

No pré-treinamento autossupervisionado, os próprios textos fornecem os rótulos: cada sequência de entrada é deslocada para criar os tokens-alvo. Isso permite aprender com grandes corpora sem anotação manual, mas transfere ao modelo padrões, lacunas e vieses presentes nas fontes. Deduplicação, normalização, remoção de dados pessoais e balanceamento entre domínios fazem parte do problema de modelagem, não apenas de limpeza.

Qualidade e diversidade importam tanto quanto volume. Um corpus deve representar idiomas, gêneros e conhecimentos relevantes para a aplicação, respeitando licenças e privacidade. Conjuntos contaminados por exemplos de avaliação tornam métricas posteriores pouco confiáveis.

O treinamento de uma rede neural é um processo de otimização estocástica. O objetivo é encontrar um conjunto de parâmetros \(\theta\) que minimize uma função de custo \(J(\theta)\).

O fluxo de dados ocorre em duas direções: 1. Forward Pass (Propagação): Os dados fluem da entrada para a saída, gerando previsões (logits). 2. Backward Pass (Retropropagação): O erro flui da saída para a entrada, calculando o gradiente da perda em relação a cada parâmetro.

11.1.1 Diagrama de Fluxo do Treinamento

O diagrama abaixo ilustra a interação entre os dados, o modelo, a função de perda e o otimizador.

flowchart TD
    subgraph Data_Preparation [Preparação de Dados]
        Batch[Batch de Treinamento]
        Targets[Labels/Targets]
    end

    subgraph Forward_Pass [Forward Pass]
        Params[(Parâmetros do Modelo)]
        Model[Arquitetura Transformer]
        Logits[Logits / Previsões]
        
        Params --> Model
        Batch --> Model
        Model --> Logits
    end

    subgraph Optimization_Cycle [Ciclo de Otimização]
        LossFn{Cross-Entropy Loss}
        Gradients[Gradientes]
        Optimizer[Otimizador AdamW]
        
        Logits --> LossFn
        Targets --> LossFn
        LossFn -- "loss.backward()" --> Gradients
        Gradients --> Optimizer
        Optimizer -- "step()" --> Params
    end

    style LossFn fill:#f96,stroke:#333,stroke-width:2px
    style Optimizer fill:#69f,stroke:#333,stroke-width:2px
    style Params fill:#eee,stroke:#333,stroke-width:2px

11.2 2. A Função de Perda: Cross-Entropy

Para modelos de linguagem (LLMs), o objetivo é prever o próximo token em uma sequência. Isso é tratado como um problema de classificação multiclasse, onde o número de classes é igual ao tamanho do vocabulário (\(V\)).

A função de perda padrão é a Cross-Entropy (Entropia Cruzada). Ela mede a divergência entre duas distribuições de probabilidade: 1. \(P\) (Distribuição Real): Uma distribuição one-hot representando o token verdadeiro. 2. \(Q\) (Distribuição Predita): As probabilidades geradas pelo modelo (após a aplicação de Softmax nos logits).

A fórmula matemática para um único exemplo é:

\[ H(P, Q) = -\sum_{x} P(x) \log Q(x) \]

No contexto de PyTorch, a classe nn.CrossEntropyLoss combina duas operações para estabilidade numérica: LogSoftmax e NLLLoss (Negative Log Likelihood Loss).

Por que Cross-Entropy? Ela penaliza fortemente previsões confiantes, porém erradas. Se o modelo atribui uma probabilidade baixa ao token correto, o logaritmo dessa probabilidade resulta em um valor negativo grande, aumentando drasticamente a perda.

11.3 3. O Otimizador: AdamW

Embora o Stochastic Gradient Descent (SGD) seja a base, modelos modernos (especialmente Transformers) utilizam variantes adaptativas. O padrão da indústria atual é o AdamW.

11.3.1 Adam vs. AdamW

  • Adam (Adaptive Moment Estimation): Calcula taxas de aprendizado adaptativas para cada parâmetro. Ele armazena médias móveis exponenciais dos gradientes passados (\(m_t\), primeiro momento) e dos gradientes ao quadrado (\(v_t\), segundo momento).
  • O Problema do Adam: A implementação original do Adam aplicava a regularização L2 (Weight Decay) junto com o cálculo do gradiente adaptativo. Isso acoplava a regularização à magnitude dos gradientes, o que muitas vezes levava a uma generalização subótima.
  • A Solução AdamW (Decoupled Weight Decay): O AdamW desacopla o decaimento de peso da atualização do gradiente. O decaimento de peso é aplicado diretamente aos parâmetros após o passo de otimização principal.

A atualização de parâmetros no AdamW segue a lógica:

  1. Atualiza parâmetros baseados nos gradientes adaptativos (Adam clássico).
  2. Aplica decaimento de peso: \(\theta_{t} = \theta_{t} - \eta \lambda \theta_{t-1}\) (onde \(\lambda\) é o fator de decaimento).

Isso permite usar taxas de decaimento de peso mais agressivas sem prejudicar o aprendizado, resultando em modelos que generalizam melhor.

11.4 4. Implementação Técnica

Abaixo apresentamos uma implementação robusta do loop de treinamento utilizando PyTorch.

11.4.1 Estrutura do Código

import torch
import torch.nn as nn
from torch.optim import AdamW
from torch.utils.data import DataLoader

def train_epoch(
    model: nn.Module, 
    dataloader: DataLoader, 
    optimizer: torch.optim.Optimizer, 
    device: torch.device,
    clip_grad: float = 1.0
) -> float:
    """
    Executa uma época de treinamento.
    
    Args:
        model: O modelo Transformer.
        dataloader: Iterador de dados em batches.
        optimizer: Otimizador configurado (AdamW).
        device: CPU ou CUDA.
        clip_grad: Valor para recorte de gradiente (Gradient Clipping).
        
    Returns:
        Média da perda (loss) na época.
    """
    model.train() # Coloca o modelo em modo de treinamento (ativa Dropout, etc.)
    total_loss = 0.0
    
    # Definição da Função de Perda
    # ignore_index é usado para ignorar tokens de padding no cálculo da perda
    criterion = nn.CrossEntropyLoss(ignore_index=0) 

    for batch_idx, (inputs, targets) in enumerate(dataloader):
        # 1. Mover dados para o dispositivo (GPU/TPU)
        inputs = inputs.to(device)
        targets = targets.to(device)
        
        # 2. Zerar os gradientes acumulados
        # Essencial: PyTorch acumula gradientes por padrão.
        optimizer.zero_grad(set_to_none=True) 
        
        # 3. Forward Pass
        # O modelo retorna logits de forma (Batch, Seq_Len, Vocab_Size)
        logits = model(inputs)
        
        # 4. Reshape para cálculo da Loss
        # CrossEntropy espera (N, C) ou (N, C, d1...)
        # Achatamos Batch e Seq_Len para tratar cada token como uma instância
        B, T, C = logits.shape
        logits = logits.view(B * T, C)
        targets = targets.view(B * T)
        
        # 5. Cálculo da Perda
        loss = criterion(logits, targets)
        
        # 6. Backward Pass (Backpropagation)
        # Calcula dLoss/dWeights para todos os parâmetros
        loss.backward()
        
        # 7. Gradient Clipping
        # Crucial em Transformers para evitar explosão de gradientes
        torch.nn.utils.clip_grad_norm_(model.parameters(), clip_grad)
        
        # 8. Passo do Otimizador
        # Atualiza os pesos: w = w - lr * grad
        optimizer.step()
        
        total_loss += loss.item()
        
        # (Opcional) Logging intra-época
        if batch_idx % 100 == 0:
            print(f"Batch {batch_idx} | Loss: {loss.item():.4f}")

    return total_loss / len(dataloader)

# Exemplo de Configuração do Otimizador
# model = MyTransformer(...)
# optimizer = AdamW(model.parameters(), lr=3e-4, betas=(0.9, 0.95), weight_decay=0.1)

11.5 5. Detalhes Críticos de Arquitetura

11.5.1 Gradient Accumulation (Acumulação de Gradientes)

Se a memória da GPU não suportar o tamanho de batch desejado (ex: 512), usamos acumulação de gradientes. Executamos o forward e backward várias vezes com micro-batches menores, mas só chamamos optimizer.step() após \(N\) passos.

11.5.2 Gradient Clipping

Transformers são profundos e sensíveis. Gradientes podem crescer exponencialmente (exploding gradients), desestabilizando o treino. O clip_grad_norm_ reescala o vetor de gradientes global se sua norma exceder um limiar (geralmente 1.0), mantendo a direção mas limitando a magnitude.

11.5.3 set_to_none=True

Ao zerar gradientes, usar optimizer.zero_grad(set_to_none=True) é ligeiramente mais eficiente em termos de memória do que definir os tensores como zero, pois evita operações de leitura/escrita de memória desnecessárias.

11.6 Resumo do Capítulo

O sucesso do treinamento de um LLM depende do equilíbrio delicado entre a arquitetura do modelo e o algoritmo de otimização. Utilizamos Cross-Entropy para guiar o modelo probabilisticamente e AdamW para navegar na superfície de erro complexa de forma eficiente, desacoplando a regularização da adaptação do passo. A implementação correta do loop, incluindo salvaguardas como Gradient Clipping, é mandatória para a convergência do modelo.

11.7 Instrumentação e depuração

Registre perda, taxa de aprendizado, norma dos gradientes, tokens processados por segundo e uso de memória. Uma perda que se torna NaN pode indicar taxa alta, overflow em precisão reduzida, dados inválidos ou máscara incorreta. A norma dos gradientes ajuda a distinguir explosão de gradiente de um problema na entrada.

Um batch efetivo com acumulação é micro_batch × passos_de_acumulação × número_de_dispositivos. Divida a perda pelos passos de acumulação antes de backward() para manter a escala do gradiente equivalente à média do lote maior.

WarningOrdem das operações

Com mixed precision, desfaça a escala dos gradientes antes do clipping. Só então execute o passo do otimizador, atualize o scaler e zere os gradientes.

11.8 Escala de modelo, dados e computação

Resultados empíricos mostram relações previsíveis entre perda, quantidade de dados, parâmetros e computação, mas aumentar apenas uma dimensão pode desperdiçar recursos. O planejamento deve equilibrar capacidade do modelo e quantidade de tokens, além de considerar inferência, energia e custo de avaliação. Capacidades observadas em escala precisam ser medidas diretamente; não devem ser presumidas apenas pela contagem de parâmetros.

11.9 Overfitting, underfitting e regularização

Underfitting ocorre quando o modelo ou o procedimento de treino não consegue reduzir adequadamente nem a perda de treinamento. Pode indicar capacidade insuficiente, poucos passos, taxa de aprendizado inadequada ou dados difíceis. Overfitting aparece quando o desempenho no treino melhora, mas a validação deixa de acompanhar ou piora.

Regularização não é uma única técnica. Weight decay desencoraja pesos excessivos; dropout perturba ativações durante o treino; parada antecipada escolhe um checkpoint anterior à degradação de validação; diversidade de dados reduz a dependência de padrões estreitos. Em LLMs, deduplicação e controle da proporção entre fontes também funcionam como regularização do corpus.

loss de treino:      3.2 -> 2.4 -> 1.8 -> 1.3
loss de validação:   3.3 -> 2.6 -> 2.5 -> 2.8
                                      ^ melhor checkpoint
WarningTreinar por mais tempo pode piorar

Uma loss de treino menor não compensa perda de generalização. Critérios de parada e escolha de checkpoint devem ser definidos antes de observar o conjunto de teste.

11.10 Ajuste de hiperparâmetros

Pesquisa em grade avalia combinações predeterminadas, mas cresce exponencialmente com o número de dimensões. Pesquisa aleatória costuma explorar melhor quando poucos hiperparâmetros dominam o resultado. Métodos bayesianos usam resultados anteriores para sugerir novas tentativas, sendo úteis quando cada treinamento é caro.

Comece por faixas logarítmicas para taxa de aprendizado e weight decay. Faça ensaios curtos, elimine configurações instáveis e só então invista em execuções longas. Compare sob o mesmo orçamento de tokens e registre configuração, código, semente e hardware.

11.11 Leis de escala como ferramenta de planejamento

Curvas de escala ajustam relações empíricas entre perda e recursos. Elas ajudam a estimar se o próximo incremento deve ir para dados, parâmetros ou passos, mas são válidas no regime medido. Mudanças de arquitetura, qualidade de dados e objetivo podem deslocar a curva. Extrapolação é uma hipótese quantitativa, não uma promessa de novas capacidades.