Una guía completa de SegFormer: por qué un transformador sin codificación posicional funciona para la segmentación.

Antes de SegFormer, ya existían intentos de introducir Transformers en la segmentación semántica, pero muchos dependían de codificaciones posicionales y decodificadores complejos. Por lo tanto, se enfrentaban a dos problemas: la precisión podía disminuir cuando la resolución de la prueba difería de la resolución de entrenamiento, y el cálculo podía resultar costoso. SegFormer (Xie, Wang, Yu, Anandkumar, Alvarez y Luo, NeurIPS 2021) aborda ambos problemas con un diseño notablemente simple: un codificador jerárquico sin codificación posicional y un decodificador All-MLP que no utiliza convolución alguna. Este artículo examina, paso a paso, por qué funciona este diseño, basándose en el artículo original (arXiv:2105.15203).

Este artículo se basa en el trabajo original “SegFormer: Diseño simple y eficiente para la segmentación semántica con Transformers” (Xie et al., NeurIPS 2021, arXiv:2105.15203).

0. Lo que aprenderá

• Los problemas que SegFormer se propuso resolver en los modelos de segmentación anteriores basados en Transformers
• Cómo funciona el codificador jerárquico MiT (Mix Transformer) sin codificación posicional
• Por qué un decodificador compuesto únicamente por MLP, sin convolución alguna, puede lograr una alta precisión
• Medición de mIoU, recuento de parámetros y coste computacional para B0–B5 en Cityscapes y ADE20K
• Cómo elegir un modelo en la práctica, incluyendo la robustez ante la corrupción de imágenes

1. ¿Qué es SegFormer?

SegFormer es un modelo de segmentación semántica que combina un codificador Transformer jerárquico (MiT) sin codificación posicional con un decodificador All-MLP ligero que no utiliza convolución. Su idea central es rediseñar la división del trabajo entre el codificador y el decodificador: una vez que el codificador ha producido características en múltiples resoluciones, el decodificador se puede simplificar enormemente.

2. ¿Por qué fue necesario este diseño?

Utilizar un Vision Transformer (ViT) directamente como codificador para la segmentación semántica genera dos problemas. Primero, las características producidas por ViT tienen una única resolución. Los objetos pueden tener muchos tamaños diferentes dentro de una misma imagen, y los modelos de segmentación convencionales basados en CNN, como FCN y U-Net, abordan esto combinando características en múltiples resoluciones. Un ViT que produce una sola resolución pierde esta capacidad de representación multiescala.

El segundo problema es la dependencia de la codificación posicional. ViT divide una imagen en parches y los introduce en un Transformer, pero este no dispone de información sobre la ubicación de cada parche en la imagen. Por lo tanto, es necesario añadir por separado una codificación posicional fija o aprendida. Dado que esta codificación se basa en la resolución de entrada utilizada durante el entrenamiento, una imagen con una resolución diferente en el momento de la inferencia requiere la interpolación de la codificación posicional, lo que puede reducir la precisión. La segmentación suele requerir entradas de mayor resolución que la detección, lo que hace que esta dependencia de la resolución sea especialmente problemática.

SegFormer aborda ambos problemas de raíz mediante el uso de un codificador jerárquico y la eliminación total de la codificación posicional.

3. ¿Cuál es la entrada?

Al igual que otros modelos de segmentación, la entrada es una imagen RGB x \in \mathbb{R}^{H \times W \times 3}. SegFormer se evalúa con imágenes de alta resolución (1024×2048) para Cityscapes e imágenes de aproximadamente 512×512 para ADE20K. Mientras que ViT realiza un submuestreo con parches gruesos de 16×16, el codificador de SegFormer comienza con parches de tan solo 4×4, preservando así más detalles al inicio de la extracción de características.

4. ¿Qué predice?

La salida es un mapa de segmentación a 1/4 de la resolución de entrada, donde la clase de cada píxel se representa como \hat{y} \in \mathbb{R}^{\frac{H}{4} \times \frac{W}{4} \times N_{cls}} (N_{cls} es el número de clases). Posteriormente, se restaura a la resolución original mediante interpolación del vecino más cercano o un método similar. La idea principal es que, si el codificador crea características multiescala suficientemente ricas, el decodificador solo necesita una combinación de capas lineales simples para integrarlas y convertirlas en una máscara precisa.

5. Arquitectura básica

SegFormer consta de dos bloques principales: un codificador MiT (Mix Transformer) que extrae características en cuatro resoluciones, y un decodificador All-MLP que combina esas características y genera una máscara de segmentación.

Figura 1: El codificador MiT produce características en cuatro resoluciones (1/4, 1/8, 1/16 y 1/32), y el decodificador All-MLP las combina utilizando únicamente capas lineales para crear una máscara de segmentación. La única convolución se encuentra en la Mix-FFN del codificador; El decodificador está compuesto, literalmente, en su totalidad por capas lineales.

Dado que el codificador genera características con cuatro resoluciones por adelantado, el decodificador puede obtener una máscara precisa con un procesamiento mínimo: alinear las características y combinarlas. Esta división del trabajo es la razón principal por la que SegFormer mantiene un número reducido de parámetros.

6. Detalles técnicos de los componentes

Fusión de parches superpuestos: una jerarquía de cuatro etapas

En lugar de dividir la imagen en parches una sola vez, como hace ViT, el codificador MiT reduce gradualmente la resolución en cuatro etapas. Las resoluciones de salida son 1/4, 1/8, 1/16 y 1/32 de la entrada, reproduciendo la estructura jerárquica utilizada por las redes neuronales convolucionales como ResNet.

El método de partición de parches también difiere de los parches no superpuestos de ViT: los parches adyacentes se superponen al fusionarse. La primera etapa utiliza un tamaño de núcleo K=7, un paso S=4 y un relleno P=3; las etapas posteriores utilizan K=3, S=2, P=1. La superposición preserva la continuidad local entre los límites de los parches (la información compartida por parches vecinos) mientras se reduce el mapa de características.

Autoatención eficiente

La autoatención multi-cabeza convencional requiere un cálculo de O(N^2) para una secuencia de longitud N. Dado que N se vuelve muy grande para la segmentación de alta resolución, esto representa un cuello de botella importante. SegFormer introduce la Reducción de Secuencia, que comprime las secuencias de Clave y Valor mediante una relación de reducción R antes de calcular la atención.

Las secuencias K, V se reducen de longitud N a N/R antes de su uso. La relación de reducción R para las etapas 1 a 4 se establece en [64, 16, 4, 1]: las etapas superficiales tienen alta resolución y secuencias largas, por lo que se comprimen agresivamente; las etapas profundas tienen menor resolución y secuencias más cortas, por lo que no se comprimen en absoluto. Esto reduce el costo computacional de O(N^2) a O(N^2/R) y hace que la autoatención sea práctica incluso para entradas de alta resolución.

Mix-FFN: cómo se elimina la codificación posicional

SegFormer puede evitar por completo la codificación posicional porque mezcla una convolución de 3×3 en la red neuronal de alimentación directa (FFN), creando así Mix-FFN.

La convolución de 3×3 utiliza relleno de ceros en el borde de la imagen. Este borde con relleno filtra indirectamente información sobre la posición en la imagen, de modo que la información posicional se integra en las características simplemente al pasar por Mix-FFN, sin una codificación posicional explícita. Como resultado, las pruebas con una resolución diferente a la de entrenamiento no requieren la interpolación de una codificación posicional, lo que podría reducir la precisión.

Decodificador All-MLP: por qué la convolución es innecesaria

El decodificador consta de los siguientes cuatro pasos, utilizando únicamente capas lineales (MLP).

1. Unificar canales: mapear cada característica de etapa F_i al mismo número de canales C con una capa lineal independiente.

2. Sobremuestreo y concatenación: sobremuestrear cada característica a 1/4 de la resolución y concatenarlas a lo largo de la dimensión del canal.

3. Fusionar: mapea la característica concatenada de 4C dimensiones a C dimensiones con una capa lineal.

4. Predecir: genera una máscara de segmentación con un canal para cada clase mediante una capa lineal final.

El artículo original respalda este diseño con un análisis del campo receptivo efectivo (ERF). Los decodificadores basados en CNN, como DeepLabv3+, requieren mecanismos complejos —convoluciones dilatadas y agrupamiento multiescala— para obtener un campo receptivo amplio y lograr una alta precisión. En contraste, el codificador MiT adquiere de forma natural ERF locales en etapas superficiales y ERF no locales que cubren toda la imagen en etapas profundas mediante autoatención, sin mecanismos adicionales. Dado que el propio codificador ya construye características que contienen información desde escalas locales hasta globales, el decodificador no necesita un mecanismo complejo de expansión del campo receptivo. Esta es la razón teórica por la que una simple combinación de capas lineales es suficiente.

Receta de entrenamiento

El artículo original utiliza una configuración de entrenamiento sencilla, sin trucos especiales. El codificador MiT se preentrena en ImageNet-1K, mientras que el decodificador se entrena a partir de una inicialización aleatoria. Se utiliza AdamW para la optimización, con una tasa de aprendizaje que disminuye desde un valor inicial de 0,00006 bajo un esquema polinomial (decaimiento de ley de potencias con coeficiente 1,0). El número de iteraciones es de 160.000 en ADE20K y Cityscapes, y de 80.000 en COCO-Stuff. El aumento de datos durante el entrenamiento se limita al redimensionamiento aleatorio con una relación de 0,5 a 2,0, el volteo horizontal aleatorio y el recorte aleatorio (512×512 para ADE20K, 1024×1024 para Cityscapes y 640×640 para B5 en ADE20K). El artículo declara explícitamente que no utiliza técnicas ampliamente adoptadas como OHEM (Online Hard Example Mining), pérdidas auxiliares o pérdidas con balance de clases. El objetivo es demostrar que el codificador MiT y el decodificador All-MLP pueden alcanzar un alto mIoU sin estas adiciones.

7. Comparación de variantes del modelo

SegFormer está disponible en seis tamaños, desde MiT-B0 hasta MiT-B5. A continuación se muestran los resultados medidos de Cityscapes y ADE20K, que se reportan en la Tabla 2 del artículo original.

Fuente: Tabla 2 del artículo original (Xie et al., NeurIPS 2021). El mIoU de Cityscapes se mide mediante pruebas multiescala y de inversión izquierda-derecha (MS). El recuento de parámetros incluye el codificador y el decodificador completos; incluso para B5, el decodificador representa solo alrededor del 4 % del total. Como referencia, B5 alcanza 82,4 mIoU con pruebas de escala única (SS) y 84,0 con pruebas MS.

El artículo compara estos resultados con el estado del arte de la época e informa que SegFormer-B4 alcanza un mIoU del 50,3 % en ADE20K con 64,1 millones de parámetros, superando al mejor método de ese período en 2,2 puntos con un modelo casi una quinta parte del tamaño. B0 mantiene un mIoU práctico de Cityscapes de 76,2 con solo 3,8 millones de parámetros, y los números respaldan su diseño orientado a los bordes.

8. Sus dificultades

La mayoría de las debilidades de SegFormer son comunes a los modelos basados en Transformer en general. En primer lugar, los modelos más grandes (B3 y superiores) tienen altos costos de entrenamiento e inferencia. El consumo de operaciones de punto flotante (FLOPs) de Cityscapes aumenta de 125,5 GB para B0 a 1460,4 GB para B5, un incremento de más de diez veces. Para la segmentación de alta resolución, este aumento puede dificultar el funcionamiento en tiempo real.

La dependencia de grandes conjuntos de datos de preentrenamiento también es importante. El codificador MiT se preentrena con ImageNet-1K, por lo que la calidad y cantidad de los datos de ajuste fino afectan significativamente la transferencia a dominios que difieren sustancialmente de los datos de preentrenamiento, como imágenes médicas o satelitales.

Los objetos muy pequeños y las estructuras delgadas, como cables y señales de tráfico, tampoco representan una tarea sencilla para SegFormer. La fusión de parches superpuestos reduce la resolución en 4x4 incluso en la primera etapa, por lo que las estructuras de tan solo unos pocos píxeles de ancho pueden perder información antes de alcanzar etapas más profundas.

Al mismo tiempo, la evaluación de robustez en Cityscapes-C, reportada en el artículo original, es una clara fortaleza. Para imágenes corruptas con ruido gaussiano añadido, SegFormer-B5 registra una mejora máxima relativa de mIoU del 588% con respecto a DeepLabv3+ con una arquitectura base Xception-71; para corrupción tipo nieve, la mejora relativa máxima es del 295%. Estas mediciones muestran una mayor tolerancia a la corrupción de imágenes que los modelos CNN convencionales. Se cree que esta robustez se debe a que la autoatención es menos susceptible a la influencia del ruido local.

9. Cómo elegir un modelo en la práctica

Para dispositivos periféricos y procesamiento en tiempo real, SegFormer-B0 o B1 son opciones realistas. B0 equilibra precisión y compacidad con 3,8 millones de parámetros y 76,2 mIoU en Cityscapes, lo que hace viable su implementación en dispositivos integrados y drones.

Para aplicaciones como el reconocimiento de áreas transitables en la conducción autónoma, donde se requieren imágenes de alta resolución y precisión constante, B2 o B3 suelen ser una solución intermedia razonable. B4 y B5 ofrecen la mayor precisión, pero sus operaciones de punto flotante (FLOPs) en Cityscapes superan los 1200 GB, por lo que la inferencia en tiempo real en una GPU integrada en el vehículo generalmente requiere optimización, como la cuantización o la conversión a TensorRT.

Para el análisis de precisión offline de imágenes médicas o satelitales, la precisión tiene prioridad sobre el rendimiento en tiempo real. B4 o B5 son apropiados, con ajustes finos en datos específicos del dominio según sea necesario.

Para entornos exteriores con frecuentes inclemencias del tiempo o cambios de iluminación, como robots agrícolas y cámaras de vigilancia exteriores, la robustez medida en Cityscapes-C representa una ventaja práctica. En comparación con los modelos basados en CNN de precisión similar, se espera que SegFormer pierda menos precisión ante la corrupción de imágenes.

Para aprender Detección de Objetos y Segmentación Semántica desde sus fundamentos, consulte «Introducción a la Detección de Objetos y la Segmentación Semántica». Para conocer las últimas tendencias en segmentación, consulte «Tendencias Tecnológicas en Segmentación Semántica».

10. Resumen en tres líneas

• SegFormer combina un codificador jerárquico (MiT) que produce características multiescala sin codificación posicional con un decodificador All-MLP que no utiliza convolución.

• La convolución 3×3 en Mix-FFN filtra implícitamente información posicional, eliminando la necesidad de codificación posicional, mientras que el amplio campo receptivo efectivo del Transformer permite un decodificador MLP simple. La familia abarca desde B0 (3,8 millones de parámetros, 76,2 mIoU en Cityscapes) hasta B5 (84,7 millones, 84,0 mIoU), y los resultados de las mediciones también confirman su robustez ante la corrupción de imágenes.

Referencias

• SegFormer: Diseño simple y eficiente para la segmentación semántica con transformadores (NeurIPS 2021, arXiv:2105.15203)
• Artículo sobre SegFormer en ar5iv (texto completo)
• Implementación oficial de SegFormer (NVlabs/SegFormer, GitHub)
• Actas de NeurIPS 2021: SegFormer