TAPIOCA: Why Task- Aware Pruning Improves OOD model Capability

Resumen Técnico: TAPIOCA – Por qué el Poda Consciente de la Tarea Mejora la Capacidad de los Modelos para Datos Fuera de Distribución

Enunciado del Problema
Trabajos recientes han demostrado que eliminar capas de Modelos de Lenguaje Grandes (LLM) durante la inferencia —específicamente mediante métodos de poda consciente de la tarea como TALE (Eliminación de Capas Consciente de la Tarea)— puede mejorar el rendimiento en tareas específicas sin reentrenamiento. Sin embargo, los mecanismos que impulsan estas ganancias permanecen poco claros. Una pregunta central es si estas mejoras provienen de beneficios generales de compresión, efectos de regularización o interacciones específicas entre las representaciones internas del modelo y la distribución de entrada. Además, se desconoce si estas ganancias son universales o específicas de ciertos tipos de desplazamientos de datos.

Metodología
Los autores investigan este fenómeno utilizando una combinación de experimentos sintéticos controlados y evaluaciones a gran escala en LLMs ajustados finamente.

1. Configuración de Regresión Controlada: El estudio comienza con un transformador de 12 capas entrenado desde cero en regresión lineal en contexto. La distribución de entrenamiento se define explícitamente (coeficientes muestreados de $U(-1, 1)$). Los autores evalúan el modelo tanto en datos dentro de la distribución (ID, misma distribución) como en datos fuera de la distribución (OOD, coeficientes muestreados de intervalos más amplios $U(-\sigma, \sigma)$ donde $\sigma > 1$). Aplican TALE, un algoritmo de poda codicioso que elimina capas para maximizar la precisión de validación en una distribución específica.
2. Evaluación a Gran Escala de LLMs: Los hallazgos se extienden a LLMs ajustados finamente (Llama 3.1 8B y GPT-OSS 120B) especializados en razonamiento matemático y generación de código. Aquí, la distribución "adaptada" es la tarea de ajuste fino (por ejemplo, MATH500), y las tareas OOD incluyen benchmarks generales como MMLU y BoolQ.
3. Análisis Geométrico y Causal: Para explicar por qué funciona la poda, los autores analizan la geometría de los estados ocultos. Miden las normas representacionales por capa y las distancias entre pares de tokens. También emplean sustitutos lineales locales para aproximar las transformaciones de capa y realizan intervenciones causales, como la inserción de mapas sustitutos inversos y el reescalado de actualizaciones residuales, para aislar el efecto de la distorsión geométrica.

Hallazgos Clave

• Ganancias Específicas de OOD: El hallazgo empírico principal es que la poda consciente de la tarea produce ganancias de rendimiento solo bajo desplazamiento de distribución. Cuando se evalúa en la distribución a la que el modelo está adaptado (ID), TALE falla en mejorar el rendimiento o lo degrada activamente al eliminar capas necesarias para una precisión óptima. Por el contrario, en datos OOD, TALE identifica y elimina consistentemente capas que mejoran la precisión.
• Desajuste de la Geometría de Representación: Los autores demuestran que las entradas OOD inducen una geometría de representación distorsionada en comparación con las entradas ID. Específicamente, las entradas OOD hacen que las normas de los estados ocultos por capa y las distancias entre pares se desvíen sistemáticamente de los perfiles observados en datos ID. En la configuración de regresión, las entradas OOD causan que las distancias se inflen a través de las capas; en LLMs, las trayectorias OOD divergen de las trayectorias ID en las capas medias a tardías.
• Mecanismo de Amplificación: El estudio identifica que ciertas capas actúan como "amplificadores dependientes de la distribución". En entradas ID, estas capas realizan refinamientos útiles. Sin embargo, en entradas OOD, las mismas capas exhiben una amplificación de alto ganancia y bajo rango, empujando las representaciones más lejos de la geometría adaptada a la tarea. TALE mejora el rendimiento al eliminar estas capas específicas, previniendo así la amplificación de la distorsión geométrica.
• Alineación, no solo Contracción: La poda no suprime simplemente las normas de activación. En experimentos donde la dirección del desplazamiento de distribución se invierte (entrenamiento en distribuciones amplias, prueba en distribuciones estrechas), la poda expande las distancias de representación para alinearse con el perfil ID. El mecanismo central es la realineación de las representaciones OOD hacia la geometría aprendida durante la adaptación, independientemente de si esto requiere contracción o expansión.
• Evidencia Causal: Las intervenciones causales confirman que la distorsión geométrica es el motor de la pérdida de rendimiento. Insertar un mapa sustituto inverso después de una capa podada (para deshacer el efecto geométrico de la capa) restaura el rendimiento OOD casi al nivel del modelo podado. De manera similar, reescalar la actualización residual de una capa seleccionada mejora continuamente la precisión a medida que se atenúa la contribución de la capa, descartando artefactos topológicos o redundancia simple como la única causa.

Significado y Afirmaciones
El artículo afirma proporcionar una explicación geométrica novedosa para el éxito de la poda en tiempo de inferencia. Argumenta que la importancia de las capas no es fija, sino condicional tanto a la tarea como a la distribución de entrada. Los autores postulan que los modelos preentrenados codifican múltiples funciones especializadas en tareas, cada una asociada con una geometría de representación local. Las entradas ID se ajustan a esta geometría, mientras que las entradas OOD la violan.

El significado de este trabajo radica en:

1. Clarificación de los Límites de la Poda: Establece que la poda consciente de la tarea no es un potenciador de rendimiento universal, sino un mecanismo correctivo para el desplazamiento de distribución.
2. Interpretación Geométrica: Cambia la comprensión de la redundancia de capas desde conteos estáticos de parámetros hacia distorsiones geométricas dinámicas dependientes de la distribución.
3. Perspectiva Mecanística: Proporciona evidencia causal de que las ganancias de rendimiento surgen de corregir la trayectoria de las representaciones internas para que coincidan con la geometría adaptada del modelo, en lugar de aumentar la capacidad del modelo o reducir el ruido.

Los autores concluyen que la poda mejora el rendimiento OOD al atenuar las capas que amplifican el desajuste entre la distribución de entrada y la geometría de tarea aprendida del modelo, realineando efectivamente el estado interno del modelo con las condiciones bajo las cuales fue optimizado.