Resumen Técnico: TAS-LoRA

Declaración del Problema

La Búsqueda de Arquitectura de Transformadores (TAS) tiene como objetivo descubrir automáticamente arquitecturas óptimas de Transformadores de Visión (ViT), reduciendo el esfuerzo manual requerido para el diseño. Los métodos actuales de TAS más avanzados dependen principalmente de enfoques one-shot, que entrenan una "supernet" que contiene todas las subredes posibles. En estas supernets con pesos entrelazados, todas las subredes comparten un conjunto común de pesos.

El artículo identifica una limitación crítica en este paradigma: el problema del colapso de características. Debido a que la supernet se optimiza para minimizar una pérdida global en todas las subredes, los pesos compartidos tienden a aprender características genéricas que funcionan adecuadamente para arquitecturas diversas pero fallan en capturar características específicas de la subred. En consecuencia, las subredes individuales no pueden aprovechar plenamente sus propiedades arquitectónicas únicas (por ejemplo, números específicos de cabezas de atención o ratios MLP), lo que conduce a un rendimiento subóptimo en comparación con las subredes entrenadas independientemente desde cero. El análisis empírico en el artículo confirma que las subredes dentro de supernets con pesos entrelazados exhiben representaciones de características casi idénticas (alta similitud coseno), mientras que las subredes entrenadas independientemente desarrollan representaciones distintas.

Metodología: TAS-LoRA

Para abordar el colapso de características manteniendo la eficiencia computacional, los autores proponen TAS-LoRA (Búsqueda de Arquitectura de Transformadores con Expertos de Mezcla de LoRA). El método integra la Adaptación de Bajo Rango (LoRA) en el proceso de entrenamiento de la supernet para permitir el aprendizaje de características especializado para cada subred sin incurrir en el costo prohibitivo de asignar parámetros únicos a cada configuración posible de subred.

1. Mezcla de Expertos LoRA (MoLE)

En lugar de asignar un módulo LoRA único a cada subred (lo cual es computacionalmente inviable debido al vasto espacio de búsqueda), TAS-LoRA emplea una estrategia de Mezcla de Expertos LoRA (MoLE).

• Arquitectura: Los expertos LoRA se aplican a las capas lineales de la supernet (incluyendo las capas de proyección QKV y las capas MLP).
• Enrutamiento Dinámico: Un enrutador ligero selecciona y combina dinámicamente un conjunto limitado de expertos LoRA basándose en la arquitectura específica de la subred muestreada.
• Eficiencia de Inferencia: En el momento de la inferencia, los parámetros LoRA seleccionados se fusionan en los pesos originales de la supernet. Dado que la configuración de la subred es fija durante la inferencia, el enrutador se ejecuta solo una vez, y los pesos fusionados introducen ninguna sobrecarga computacional adicional.

2. Diseño del Enrutador

El enrutador está diseñado para capturar las dependencias secuenciales de los bloques transformadores:

• Entrada: Toma tanto atributos a nivel de bloque (por ejemplo, número de cabezas de atención, ratio MLP) como atributos a nivel de subred (por ejemplo, dimensión de incrustación, profundidad) como entrada.
• Mecanismo: Estos atributos son codificados por una capa de incrustación aprendible y procesados por una LSTM para capturar dependencias secuenciales a través de los bloques. La LSTM produce representaciones ocultas, que se pasan a través de una "cabeza de enrutador" totalmente conectada para generar una distribución de probabilidad sobre los expertos LoRA disponibles para cada capa.

3. Inicialización del Enrutador por Grupos

Un desafío clave identificado es que, durante las primeras etapas del entrenamiento, el enrutador tiende a asignar pesos similares a todos los expertos, lo que lleva a redundancia y previene la especialización. Para mitigar esto, los autores introducen una técnica de inicialización del enrutador por grupos:

• Agrupación: Los bloques transformadores a través de diferentes subredes se agrupan basándose en similitudes arquitectónicas (específicamente conteo de cabezas de atención, ratio MLP y dimensión de incrustación).
• Inicialización: El número de expertos LoRA se establece para coincidir con el número de grupos. El enrutador se inicializa de tal manera que cada experto está sesgado hacia un grupo distinto. Esto se logra estableciendo términos de sesgo específicos en la cabeza del enrutador a un valor positivo $\beta$ para el par grupo-experto correspondiente, mientras que los demás se ponen a cero.
• Efecto: Esto obliga al enrutador a asignar expertos distintos a diferentes grupos arquitectónicos al inicio del entrenamiento, fomentando un aprendizaje de características diverso. A medida que avanza el entrenamiento, el enrutador actualiza dinámicamente estas asignaciones basándose en las propiedades de la subred.

Contribuciones Clave

1. Identificación del Colapso de Características en TAS: El artículo identifica y analiza formalmente el problema del colapso de características en supernets con pesos entrelazados, demostrando que los pesos compartidos impiden que las subredes aprendan representaciones especializadas.
2. Marco TAS-LoRA: Un método novedoso que combina LoRA con un mecanismo de enrutamiento dinámico (MoLE) para aprender características específicas de la subred de manera eficiente. Permite que las subredes refinen características compartidas basándose en sus propiedades arquitectónicas sin reentrenar toda la supernet desde cero.
3. Inicialización del Enrutador por Grupos: Una técnica para prevenir la redundancia de expertos durante las etapas tempranas del entrenamiento, asegurando que los expertos LoRA aprendan representaciones de características distintas.
4. Rendimiento de Vanguardia: Experimentos extensivos demuestran que TAS-LoRA supera significativamente a los métodos TAS existentes (por ejemplo, AutoFormer, DYNAS, FocusFormer) en diversos espacios de búsqueda y conjuntos de datos.

Resultados Experimentales

Los autores evaluaron TAS-LoRA en el conjunto de datos ImageNet y varias pruebas de aprendizaje por transferencia (CIFAR-10/100, Flowers, CARS, INAT-19).

• Rendimiento en ImageNet: TAS-LoRA mejoró consistentemente la precisión top-1 de supernets entrenadas con AutoFormer y DYNAS en los espacios de búsqueda Tiny, Small y Base.
  • Por ejemplo, en el espacio AutoFormer-T, TAS-LoRA logró una precisión del 75.7%, superando al AutoFormer original (74.7%) y a otros métodos SOTA como PreNAS-Ti (77.1% pero con menos flexibilidad).
  • En el espacio AutoFormer-B, TAS-LoRA-B logró una precisión del 82.6% con 54M de parámetros, superando a AutoFormer-B (82.4%) y igualando o superando modelos más grandes como GLiT-B (82.3% con 96M de parámetros).
• Aprendizaje por Transferencia: TAS-LoRA demostró una generalización superior, logrando resultados competitivos o de vanguardia en CIFAR-10/100, Flowers, CARS e INAT-19.
• Estudios de Ablación:
  • Inicialización: La inicialización por grupos superó consistentemente a la inicialización aleatoria, confirmando su papel en la reducción de la redundancia de expertos.
  • Similitud de Características: El análisis de similitud coseno mostró que TAS-LoRA con inicialización por grupos produjo representaciones de características significativamente más diversas entre expertos en comparación con la inicialización aleatoria o los métodos estándar con pesos entrelazados.
  • Comparación con Ajuste Fino: TAS-LoRA superó a una supernet completamente ajustada (donde se actualizan todos los pesos) a pesar de entrenar solo durante 50 épocas en las capas MoLE. Esto indica que las ganancias provienen de la capacidad de aprendizaje de características específica de la subred de MoLE, no meramente de un entrenamiento extendido.

Significado y Afirmaciones

El artículo afirma que TAS-LoRA representa un avance significativo en la Búsqueda de Arquitectura de Transformadores al resolver eficazmente el problema del colapso de características inherente a las supernets con pesos entrelazados.

• Eficiencia vs. Rendimiento: A diferencia de los métodos que requieren entrenar múltiples supernets (por ejemplo, MoS) o reentrenar desde cero cuando cambian las restricciones de hardware (por ejemplo, PreNAS), TAS-LoRA entrena una sola supernet con capas de adaptación específicas de la subred. Esto permite una adaptación flexible a nuevas restricciones de hardware (por ejemplo, FLOPs, latencia) sin reentrenar todo el modelo.
• Escalabilidad: Al utilizar un conjunto limitado de expertos y un enrutador dinámico, el método escala a supernets con números enormes de configuraciones posibles de subred sin un aumento lineal en parámetros o costo computacional.
• Generalizabilidad: El enfoque se presenta como una solución general aplicable a cualquier supernet con pesos entrelazados, mejorando el rendimiento independientemente de la estrategia de entrenamiento subyacente (por ejemplo, AutoFormer o DYNAS).

Los autores concluyen que TAS-LoRA mitiga con éxito el colapso de características, permitiendo que las subredes aprendan representaciones especializadas y logren un rendimiento de vanguardia tanto en tareas de búsqueda de arquitectura como en aprendizaje por transferencia.