O otimizador LAMB (Layer-wise Adaptive Moments for Batch training) é um algoritmo de otimização estocástica para treinar redes neurais profundas que estende o otimizador Adam com uma etapa de normalização por camada. Introduzido em 2019 por pesquisadores do Google (especificamente por Yang You, Jing Li, Jonathan Hseu, entre outros), o LAMB foi projetado para permitir o uso eficiente de tamanhos de mini-lote muito grandes (por exemplo, 32.768 ou mais) sem degradar a precisão do modelo ou exigir ajuste exaustivo de hiperparâmetros. Ele alcança isso escalando a magnitude da atualização para cada camada com base na razão entre a norma dos pesos da camada e a norma de sua atualização, efetivamente desacoplando a taxa de aprendizado da escala dos parâmetros de cada camada.
O LAMB tem sido particularmente influente no treinamento de modelos baseados em Transformer (architecture), incluindo os primeiros modelos de linguagem de grande escala (LLMs) e arquiteturas de visão como ResNet. Sua adoção em estruturas de treinamento distribuído, como TensorFlow (via tf.keras.optimizers.LAMB) e PyTorch (via a implementação do LAMB em bibliotecas como NVIDIA, Hugging Face e FairScale), tornou-o uma ferramenta padrão para escalar execuções de treinamento. Ao permitir lotes maiores, o LAMB reduz o tempo de relógio necessário para treinar modelos de última geração, o que é crítico para organizações como OpenAI, Anthropic e Google DeepMind que dependem de clusters de computação massivos.
Motivação e Contexto
Treinar redes neurais profundas com mini-lotes grandes reduz o número de atualizações de parâmetros por época e etapas computacionais, mas o escalonamento ingênuo do tamanho do lote frequentemente leva a uma generalização ruim e convergência instável. Isso é conhecido como o "problema do treinamento com lote grande". Otimizadores padrão como a descida do gradiente estocástica (SGD) com momentum ou Adam exigem ajuste cuidadoso da taxa de aprendizado quando o tamanho do lote aumenta, e mesmo assim, a precisão frequentemente degrada. O otimizador LAMB foi desenvolvido para resolver isso tornando o otimizador mais robusto a mudanças no tamanho do lote.
A percepção chave por trás do LAMB é que diferentes camadas em uma rede profunda exibem escalas vastamente diferentes de normas de gradientes e pesos. Por exemplo, camadas convolucionais iniciais têm pesos pequenos, enquanto camadas totalmente conectadas posteriores têm pesos grandes. Uma única taxa de aprendizado global no Adam pode levar a atualizações muito grandes para algumas camadas (causando divergência) ou muito pequenas para outras (desacelerando a convergência). O LAMB introduz uma taxa adaptativa por camada que normaliza a atualização com base na razão entre as normas dos pesos e gradientes da camada, garantindo que cada camada se mova a um ritmo estável relativo à sua magnitude.
Descrição do Algoritmo
O LAMB pode ser visto como uma variante do Adam com uma etapa adicional de normalização. Seja \(\theta_t\) os parâmetros na iteração \(t\), e \(g_t\) o gradiente da perda em relação a \(\theta_t\). O LAMB mantém os primeiros e segundos momentos (\(m_t\) e \(v_t\)) dos gradientes, semelhante ao Adam, com taxas de decaimento exponencial \(\beta_1\) e \(\beta_2\) (tipicamente 0,9 e 0,999). Após a correção de viés, ele calcula uma atualização Adam \(\frac{m_t}{\sqrt{v_t} + \epsilon}\).
A diferença crucial é a razão de confiança \(\phi\): para cada camada \(i\), \(\phi_i = \frac{||\theta_{t,i}||}{||r_{t,i}||}\), onde \(r_{t,i} = \frac{m_{t,i}}{\sqrt{v_{t,i}} + \epsilon}\) é a atualização Adam (sem taxa de aprendizado) para essa camada, e \(||\cdot||\) denota a norma L2. A atualização final para a camada \(i\) é \(\theta_{t+1,i} = \theta_{t,i} - \eta \cdot \phi_i \cdot r_{t,i}\), onde \(\eta\) é a taxa de aprendizado global. Essa razão de confiança escala a atualização proporcionalmente à norma dos pesos da camada, então camadas pequenas recebem atualizações pequenas (mas não desprezíveis), e camadas grandes recebem atualizações maiores, porém estáveis.
Na prática, uma pequena constante (por exemplo, 1e-6) é adicionada ao denominador para evitar divisão por zero. O algoritmo também inclui decaimento de peso opcional (regularização L2) fundido na atualização, seguindo a abordagem de decaimento de peso desacoplado usada no AdamW. Quando a razão de confiança é definida como 1 para todas as camadas, o LAMB reduz-se ao Adam padrão (com correção de viés).
Hiperparâmetros e Ajuste
O LAMB herda a maioria dos hiperparâmetros do Adam: \(\beta_1\) (momentum), \(\beta_2\) (decaimento de variância), \(\epsilon\) (estabilidade numérica) e taxa de decaimento de peso. O principal novo hiperparâmetro é a taxa de aprendizado global \(\eta\), que é frequentemente definida na faixa de 0,01-0,1 para treinamento com lote grande, significativamente maior do que o típico para o Adam (por exemplo, 1e-3). Os autores descobriram que para lotes muito grandes (por exemplo, 32.768 para BERT), uma taxa de aprendizado de 0,01 com aquecimento linear sobre os primeiros 10% das etapas funciona bem, e eles também recomendaram usar um decaimento de cosseno esquema de taxa de aprendizado sobre as etapas restantes.
Além disso, a escolha de \(\beta_2\) pode afetar a estabilidade; para modelos com gradientes esparsos, um \(\beta_2\) mais alto (por exemplo, 0,99) pode ser usado. Os autores também sugeriram que o tamanho do lote pode ser escalado proporcionalmente com a taxa de aprendizado (regra de escalonamento linear), uma diretriz que funciona bem com o LAMB. Por exemplo, se o tamanho do lote for dobrado, a taxa de aprendizado também pode ser dobrada sem perda de precisão.
Desempenho e Benchmarks
No artigo original, o LAMB foi avaliado em duas tarefas principais: treinar ResNet-50 no ImageNet (classificação de imagens) e BERT (um modelo de linguagem baseado em transformer) para modelagem de linguagem mascarada. Usando o LAMB, os autores alcançaram precisão top-1 no ImageNet de 76,0% em apenas 2.048 iterações com um tamanho de lote de 32.768, igualando a precisão de última geração alcançada com lotes menores (por exemplo, 256) em muito menos épocas. Para o BERT, eles treinaram o modelo até a mesma precisão (por exemplo, pontuação F1 de 1,0 no SQUAD) em cerca de 3,5 minutos usando 1.024 TPUs, uma aceleração de 10x em relação a métodos anteriores.
Posteriormente, o LAMB tornou-se o otimizador padrão para treinar modelos baseados em BERT nos fluxos de trabalho internos do Google. O artigo relatou que o LAMB superou tanto o Adam quanto o SGD com momentum ao escalar tamanhos de lote de 1.024 a 65.536. Os autores também mostraram que o LAMB funciona bem com recorte de gradiente (usado para prevenir gradientes explosivos) e é compatível com treinamento de precisão mista, como usado em hardware moderno como GPUs NVIDIA e TPUs Google.
Aplicações em Treinamento em Grande Escala
A principal aplicação do LAMB é no treinamento distribuído, onde o tamanho do lote é grande demais para caber na memória de um único dispositivo. Em tais configurações, os gradientes são calculados em média em múltiplas GPUs ou TPUs usando paralelismo de dados. Por exemplo, OpenAI e Google DeepMind usam otimizadores análogos ao LAMB ao treinar grandes modelos Transformer (architecture) com comprimentos de sequência na casa dos milhares. Embora otimizadores mais novos como o LAMB (e seu sucessor LAMB2) tenham sido propostos, o LAMB continua sendo uma escolha confiável em muitos esforços de código aberto, incluindo o treinamento de transformers de visão e LLMs por grupos de pesquisa e empresas como AI21 Labs e SambaNova.
No contexto de aprendizado de máquina em Amazon Web Services (com hardware AWS Trainium), o LAMB é suportado em kernels personalizados para eficiência. Da mesma forma, Intel e AMD avaliaram o LAMB em seus aceleradores. A capacidade do otimizador de lidar com tamanhos de lote extremos o torna valioso para pré-treinar modelos em conjuntos de dados massivos, onde o custo de uma única época é alto, e reduzir épocas é fundamental.
Relação com Outros Otimizadores
O LAMB faz parte de uma família de otimizadores adaptativos que inclui variantes de SGD, Adam e seus sucessores como AdamW (decaimento de peso desacoplado) e LARS (Layer-wise Adaptive Rate Scaling). O LARS, introduzido por You et al. em 2017 para treinamento com lote grande de CNNs, usa uma razão de confiança por camada semelhante, mas não mantém segundos momentos; ele depende de primeiros momentos (momentum) e normas de gradiente. O LAMB combina os benefícios do LARS (escalonamento por camada) com as taxas de aprendizado adaptativas por parâmetro do Adam, tornando-o mais robusto para modelos com gradientes esparsos (como transformers).
Outro otimizador intimamente relacionado é o NVLAMB (da Nvidia), que incorpora redução de variância. No entanto, o LAMB permanece mais simples e amplamente usado. Para tarefas sequência a sequência com busca em feixe, o LAMB não afeta diretamente a inferência, mas ajuda na convergência do treinamento, o que indiretamente melhora a decodificação de sequências.
Extensões e Variantes
Desde sua introdução, várias variantes foram propostas. LAMB2 (também do Google) adiciona um fator de normalização baseado na variância do gradiente, melhorando a estabilidade para certos problemas. LARS com correção de viés e outras modificações também são comuns. Na prática, muitas estruturas implementam o LAMB com correção de viés opcional para os momentos, o que é benéfico durante as primeiras etapas. Algumas implementações, como em PyTorch's torch.optim.Lamb (no pacote torch_optimizer), permitem ajustar o parâmetro da razão de confiança ou usar uma taxa de aprendizado personalizada por camada.
Apesar do surgimento de novos otimizadores como AdamW com diferentes estratégias de escalonamento (por exemplo, esquemas de 1 ciclo), o LAMB continua sendo uma linha de base forte para treinamento com lote grande. A comunidade de pesquisa explorou combinar o LAMB com aumento de dados e recorte de gradiente para melhorar ainda mais a generalização.
Considerações Práticas e Limitações
Embora o LAMB se destaque em configurações de lote grande, ele nem sempre é a melhor escolha para tamanhos de lote pequenos (por exemplo, abaixo de 1.024). Em tais regimes, o Adam padrão ou SGD com momentum podem ser mais simples e igualmente eficazes. O LAMB adiciona uma sobrecarga computacional de calcular normas por camada, que é desprezível em hardware moderno, mas pode ser não trivial para modelos com muitas camadas pequenas (por exemplo, arquiteturas U-Net).
Outra limitação é que a razão de confiança do LAMB pode ocasionalmente levar a treinamento instável se algumas camadas tiverem normas de pesos muito pequenas (próximas de zero). Isso é tipicamente mitigado adicionando um termo épsilon ao denominador e usando decaimento de peso, que impede que os pesos se desviem para zero. Além disso, o LAMB requer ajuste cuidadoso da taxa de aprendizado e etapas de aquecimento; um esquema inadequado pode levar à divergência.
O uso de memória é semelhante ao Adam (dois vetores de momento por parâmetro), então não é mais faminto por memória. Para modelos muito grandes, poda de modelo ou recorte de gradiente podem ser usados junto com o LAMB, mas essas são técnicas ortogonais.
Conclusão
O LAMB tornou-se uma pedra angular na caixa de ferramentas de algoritmos de otimização para aprendizado profundo em grande escala. Ao permitir treinamento eficaz com mini-lotes massivos, ele acelerou o desenvolvimento de muitos modelos de referência e reduziu o custo da experimentação. Seu princípio de adaptação por camada influenciou designs subsequentes de otimizadores e continua sendo uma solução prática e bem compreendida para profissionais que enfrentam os desafios do treinamento distribuído. À medida que a inteligência artificial continua a crescer, otimizadores como o LAMB provavelmente evoluirão, mas suas ideias centrais de confiança por camada e momentos adaptativos vieram para ficar.