Reg4Pru: Regularisation Through Random Token Routing for Token Pruning

Resumen Técnico: Reg4Pru – Regularización mediante el enrutamiento aleatorio de tokens para la poda de tokens

Planteamiento del problema
Los Vision Transformers (ViTs) se han convertido en fundamentales para los modelos de visión modernos debido a sus capacidades de escalabilidad y generalización. Sin embargo, su dependencia de los mecanismos de autoatención resulta en una complejidad computacional que escala cuadráticamente ($O(N^2)$) con el número de tokens $N$. Este cuello de botella es particularmente agudo en la segmentación de imágenes médicas, donde se requieren entradas de alta resolución para preservar detalles diagnósticos críticos.

Aunque los métodos de poda de tokens (eliminar tokens redundantes) y las estrategias de reactivación (reintroducir tokens podados durante la inferencia para obtener predicciones densas píxel a píxel) ofrecen un camino hacia la eficiencia, los autores observan una degradación significativa del rendimiento en las capas profundas. Específicamente, cuando los tokens podados se reactivan durante la inferencia, exhiben representaciones drásticamente diferentes a su estado durante el entrenamiento. Este cambio de distribución entre el entrenamiento y la prueba provoca inestabilidad y un pobre rendimiento en la predicción densa, causando contraintuitivamente que el rendimiento disminuya a medida que aumenta la profundidad del modelo.

Metodología
El artículo propone Reg4Pru, una técnica de regularización durante el entrenamiento diseñada para mitigar la inestabilidad asociada con la poda y la reactivación de tokens. El mecanismo central consiste en el enrutamiento aleatorio de tokens durante la fase de entrenamiento para simular los efectos de la poda sin descartar permanentemente los tokens.

1. Enrutamiento Aleatorio de Tokens: Durante la pasada hacia adelante (forward pass) del entrenamiento, el modelo selecciona un rango de enrutamiento aleatorio $[l, n]$ entre los bloques intermedios del transformer. Se elige un subconjco aleatorio de tokens ($\tau_r$) para omitir el cómputo de los bloques desde $l$ hasta $n$, mientras que los tokens restantes ($\tau_k$) se procesan normalmente.
2. Flujo de Gradiente y Representación: Los tokens enrutados omiten bloques específicos, cortando efectivamente sus gradientes con respecto a esas capas. Sin embargo, al aleatorizar tanto el subconjunto de tokens enrutados como el rango de enrutamiento $[l, n]$ en cada paso, diferentes tokens contribuyen con gradientes a diferentes bloques a través de las iteraciones. Esto obliga al modelo a aprender representaciones invariantes a la profundidad, asegurando que los tokens puedan ser reactivados en cualquier profundidad durante la inferencia sin sufrir un desajuste de representación.
3. Estrategia de Inferencia: En la inferencia, una red de política separada predice qué tokens podar basándose en un esquema jerárquico de ratio de permanencia (keep-ratio). Los tokens podados se eliminan temporalmente para ahorrar computación y se reactivan en el bloque final para generar máscaras de segmentación densas.
4. Funciones de Pérdida: El objetivo de entrenamiento combina una pérdida estándar basada en ratios (para imponer presupuestos de tokens) y, opcionalmente, una pérdida de política informada ($L_{pol}$) que utiliza las máscaras de segmentación de verdad de campo (ground truth) para guiar las decisiones de poda. El mecanismo de enrutamiento actúa como un regularizador junto con estas pérdidas.

Contribuciones Clave

• Análisis de Inestabilidad por Profundidad: Los autores identifican y caracterizan una inestabilidad específica en la poda de tokens para la predicción densa, donde el rendimiento se degrada en las capas más profundas debido a la reactivación de tokens que no han sido expuestos a la profundidad total de la red durante el entrenamiento.
• Marco Reg4Pru: Introducen una estrategia de enrutamiento de tokens durante el entrenamiento con límites aleatorios. A diferencia de otros métodos de enrutamiento que utilizan límites fijos o políticas de enrutamiento aprendidas, Reg4Pru aleatoriza el rango de enrutamiento para aproximar la distribución de características a través de todas las profundidades de poda, actuando así como un regularizador.
• Validación Empírica: El método se valida en el conjunto de datos de segmentación de vasos sanguíneos retinianos FIVES, demostrando que el enrutamiento aleatorio estabiliza significativamente el rendimiento en comparación con el enrutamiento fijo o sin enrutamiento.

Resultados
Los experimentos se realizaron en el conjunto de datos FIVES (imágenes de 2048$\times$2048, redimensionadas a cuadrantes para pruebas de rendimiento/throughput) utilizando una arquitectura EoMT inicializada con DINOv2-small.

• Ganancias de Rendimiento: En comparación con un modelo podado entrenado sin enrutamiento, Reg4Pru mejoró la Precisión Media (AP) en un 46% absoluto (de 25.60% a 71.67% para la configuración $L_{ratio}$).
• Eficiencia: La configuración Reg4Pru logró una aceleración relativa del 29% en el tiempo de ejecución (wall-clock time) en comparación con la línea base sin poda, manteniendo una alta calidad de segmentación.
• Estabilidad: El análisis de profundidad (Figura 5) mostró que, mientras los modelos podados sin enrutamiento sufrían una fuerte caída en la AP en los bloques más profundos, Reg4Pru mantuvo un rendimiento estable. El enrutamiento de límites fijos mitigó parte de la inestabilidad pero no logró evitar las caídas de rendimiento después de que terminaba el rango fijo, confirmando la necesidad de la aleatorización.
• Comparación: Reg4Pru superó a los métodos de Token Merging (ToMe) en términos de AP al comparar configuraciones de rendimiento similares, particularmente cuando ToMe se aplicó tanto a las capas MHSA como a las MLP.

Significación y Reivindicaciones
El artículo sostiene que Reg4Pru es un regularizador valioso para las estrategias de reducción de tokens en tareas de predicción densa. Los autores argumentan que la inestabilidad observada en la poda es el resultado de un cambio de distribución de entrenamiento-prueba respecto a las representaciones de los tokens en diferentes profundidades. Al aleatorizar el rango de enrutamiento durante el entrenamiento, Reg4Pru obliga al modelo a aprender representaciones que sean robustas al ser omitidas en diversas profundidades, cerrando eficazmente la brecha entre el entrenamiento y la inferencia.

Los autores mantienen la modestia respecto a la generalizabilidad de sus hallazgos, señalando que los experimentos se realizaron en un único conjunto de datos médicos de alta resolución. Sugieren que, si bien el método es efectivo para la segmentación, se requiere una evaluación adicional para evaluar su aplicabilidad en otras tareas de predicción densa (por ejemplo, estimación de profundidad) y su potencial como mecanismo de regularización general cuando se combina con técnicas como el dropout o la profundidad estocástica (stochastic depth). También reconocen que, aunque exploraron políticas informadas, las políticas basadas en ratios actualmente ofrecen un rendimiento superior en su configuración específica.