Modelos de lenguaje grande (LLM) como PaLM o GPT-3 mostró que escalar transformadores a cientos de miles de millones de parámetros mejora el rendimiento y desbloquea habilidades emergentes. Sin embargo, los modelos densos más grandes para la comprensión de imágenes han alcanzado solo 4 mil millones de parámetros, a pesar de que las investigaciones indican que los modelos multimodales prometedores como PaLI continúan beneficiándose de los modelos de visión escalables junto con sus contrapartes lingüísticas. Motivados por esto y los resultados de escalar los LLM, decidimos emprender el siguiente paso en el viaje de escalar Vision Transformer.
En “Transformadores de visión escalables a 22 mil millones de parámetros”, presentamos el modelo de visión densa más grande, ViT-22B. Es 5,5 veces más grande que la columna vertebral de visión más grande anterior, ViT-e, que tiene 4 mil millones de parámetros. Para habilitar este escalado, ViT-22B incorpora ideas de escalado de modelos de texto como PaLM, con mejoras tanto en la estabilidad del entrenamiento (usando normalización QK) y la eficiencia del entrenamiento (con un enfoque novedoso llamado operaciones lineales paralelas asíncronas). Como resultado de su arquitectura modificada, su receta de fragmentación eficiente y su implementación personalizada, se pudo capacitar en TPU en la nube con una alta utilización de hardware1. ViT-22B avanza el estado del arte en muchas tareas de visión utilizando representaciones congeladas o con un ajuste completo. Además, el modelo también se utilizó con éxito en PaLM-e, que demostró que un modelo grande que combina ViT-22B con un modelo de lenguaje puede avanzar significativamente en el estado del arte en tareas de robótica.
Arquitectura
Nuestro trabajo se basa en muchos avances de LLM, como PaLM y GPT-3. En comparación con la arquitectura estándar de Vision Transformer, utilizamos capas paralelasun enfoque en el que atención y bloques MLP se ejecutan en paralelo, en lugar de secuencialmente como en el Transformador estándar. Este enfoque se utilizó en PaLM y redujo el tiempo de capacitación en un 15 %.
En segundo lugar, ViT-22B omite sesgos en el QKV proyecciones, parte del mecanismo de autoatención, y en el Normas de capa, lo que aumenta la utilización en un 3%. El siguiente diagrama muestra la arquitectura de transformador modificada utilizada en ViT-22B:
 |
| La arquitectura del codificador de transformador ViT-22B utiliza capas paralelas de avance, omite sesgos en las capas QKV y LayerNorm y normaliza las proyecciones de Query y Key. |
Los modelos a esta escala requieren “fragmentación”: distribuir los parámetros del modelo en diferentes dispositivos informáticos. Junto a esto, también fragmentamos el activaciones (las representaciones intermedias de una entrada). Incluso algo tan simple como la multiplicación de una matriz requiere un cuidado especial, ya que tanto la entrada como la propia matriz se distribuyen entre dispositivos. Desarrollamos un enfoque llamado operaciones lineales paralelas asíncronas, mediante el cual las comunicaciones de activaciones y pesos entre dispositivos ocurren al mismo tiempo que los cálculos en la unidad de multiplicación matricial (la parte de la TPU que contiene la gran mayoría de la capacidad de cómputo). Este enfoque asincrónico minimiza el tiempo de espera en la comunicación entrante, lo que aumenta la eficiencia del dispositivo. La siguiente animación muestra un patrón de cálculo y comunicación de ejemplo para una multiplicación de matrices.
 |
| Operación lineal paralela asíncrona. El objetivo es calcular la multiplicación de la matriz y = Ax, pero tanto la matriz A como la activación x se distribuyen en diferentes dispositivos. Aquí ilustramos cómo se puede hacer con comunicación y computación superpuestas entre dispositivos. La matriz A está fragmentada en columnas a través de los dispositivos, cada uno con un segmento contiguo, cada bloque representado como Ayo. Más detalles están en el papel. |
Al principio, la nueva escala del modelo resultó en graves inestabilidades de entrenamiento. El enfoque de normalización de Gilmer et al. (2023, próximamente) resolvió estos problemas, lo que permitió un entrenamiento de modelos fluido y estable; esto se ilustra a continuación con progresiones de entrenamiento de ejemplo.
 |
| El efecto de normalizar las consultas y claves (normalización QK) en la capa de autoatención en la dinámica de entrenamiento. Sin la normalización de QK (rojo), los gradientes se vuelven inestables y la pérdida de entrenamiento diverge. |
Resultados
Aquí destacamos algunos resultados de ViT-22B. Tenga en cuenta que en el documento también exploramos varios otros dominios de problemas, como clasificación de videos, estimación de profundidady segmentación semántica.
Para ilustrar la riqueza de la representación aprendida, entrenamos un modelo de texto para producir representaciones que alinean las representaciones de texto e imagen (usando LiT-tuning). A continuación mostramos varios resultados de imágenes fuera de distribución generadas por Fiesta y Imagen:
 |
| Ejemplos de comprensión de imagen+texto para ViT-22B junto con un modelo de texto. El gráfico muestra la distribución de probabilidad normalizada para cada descripción de una imagen. |
Alineación de reconocimiento de objetos humanos
Para averiguar qué tan alineadas están las decisiones de clasificación de ViT-22B con las decisiones de clasificación humana, evaluamos ViT-22B ajustado con diferentes resoluciones en conjuntos de datos fuera de distribución (OOD) para los cuales los datos de comparación humanos están disponibles a través del caja de herramientas modelo-vs-humano. Esta caja de herramientas mide tres métricas clave: ¿Qué tan bien manejan los modelos las distorsiones (precisión)? ¿Cuán diferentes son las precisiones humanas y del modelo (diferencia de precisión)? Finalmente, ¿qué tan similares son los patrones de error humano y del modelo (consistencia del error)? Si bien no todas las resoluciones de ajuste fino desempeñarse igualmente bien, las variantes ViT-22B son lo último en tecnología para las tres métricas. Además, los modelos ViT-22B también tienen la más alta jamás registrada sesgo de forma en modelos de visión. Esto significa que utilizan principalmente la forma del objeto, en lugar de la textura del objeto, para informar las decisiones de clasificación, una estrategia conocida por la percepción humana (que tiene un sesgo de forma del 96%). Modelos estándar (p. ej., ResNet-50, que tiene un sesgo de forma de ~20–30%) a menudo clasifican imágenes como el gato con textura de elefante debajo según la textura (elefante); los modelos con un alto sesgo de forma tienden a centrarse en la forma (gato). Si bien todavía hay muchas diferencias importantes entre la percepción humana y la del modelo, ViT-22B muestra mayores similitudes con el reconocimiento de objetos visuales humanos.
 |
| ¿Gato o elefante? ¿Coche o reloj? ¿Pájaro o bicicleta? Imágenes de ejemplo con la forma de un objeto y la textura de un objeto diferente, utilizado para medir el sesgo de forma/textura. |
 |
| Evaluación del sesgo de forma (más alto = más sesgo de forma). Muchos modelos de visión tienen un sesgo de forma baja / textura alta, mientras que ViT-22B se ajustó en ImageNet (rojo, verde, azul entrenadas en imágenes 4B como lo indican los corchetes después de los nombres de los modelos, a menos que entrenen solo en ImageNet) tienen el sesgo de forma más alto registrado en un modelo ML hasta la fecha, acercándolos a un sesgo de forma similar al humano. |
Rendimiento fuera de distribución
Medir el rendimiento en conjuntos de datos OOD ayuda a evaluar la generalización. En este experimento, construimos mapas de etiquetas (asignaciones de etiquetas entre conjuntos de datos) a partir de JFT a ImageNet y también de ImageNet a diferentes conjuntos de datos fuera de distribución como ObjectNet (resultados después del entrenamiento previo con estos datos que se muestran en la curva izquierda a continuación). Luego, los modelos se ajustan por completo en ImageNet.
Observamos que escalar Vision Transformers aumenta el rendimiento de OOD: aunque la precisión de ImageNet se satura, vemos un aumento significativo en ObjectNet de ViT-e a ViT-22B (que se muestra con los tres puntos naranjas en la esquina superior derecha).
 |
| Aunque la precisión de ImageNet se satura, observamos un aumento significativo en el rendimiento de ObjectNet de ViT-e/14 a ViT-22B. |
Sonda lineal
La sonda lineal es una técnica en la que se entrena una sola capa lineal sobre un modelo congelado. En comparación con el ajuste fino completo, esto es mucho más barato de entrenar y más fácil de configurar. Observamos que el rendimiento de la sonda lineal de ViT-22B se aproxima al del ajuste fino completo de última generación de modelos más pequeños utilizando imágenes de alta resolución (el entrenamiento con una resolución más alta es generalmente mucho más costoso, pero para muchas tareas produce Mejores resultados). Estos son los resultados de una sonda lineal entrenada en el conjunto de datos de ImageNet y evaluada en el conjunto de datos de validación de ImageNet y otros conjuntos de datos de OOD ImageNet.
Destilación
El conocimiento del modelo más grande se puede transferir a un modelo más pequeño usando el método de destilación. Esto es útil ya que los modelos grandes son más lentos y más caros de usar. Descubrimos que el conocimiento de ViT-22B se puede transferir a modelos más pequeños como ViT-B/16 y ViT-L/16, logrando un nuevo estado del arte en ImageNet para esos tamaños de modelo.
Equidad y sesgo
Los modelos de ML pueden ser susceptibles a sesgos injustos no deseados, como detectar correlaciones falsas (medidas mediante la paridad demográfica) o tener brechas de desempeño entre los subgrupos. Mostramos que aumentar el tamaño ayuda a mitigar tales problemas.
En primer lugar, la escala ofrece una frontera de compensación más favorable: el rendimiento mejora con la escala, incluso cuando el modelo se procesa posteriormente después del entrenamiento para controlar su nivel de paridad demográfica por debajo de un nivel prescrito y tolerable. Es importante destacar que esto es válido no solo cuando el rendimiento se mide en términos de precisión, sino también en otras métricas, como la calibración, que es una medida estadística de la veracidad de las probabilidades estimadas del modelo. En segundo lugar, la clasificación de todos los subgrupos tiende a mejorar con la escala, como se demuestra a continuación. En tercer lugar, ViT-22B reduce la brecha de rendimiento entre los subgrupos.
 |
 |
| Arriba: Precisión para cada subgrupo en celeb antes de eliminar el sesgo. Abajo: El eje Y muestra la diferencia absoluta en el desempeño de los dos subgrupos específicos resaltados en este ejemplo: mujeres y hombres. ViT-22B tiene una pequeña brecha en el rendimiento en comparación con las arquitecturas ViT más pequeñas. |
Conclusiones
Hemos presentado ViT-22B, actualmente el modelo de transformador de visión más grande con 22 mil millones de parámetros. Con cambios pequeños pero críticos en la arquitectura original, logramos una excelente utilización del hardware y estabilidad en el entrenamiento, lo que generó un modelo que avanza el estado del arte en varios puntos de referencia. Se puede lograr un gran rendimiento utilizando el modelo congelado para producir incrustaciones y luego formar capas finas encima. Nuestras evaluaciones muestran además que ViT-22B muestra mayores similitudes con la percepción visual humana en lo que respecta al sesgo de forma y textura, y ofrece beneficios en equidad y robustez, en comparación con los modelos existentes.
Agradecimientos
Este es un trabajo conjunto de Mostafa Dehghani, Josip Djolonga, Basil Mustafa, Piotr Padlewski, Jonathan Heek, Justin Gilmer, Andreas Steiner, Mathilde Caron, Robert Geirhos, Ibrahim Alabdulmohsin, Rodolphe Jenatton, Lucas Beyer, Michael Tschannen, Anurag Arnab, Xiao Wang , Carlos Riquelme, Matthias Minderer, Joan Puigcerver, Utku Evci, Manoj Kumar, Sjoerd van Steenkiste, Gamaleldin Fathy, Elsayed Aravindh Mahendran, Fisher Yu, Avital Oliver, Fantine Huot, Jasmijn Bastings, Mark Patrick Collier, Alexey Gritsenko, Vighnesh Birodkar, Cristina Vasconcelos, Yi Tay, Thomas Mensink, Alexander Kolesnikov, Filip Pavetić, Dustin Tran, Thomas Kipf, Mario Lučić, Xiaohua Zhai, Daniel Keysers Jeremiah Harmsen y Neil Houlsby
Nos gustaría agradecer a Jasper Uijlings, Jeremy Cohen, Arushi Goel, Radu Soricut, Xingyi Zhou, Lluis Castrejon, Adam Paszke, Joelle Barral, Federico Lebron, Blake Hechtman y Peter Hawkins. Su experiencia y apoyo inquebrantable jugaron un papel crucial en la realización de este documento. También reconocemos la colaboración y la dedicación de los talentosos investigadores e ingenieros de Google Research.
1Nota: ViT-22B tiene una utilización de FLOP del modelo (MFU) del 54,9 %, mientras que PaLM informó un 46,2 % de MFU y medimos un 44,0 % de MFU para ViT-e en el mismo hardware. ↩
