Resumen Técnico: ZO-Act

Planteamiento del Problema

El ajuste fino (fine-tuning) de Modelos de Lenguaje de Gran Escala (LLMs) depende típicamente de la optimización de primer orden (retropropagación o backpropagation), la cual se vuelve prohibitiva en términos de memoria a medida que los modelos escalan a miles de millones de parámetros debido a los requisitos de almacenamiento para activaciones, estados del optimizador y cálculos de gradientes. La optimización de orden cero (ZO) ofrece una alternativa de solo paso hacia adelante (forward-only) al estimar las direcciones de actualización utilizando únicamente evaluaciones de pérdida, evitando así la retropropagación. Sin embargo, los métodos ZO existentes enfrentan desafíos significativos:

1. Alta Varianza: Perturbar el espacio completo de parámetros (como en MeZO) resulta en estimaciones de gradiente con alta varianza, lo que conduce a una convergencia lenta e inestabilidad.
2. Subespacios Subóptimos: Aunque algunos métodos restringen las perturbaciones a subespacios de baja dimensión (por ejemplo, LOZO, SubZero), a menudo dependen de subespacios construidos aleatoriamente o perturbaciones aleatorias que pueden no alinearse con el flujo real del gradiente.
3. Compatibilidad con Optimizadores: Muchos métodos ZO tienen dificultades para utilizar eficazmente optimizadores estándar basados en momentum (como Adam), debido a que la estructura de perturbación no expone variables entrenables explícitas y ligeras.
4. Restricciones de Cuantización: El ajuste fino de LLMs cuantizados (por ejemplo, INT4) es difícil con métodos que requieren la materialización de perturbaciones de peso completo.

Metodología: ZO-Act

Los autores proponen ZO-Act, un método de ajuste fino ZO de un solo paso (one-shot) informado por la activación. La idea central es que los componentes dominantes de los gradientes de peso en las capas lineales residen dentro del subespacio generado por los vectores singulares derechos superiores de la matriz de activación de entrada.

Mecanismo Central

1. Inicialización de un Solo Paso (One-Shot): Para cada capa lineal, ZO-Act realiza un único paso hacia adelante en un lote de calibración para calcular la matriz de activación de entrada $X$. Posteriormente, calcula los $r$ vectores singulares derechos superiores ($V_r$) de $X$ mediante una SVD (Descomposición en Valores Singulares) delgada.
2. Parametrización de Subespacio Fijo: La actualización de peso se restringe a un subespacio de bajo rango definido por estos vectores de activación congelados. La actualización de peso efectiva se parametriza como:
   $$\Delta W = V_r B$$
   donde $V_r \in \mathbb{R}^{m \times r}$ es la base de activación congelada, y $B \in \mathbb{R}^{r \times n}$ es una matriz de coeficientes ligera y entrenable. Los pesos preentrenados originales $W$ y la base $V_r$ permanecen congelados durante todo el entrenamiento.
3. Optimización en el Espacio de Coeficientes: En lugar de perturbar la matriz de peso completa $W$, ZO-Act toma muestras de perturbaciones aleatorias $Z$ en el espacio de coeficientes de baja dimensión ($Z \sim \mathcal{N}(0, I)$). El peso efectivo perturbado se convierte en $W_{eff} = W + V_r(B + \mu Z)$.
4. Estimación de Solo Paso Hacia Adelante (Forward-Only): El gradiente para $B$ se estima utilizando evaluaciones de pérdida de diferencia finita hacia adelante:
   $$\hat{g}_B = \frac{L(B + \mu Z) - L(B)}{\mu} Z$$
   Este gradiente estimado se utiliza para actualizar $B$ empleando optimizadores estándar de primer orden como Adam.

Análisis Teórico

El artículo analiza ZO-Act como una optimización de orden cero sobre un espacio de coeficientes restringido $\beta$ (dimensión $k$) en lugar del espacio de pesos completo (dimensión $d$).

• Reducción de Varianza: La varianza del estimador ZO escala linealmente con la dimensión de la perturbación. Al reducir la dimensión de $d$ (pesos completos) a $k = \sum r n_\ell$ (coeficientes, donde $r \ll m$), ZO-Act reduce significativamente el término de convergencia dependiente de la varianza y el error de diferencia finita.
• Compromiso Sesgo-Varianza (Bias-Variance Trade-off): Restringir las actualizaciones al subespacio informado por la activación introduce un sesgo de aproximación de subespacio. Sin embargo, los autores argumentan que las activaciones y los gradientes de los LLM exhiben estructuras de bajo rango, lo que significa que las direcciones de actualización dominantes son capturadas adecuadamente por los vectores singulares superiores de las activaciones de entrada. Por lo tanto, el sesgo está controlado y mitigado por la naturaleza informada de los datos del subespacio.

Contribuciones Clave

1. Subespacio Informado por la Activación: A diferencia de los métodos de subespacio previos que utilizan matrices de bajo rango aleatorias o estáticas, ZO-Act deriva el subespacio de perturbación directamente de las activaciones de entrada, asegurando que el subespacio se alinee con las direcciones de activación dominantes de la capa.
2. Optimización de Coeficientes Explícita: Al congelar la base y optimizar únicamente la matriz de coeficientes $B$, ZO-Act transforma el ajuste fino ZO en un problema de optimización explícito sobre un espacio de baja dimensión. Esto permite la aplicación directa de optimizadores basados en momentum (por ejemplo, Adam) y evita la materialización de perturbaciones de peso completo.
3. Compatibilidad con Cuantización: El método soporta naturalmente LLMs cuantizados (por ejemplo, INT4) porque los pesos de bajos bits permanecen congelados y solo se actualizan las matrices de coeficientes de bajos bits.
4. Garantías Teóricas: El artículo proporciona un análisis de convergencia que muestra cómo ZO-Act reduce tanto la varianza del estimador de gradiente como el error de diferencia finita en comparación con el ZO de peso completo, a costa de un sesgo de subespacio manejable.

Resultados Experimentales

Los autores evaluaron ZO-Act en Llama-3-8B, OPT-13B y un Llama-3-8B cuantizado en INT4 a través de tareas de comprensión de lenguaje, respuesta a preguntas y razonamiento de sentido común.

• Rendimiento: ZO-Act superó consistentemente a los fuertes baselines de ZO (incluyendo MeZO, LOZO, SubZero, AGZO y ZO-Muon).
  • En Llama-3-8B, logró los mejores resultados de ZO de precisión completa en SST-2, RTE, BoolQ y WiC. Notablemente, mejoró el rendimiento de RTE de 81.2 (ZO-Muon) a 87.0 y de CB de 69.6 a 89.3.
  • En OPT-13B, alcanzó resultados similares de primer nivel en SST-2, RTE, CB y BoolQ.
  • En Modelos Cuantizados en INT4, ZO-Act se mantuvo competitivo con los baselines de ZO de precisión completa, demostrando su eficacia para el ajuste fino con restricciones de memoria.
• Eficiencia: ZO-Act igualó el uso de memoria más bajo entre los métodos comparados (aproximadamente 16.1 GB para Llama-3-8B, cayendo a 5.8 GB para INT4) y logró el tiempo de ejecución más rápido (42.0 minutos frente a 44.5+ minutos para otros) bajo un presupuesto fijo de consultas de paso hacia adelante (forward-query budget).
• Sensibilidad al Rango: Los experimentos mostraron que un rango de $r=1$ a menudo era suficiente para modelos de precisión completa, obteniendo el mejor rendimiento de precisión mientras se minimizaba la dimensión de la perturbación. Rangos más altos ($r=32$) fueron beneficiosos para modelos cuantizados para compensar la reducción de capacidad.
• Estabilidad: Los estudios diagnósticos confirmaron que el subespacio informado por la activación permanece estable durante el ajuste fino y captura una parte significativa de la energía del gradiente completo, validando el diseño de inicialización de un solo paso (one-shot).

Significado y Reivindicaciones

El artículo afirma que ZO-Act representa un paso significativo para hacer que el ajuste fino de orden cero sea una alternativa práctica y de alto rendimiento para LLMs grandes y cuantizados. Su importancia radica en:

• Cerrar la Brecha: Reduce la brecha de rendimiento entre el ajuste fino de orden cero (ZO) y el de primer orden mediante la reducción de la varianza del estimador a través de subespacios informados por los datos.
• Habilitar Optimizadores: Resuelve la dificultad de aplicar optimizadores estándar como Adam a los métodos ZO al exponer variables entrenables explícitas y de baja dimensión.
• Escalabilidad: Proporciona una solución eficiente en memoria para el ajuste fino de modelos cuantizados, lo cual es crítico para el despliegue en hardware con recursos limitados.

Los autores reconocen las limitaciones, señalando que ZO-Act aún no cierra completamente la brecha de rendimiento con los métodos de primer orden debido al ruido inherente de la estimación de orden cero y el sesgo del subespacio. Además, ZO-Act sigue siendo más lento en tiempo de ejecución (wall-clock time) que los métodos basados en retropropagación debido a los múltiples pasos hacia adelante requeridos por cada actualización. Sin embargo, el método se posiciona como una opción superior para escenarios donde la retropropagación es prohibitiva en términos de memoria o no está disponible.