El optimizador LAMB (Layer-wise Adaptive Moments para entrenamiento por lotes) es un algoritmo de optimización estocástica para entrenar redes neuronales profundas que extiende el optimizador Adam con un paso de normalización por capa. Introducido en 2019 por investigadores de Google (específicamente por Yang You, Jing Li, Jonathan Hseu, et al.), LAMB fue diseñado para permitir el uso eficiente de tamaños de mini-lote muy grandes (por ejemplo, 32,768 o más) sin degradar la precisión del modelo ni requerir un ajuste exhaustivo de hiperparámetros. Logra esto escalando la magnitud de la actualización para cada capa basándose en la relación entre la norma de los pesos de la capa y la norma de su actualización, desacoplando efectivamente la tasa de aprendizaje de la escala de los parámetros de cada capa.
LAMB ha sido particularmente influyente en el entrenamiento de modelos basados en Transformer (architecture), incluidos los primeros modelos de lenguaje grandes (LLMs) y arquitecturas de visión como ResNet. Su adopción en marcos de entrenamiento distribuido, como TensorFlow (a través de tf.keras.optimizers.LAMB) y PyTorch (a través de la implementación de LAMB en bibliotecas como NVIDIA, Hugging Face y FairScale), lo convirtió en una herramienta estándar para escalar ejecuciones de entrenamiento. Al permitir lotes más grandes, LAMB reduce el tiempo de reloj necesario para entrenar modelos de última generación, lo cual es crítico para organizaciones como OpenAI, Anthropic y Google DeepMind que dependen de clústeres de cómputo masivos.
Motivación y Antecedentes
Entrenar redes neuronales profundas con mini-lotes grandes reduce el número de actualizaciones de parámetros por época y los pasos computacionales, pero escalar ingenuamente el tamaño del lote a menudo conduce a una mala generalización y a una convergencia inestable. Esto se conoce como el "problema del entrenamiento con lotes grandes". Los optimizadores estándar como el descenso de gradiente estocástico (SGD) con momento o el Adam requieren un ajuste cuidadoso de la tasa de aprendizaje cuando aumenta el tamaño del lote, e incluso así, la precisión a menudo se degrada. El optimizador LAMB fue desarrollado para abordar esto haciendo que el optimizador sea más robusto a los cambios en el tamaño del lote.
La idea clave detrás de LAMB es que diferentes capas en una red profunda exhiben escalas muy diferentes de gradientes y normas de pesos. Por ejemplo, las capas convolucionales tempranas tienen pesos pequeños, mientras que las capas totalmente conectadas posteriores tienen pesos grandes. Una única tasa de aprendizaje global en Adam puede llevar a actualizaciones demasiado grandes para algunas capas (causando divergencia) o demasiado pequeñas para otras (ralentizando la convergencia). LAMB introduce una tasa adaptativa por capa que normaliza la actualización basándose en la relación entre las normas de los pesos y gradientes de la capa, asegurando que cada capa se mueva a un ritmo estable relativo a su magnitud.
Descripción del Algoritmo
LAMB puede verse como una variante de Adam con un paso de normalización adicional. Sea \(\theta_t\) los parámetros en la iteración \(t\), y \(g_t\) el gradiente de la pérdida con respecto a \(\theta_t\). LAMB mantiene los primeros y segundos momentos (\(m_t\) y \(v_t\)) de los gradientes, similar a Adam, con tasas de decaimiento exponencial \(\beta_1\) y \(\beta_2\) (típicamente 0.9 y 0.999). Después de la corrección de sesgo, calcula una actualización de Adam \(\frac{m_t}{\sqrt{v_t} + \epsilon}\).
La diferencia crucial es la relación de confianza \(\phi\): para cada capa \(i\), \(\phi_i = \frac{||\theta_{t,i}||}{||r_{t,i}||}\), donde \(r_{t,i} = \frac{m_{t,i}}{\sqrt{v_{t,i}} + \epsilon}\) es la actualización de Adam (sin tasa de aprendizaje) para esa capa, y \(||\cdot||\) denota la norma L2. La actualización final para la capa \(i\) es \(\theta_{t+1,i} = \theta_{t,i} - \eta \cdot \phi_i \cdot r_{t,i}\), donde \(\eta\) es la tasa de aprendizaje global. Esta relación de confianza escala la actualización proporcionalmente a la norma de los pesos de la capa, de modo que las capas pequeñas reciben actualizaciones pequeñas (pero no despreciables), y las capas grandes reciben actualizaciones más grandes, pero estables.
En la práctica, se añade una pequeña constante (por ejemplo, 1e-6) al denominador para evitar la división por cero. El algoritmo también incluye un decaimiento de pesos opcional (regularización L2) fusionado en la actualización, siguiendo el enfoque de decaimiento de pesos desacoplado utilizado en AdamW. Cuando la relación de confianza se establece en 1 para todas las capas, LAMB se reduce al Adam estándar (con corrección de sesgo).
Hiperparámetros y Ajuste
LAMB hereda la mayoría de los hiperparámetros de Adam: \(\beta_1\) (momento), \(\beta_2\) (decaimiento de varianza), \(\epsilon\) (estabilidad numérica) y la tasa de decaimiento de pesos. El nuevo hiperparámetro principal es la tasa de aprendizaje global \(\eta\), que a menudo se establece en el rango de 0.01-0.1 para entrenamiento con lotes grandes, significativamente más alta que la típica para Adam (por ejemplo, 1e-3). Los autores encontraron que para lotes muy grandes (por ejemplo, 32,768 para BERT), una tasa de aprendizaje de 0.01 con calentamiento lineal durante el primer 10% de los pasos funciona bien, y también recomendaron usar un decaimiento de coseno programa de tasa de aprendizaje en los pasos restantes.
Además, la elección de \(\beta_2\) puede afectar la estabilidad; para modelos con gradientes dispersos, se puede usar un \(\beta_2\) más alto (por ejemplo, 0.99). Los autores también sugirieron que el tamaño del lote puede escalarse proporcionalmente con la tasa de aprendizaje (regla de escalado lineal), una guía que funciona bien con LAMB. Por ejemplo, si el tamaño del lote se duplica, la tasa de aprendizaje también puede duplicarse sin pérdida de precisión.
Rendimiento y Evaluaciones Comparativas
En el artículo original, LAMB fue evaluado en dos tareas principales: entrenar ResNet-50 en ImageNet (clasificación de imágenes) y BERT (un modelo de lenguaje basado en transformer) para modelado de lenguaje enmascarado. Usando LAMB, los autores lograron una precisión top-1 en ImageNet del 76.0% en solo 2,048 iteraciones con un tamaño de lote de 32,768, igualando la precisión de última generación lograda con lotes más pequeños (por ejemplo, 256) en muchas menos épocas. Para BERT, entrenaron el modelo hasta la misma precisión (por ejemplo, puntuación F1 de 1.0 en SQUAD) en aproximadamente 3.5 minutos usando 1,024 TPUs, una aceleración de 10x sobre métodos anteriores.
Posteriormente, LAMB se convirtió en el optimizador predeterminado para entrenar modelos basados en BERT en los flujos de trabajo internos de Google. El artículo informó que LAMB superó tanto a Adam como a SGD con momento al escalar tamaños de lote de 1,024 a 65,536. Los autores también mostraron que LAMB funciona bien con recorte de gradientes (utilizado para prevenir gradientes explosivos) y es compatible con el entrenamiento de precisión mixta, como se usa en hardware moderno como GPUs NVIDIA y TPUs de Google.
Aplicaciones en Entrenamiento a Gran Escala
La aplicación principal de LAMB es en el entrenamiento distribuido donde el tamaño del lote es demasiado grande para caber en la memoria de un solo dispositivo. En tales configuraciones, los gradientes se promedian entre múltiples GPUs o TPUs usando paralelismo de datos. Por ejemplo, OpenAI y Google DeepMind usan optimizadores análogos a LAMB al entrenar grandes modelos Transformer (architecture) con longitudes de secuencia en los miles. Aunque se han propuesto optimizadores más nuevos como LAMB (y su sucesor LAMB2), LAMB sigue siendo una opción confiable en muchos esfuerzos de código abierto, incluido el entrenamiento de transformers de visión y LLMs por grupos de investigación y empresas como AI21 Labs y SambaNova.
En el contexto del aprendizaje automático en Amazon Web Services (con hardware AWS Trainium), LAMB está soportado en kernels personalizados para mayor eficiencia. De manera similar, Intel y AMD han evaluado LAMB en sus aceleradores. La capacidad del optimizador para manejar tamaños de lote extremos lo hace valioso para preentrenar modelos en conjuntos de datos masivos, donde el costo de una sola época es alto y reducir épocas es primordial.
Relación con Otros Optimizadores
LAMB es parte de una familia de optimizadores adaptativos que incluyen variantes de SGD, Adam y sus sucesores como AdamW (decaimiento de pesos desacoplado) y LARS (Escalado de Tasa Adaptativa por Capa). LARS, introducido por You et al. en 2017 para entrenamiento con lotes grandes de CNNs, usa una relación de confianza por capa similar pero no mantiene segundos momentos; depende de primeros momentos (momento) y normas de gradientes. LAMB combina los beneficios de LARS (escalado por capa) con las tasas de aprendizaje adaptativas por parámetro de Adam, haciéndolo más robusto para modelos con gradientes dispersos (como transformers).
Otro optimizador estrechamente relacionado es NVLAMB (de Nvidia), que incorpora reducción de varianza. Sin embargo, LAMB sigue siendo más simple y ampliamente utilizado. Para tareas secuencia a secuencia con búsqueda de haz, LAMB no afecta directamente la inferencia, pero ayuda a la convergencia del entrenamiento, lo que indirectamente mejora la decodificación de secuencias.
Extensiones y Variantes
Desde su introducción, se han propuesto varias variantes. LAMB2 (también de Google) añade un factor de normalización basado en la varianza del gradiente, mejorando la estabilidad para ciertos problemas. LARS con corrección de sesgo y otras modificaciones también son comunes. En la práctica, muchos marcos implementan LAMB con corrección de sesgo opcional para los momentos, lo cual es beneficioso durante los primeros pasos. Algunas implementaciones, como en PyTorch's torch.optim.Lamb (en el paquete torch_optimizer), permiten ajustar el parámetro de relación de confianza o usar una tasa de aprendizaje personalizada por capa.
A pesar de la aparición de nuevos optimizadores como AdamW con diferentes estrategias de escalado (por ejemplo, programas de 1 ciclo), LAMB sigue siendo una línea base sólida para el entrenamiento con lotes grandes. La comunidad de investigación ha explorado combinar LAMB con aumento de datos y recorte de gradientes para mejorar aún más la generalización.
Consideraciones Prácticas y Limitaciones
Aunque LAMB sobresale en configuraciones de lotes grandes, no siempre es la mejor opción para tamaños de lote pequeños (por ejemplo, por debajo de 1,024). En tales regímenes, Adam estándar o SGD con momento pueden ser más simples e igualmente efectivos. LAMB añade una sobrecarga computacional de calcular normas por capa, que es insignificante en hardware moderno pero puede ser no trivial para modelos con muchas capas pequeñas (por ejemplo, arquitecturas U-Net).
Otra limitación es que la relación de confianza de LAMB puede ocasionalmente conducir a un entrenamiento inestable si algunas capas tienen normas de pesos muy pequeñas (cerca de cero). Esto se mitiga típicamente añadiendo un término epsilon al denominador y usando decaimiento de pesos, que evita que los pesos se desvíen a cero. Además, LAMB requiere un ajuste cuidadoso de la tasa de aprendizaje y los pasos de calentamiento; un programa inapropiado puede llevar a divergencia.
El uso de memoria es similar a Adam (dos vectores de momento por parámetro), por lo que no consume más memoria. Para modelos muy grandes, poda de modelos o recorte de gradientes pueden usarse junto con LAMB, pero estas son técnicas ortogonales.
Conclusión
LAMB se ha convertido en una piedra angular en la caja de herramientas de algoritmos de optimización para el aprendizaje profundo a gran escala. Al permitir un entrenamiento efectivo con mini-lotes masivos, ha acelerado el desarrollo de muchos modelos de referencia y reducido el costo de la experimentación. Su principio de adaptación por capa ha influido en diseños de optimizadores posteriores y sigue siendo una solución práctica y bien comprendida para los profesionales que enfrentan los desafíos del entrenamiento distribuido. A medida que la inteligencia artificial continúa creciendo, optimizadores como LAMB probablemente evolucionarán, pero sus ideas centrales de confianza por capa y momentos adaptativos están aquí para quedarse.