Resumen Técnico: CORP - Poda Estructurada de Representación Preservada en Un Solo Paso de Forma Cerrada para Transformers

1. Declaración del Problema

Los modelos basados en Transformers logran un rendimiento de vanguardia, pero incurren en costos computacionales y de memoria prohibitivos, lo que dificulta su despliegue en hardware con recursos limitados. Si bien la poda estructurada (eliminación de dimensiones o subestructuras completas) ofrece aceleraciones prácticas, la mayoría de los métodos existentes dependen de reentrenamiento, optimización iterativa o acceso a datos etiquetados y gradientes. En escenarios estrictos de despliegue post-entrenamiento, estos recursos a menudo no están disponibles.

La poda estructurada ingenua en un solo paso falla bajo estas restricciones porque eliminar canales o dimensiones de atención introduce errores sistemáticos de representación que se acumulan a través de las capas, lo que lleva a un colapso rápido de la precisión incluso en niveles moderados de dispersión. Los autores argumentan que este fracaso no se debe meramente a la calidad de la clasificación por importancia, sino a la falta de recuperación explícita de errores para los componentes eliminados.

2. Metodología: CORP

Los autores proponen CORP (Poda Estructurada de Representación Preservada en Un Solo Paso de Forma Cerrada), un marco que replantea la poda como un problema de recuperación de representación. CORP opera sin etiquetas, gradientes o ajuste fino, utilizando únicamente un pequeño conjunto de datos de calibración no etiquetados.

Filosofía Central

En lugar de ver la poda únicamente como una tarea de selección, CORP modela los componentes eliminados como funciones afines de los componentes retenidos. Deriva soluciones de regresión ridge de forma cerrada para estimar las contribuciones faltantes e incorpora estas compensaciones directamente en los pesos de las capas restantes. Esto preserva la distribución de salida de la capa sin requerir optimización iterativa ni módulos de inferencia adicionales.

Componentes Técnicos

El método se dirige a dos estructuras principales en los Transformers:

A. Poda de Canales MLP

Para una capa MLP definida por $y = Wx + b$, donde la entrada $x$ se divide en canales retenidos ($x_S$) y podados ($x_P$):

1. Modelado: Las activaciones podadas $x_P$ se aproximan como una función afín de las activaciones retenidas: $x_P \approx Bx_S + c$.
2. Solución de Forma Cerrada: Los parámetros $(B, c)$ se derivan mediante regresión ridge sobre el conjunto de calibración $D$:
   $$B = \Sigma_{PS}(\Sigma_{SS} + \lambda I)^{-1}, \quad c = \mu_P - B\mu_S$$
   donde $\Sigma$ representa las matrices de covarianza de las activaciones.
3. Plegado de Pesos: La compensación se pliega en los pesos restantes:
   $$\hat{W}_S = W_S + W_P B, \quad \hat{b} = b + W_P c$$
   Esto resulta en una capa lineal estándar con dimensiones reducidas que minimiza el error de reconstrucción de la salida original.

B. Poda de Dimensiones de Atención (Consulta/Llave)

Para la auto-atención, CORP poda dimensiones en las proyecciones de Consulta ($Q$) y Llave ($K$), dejando intacta la proyección de Valor ($V$) para evitar problemas complejos de acoplamiento.

1. Recuperación de Logits: Los logits de atención originales son $L = QK^\top = Q_S K_S^\top + Q_P K_P^\top$. El término $Q_P K_P^\top$ (logits faltantes) se aproxima utilizando los subespacios retenidos: $Q_P K_P^\top \approx Q_S M K_S^\top$, donde $M$ es una pequeña matriz de compensación.
2. Solución de Forma Cerrada: $M$ se estima resolviendo un problema de regresión ridge sobre el conjunto de calibración, formulado como una ecuación de Sylvester o un sistema lineal vectorizado:
   $$\left[ \sum_b (K_{S,b}^\top K_{S,b}) \otimes (Q_{S,b}^\top Q_{S,b}) + \lambda I \right] \text{vec}(M) = \sum_b \text{vec}((Q_{S,b}^\top Q_{P,b})(K_{P,b}^\top K_{S,b}))$$
3. Plegado de Pesos: La matriz $I + M$ se factoriza mediante SVD ($U\Sigma V^\top$) y se pliega en las proyecciones de Consulta y Llave:
   $$\hat{W}_{Q,S} = W_{Q,S} U \Sigma^{1/2}, \quad \hat{W}_{K,S} = W_{K,S} V \Sigma^{1/2}$$

Estrategia de Clasificación

Antes de la compensación, CORP clasifica los canales/dimensiones para su eliminación utilizando métricas simples basadas en datos:

• MLP: Una métrica combinada de energía de activación ($E_i = E[x_i^2]$) y magnitud de pesos ($|W_{:,i}|^2$).
• Atención: Energía de logits esperada ($s_j = E[\|q_j\|^2 \|k_j\|^2]$).
  Se podan los canales con las puntuaciones más bajas.

3. Contribuciones Clave

1. Replanteamiento de la Poda: Los autores desplazan el enfoque de la poda en un solo paso de la mera "clasificación" a la "recuperación de representación", demostrando que la compensación es crítica para preservar la precisión.
2. Marco de Forma Cerrada: CORP proporciona una solución unificada y de forma cerrada tanto para la poda de canales MLP como para la poda de dimensiones de atención Q/K, que no requiere gradientes, etiquetas ni ajuste fino.
3. Compensación Unificada: A diferencia de trabajos anteriores que pueden compensar solo la salida de un bloque, CORP realiza recuperación de representación a nivel de capa e introduce compensación específica en el espacio de logits para las dimensiones de atención, incorporando todas las correcciones directamente en las proyecciones reducidas.

4. Resultados Experimentales

Los experimentos se realizaron en ImageNet (modelos DeiT y ViT), WikiText-2 (OPT-1.3B) y tareas posteriores (DINOv2 en NYU Depth V2 y ADE20k).

• Preservación de la Precisión:

  • DeiT-Huge: Logra una precisión Top-1 del 83.27% después de podar el 50% de las estructuras tanto de MLP como de atención (una caída de solo ~1.7% desde la línea base de 84.97%).
  • Escalabilidad: Los modelos más grandes (por ejemplo, DeiT-Huge) muestran una mayor redundancia y una mejor retención de precisión bajo dispersión agresiva en comparación con modelos más pequeños (por ejemplo, DeiT-Tiny).
  • Comparación: CORP supera o iguala a métodos de vanguardia como GRAIL, SNOWS y VBP. Por ejemplo, en DeiT-Base con 50% de dispersión MLP, CORP logra un 72.0% Top-1, en comparación con el 68.6% de GRAIL y el 66.4% de VBP.
  • Sin Ajuste Fino: CORP iguala el rendimiento de DC-ViT (que utiliza ajuste fino de imitación de características con pocos ejemplos) sin ninguna actualización de gradientes.

• Ganancias de Eficiencia:

  • DeiT-Huge: Con un 50% de dispersión, los parámetros y los FLOPs se reducen en un ~40–45%, lo que resulta en una aceleración de rendimiento de 1.6× en una NVIDIA RTX 3090.
  • OPT-1.3B: Podar el 30% de las estructuras tanto de MLP como de atención reduce los FLOPs en un 27.2% y los parámetros en un 25.0%, con un aumento modesto de la perplejidad (de 14.60 a 25.41 para la poda conjunta).
  • Tareas Posteriores: En DINOv2, la poda del 50% de la columna vertebral preserva el rendimiento posterior en estimación de profundidad (el RMSE mejora de 0.3308 a 0.3167 para DINOv2-Huge) y segmentación.

• Eficiencia de Calibración:

  • El método es robusto al tamaño del conjunto de calibración; los modelos más grandes permanecen estables incluso con conjuntos de calibración pequeños (por ejemplo, 1000 muestras).
  • El análisis del tiempo de ejecución muestra que la calibración domina el costo, mientras que la clasificación y la compensación de forma cerrada añaden una sobrecarga insignificante.

5. Significado y Afirmaciones

El artículo afirma que CORP demuestra que la poda estructurada puede realizarse efectivamente en un entorno estricto de un solo paso y post-entrenamiento tratándola como un problema de recuperación de representación. El significado clave radica en:

• Desplegabilidad Práctica: Permite la compresión en hardware donde el reentrenamiento es imposible debido a la falta de etiquetas o gradientes.
• Perspectiva Teórica: Destaca que el "error sistemático" introducido al eliminar componentes aditivos es la causa principal del colapso de la precisión, y que una compensación afín/de logits simple puede recuperar este error sin optimización compleja.
• Amplia Aplicabilidad: El método muestra resultados sólidos en modelos de visión (DeiT, ViT, DINOv2) y modelos de lenguaje (OPT), sugiriendo una fuerte redundancia tanto en las representaciones MLP como en las de atención a través de modalidades.

Los autores notan modestamente las limitaciones, incluida la suposición de compensación de módulos independientes (ignorando las interacciones entre módulos), el uso de ratios de dispersión uniformes y la falta de optimización para la composición de cuantización-poda. Afirman que el rendimiento puede degradarse bajo fuertes desplazamientos de distribución entre los datos de calibración y los de despliegue.