From Words to Amino Acids: Does the Curse of Depth Persist?

1. Planteamiento del Problema

Análisis recientes de Modelos de Lenguaje Grandes (LLM) han identificado el "Flujo de la Profundidad" (Curse of Depth), un fenómeno en el que las capas más profundas en transformadores autoregresivos contribuyen mínimamente a las predicciones finales, mientras que la mayoría de los cálculos relevantes para la tarea ocurren en capas tempranas o intermedias. Las capas posteriores sirven principalmente para el refinamiento incremental de la distribución de salida, en lugar de generar nueva información.

La pregunta central de este trabajo es si esta ineficiencia de profundidad también existe en los Modelos de Lenguaje de Proteínas (PLM). A diferencia de los LLM, los PLM utilizan frecuentemente:

• Objetivos no autoregresivos: Como el Modelado de Lenguaje enmascarado (MLM) y la Difusión Discreta.
• Entradas multimodales: La inclusión de estructura y función de proteínas junto con la secuencia.
• Paradigmas de escalado diferentes: Incluyendo arquitecturas de Mezcla de Expertos (MoE).

Los autores investigan si el "Flujo de la Profundidad" es una propiedad universal de los modelos profundos basados en transformadores o si las restricciones únicas de la modelación de proteínas (estructura biológica, restricciones evolutivas) influyen en cómo se utiliza la profundidad.

2. Metodología

Los autores realizaron un análisis sistemático de la profundidad en 7 familias de PLM con 26 variantes de modelos, que cubren un amplio espectro de tamaños (8M a 15B parámetros) y objetivos de entrenamiento (Autoregresivo, Enmascarado, Difusión y Multimodal).

Utilizaron un marco unificado que combina tres tipos de mediciones:

A. Análisis de Intervención (Salto de Capas)

Adaptando el método "skiplayer" de la investigación en LLM, los autores midieron el efecto propagado de saltar capas específicas:

• Modelos Autoregresivos: Se saltaron capas para tokens tempranos y se midió el impacto en las predicciones de tokens futuros.
• Modelos Enmascarados/Difusión: Se saltaron capas en un subconjunto de tokens enmascarados y se midió el impacto en los tokens enmascarados restantes no intervenidos.
• Métrica: El cambio relativo en los estados ocultos posteriores y las normas de salida.

B. Análisis de Similitud de Representación

• Similitud Capa a Capa: Cálculo de la CKA Lineal (Alineación de Kernel Centrado) entre las representaciones de diferentes capas para cuantificar cómo cambian las información a medida que se propagan a través de la red.
• Similitud Cruzada de Modalidades: Para modelos multimodales (ESM3, DPLM2), calcularon la Similitud Coseno entre representaciones basadas en secuencia y basadas en estructura en capas correspondientes, para entender cómo se alinean y diferencian las modalidades.

C. Evolución de la Distribución de Salida

Utilizando LogitLens y TunedLens, los autores extrajeron la distribución implícita de predicción de tokens en cada capa.

• Métricas: Medición de la Divergencia KL entre las distribuciones de las capas medias y la salida final, así como la Superposición Top-1 (coincidencia en el token más probable).

D. Evaluación del Utilidad en Tareas Posteriores

Para vincular las mediciones intrínsecas con el rendimiento práctico, evaluaron el rendimiento por capas en ProteinGym, un benchmark a gran escala para la predicción Zero-Shot de efectos de mutaciones (217 ensayos de Escaneo Mutacional Profundo). midieron la correlación de rango de Spearman entre las predicciones realizadas por capas de Salida Temprana (Early-Exit) y los datos experimentales.

3. Contribuciones Principales

1. Primer estudio sistemático de profundidad en PLM: El artículo ofrece el primer análisis exhaustivo del uso de la profundidad a través de diversas familias de PLM, objetivos y modalidades.
2. Marco unificado: Desarrollo de una metodología consistente para comparar modelos autoregresivos, enmascarados, basados en difusión y multimodales, lo que demuestra que la ineficiencia de profundidad no se limita a la predicción del siguiente token.
3. Descubrimiento de una ineficiencia de profundidad generalizada: Se demostró que el "Flujo de la Profundidad" es una característica común de los PLM modernos, donde un subconjunto de capas (a menudo las intermedias) capturan la mayoría de las señales relevantes para tareas posteriores, mientras que las capas posteriores ofrecen rendimientos marginales decrecientes.
4. Hallazgos específicos de modalidad: Revelación de perfiles de profundidad diferentes en modelos multimodales (especialmente ESM3) y explicación de su comportamiento mediante dinámicas de alineación cruzada de modalidades.

4. Resultados Principales

A. El Flujo de la Profundidad persiste en los PLM

A través de todas las familias de modelos (ESM2, DPLM, ProGen, etc.), los autores observaron un patrón consistente:

• Capas tempranas a intermedias: Saltar estas capas causa perturbaciones significativas en los cálculos posteriores y reduce drásticamente el rendimiento en tareas posteriores.
• Capas tardías: Saltar un gran porcentaje de las últimas capas (por ejemplo, las últimas 30–50 %) conduce a cambios mínimos en las representaciones posteriores y solo a pérdidas marginales en el rendimiento de ProteinGym.
• Efecto de escalado: A medida que aumenta el tamaño del modelo, esta ineficiencia se vuelve más pronunciada. En modelos grandes (por ejemplo, ESM2-15B), el "codo" donde el rendimiento se satura ocurre antes en la profundidad relativa.

B. Refinamiento de la Distribución de Salida

• Análisis TunedLens: La Divergencia KL entre las distribuciones de las capas medias y la salida final disminuye constantemente con la profundidad.
• Estabilidad: La predicción Top-1 a menudo se estabiliza mucho antes de la última capa. Las capas posteriores refinan principalmente la confianza (masa de probabilidad) de la predicción, en lugar de cambiar el token predicho en sí.

C. Modelos Multimodales y de Difusión

• Generalización: El patrón de ineficiencia de profundidad se aplica a modelos de difusión (DPLM, DPLM2) y modelos multimodales (ESM3, DPLM2) al procesar entradas puramente basadas en estructura o combinaciones de secuencia-estructura.
• Anomalía de ESM3: ESM3 muestra un perfil de profundidad pronunciado. A diferencia de otros modelos donde las capas tempranas son críticas, ESM3 presenta efectos propagados débiles en las capas tempranas, con contribuciones más fuertes en capas medias a tardías.
  • Explicación: El análisis de similitud cruzada de modalidades reveló que las primeras ~60 % de las capas de ESM3 sirven para la alineación de representaciones de secuencia y estructura en un espacio común. Una vez alineadas, las capas posteriores realizan el cálculo específico de la tarea.

D. Arquitecturas Dispersas (MoE)

• ProGen3: Como modelo disperso de Mezcla de Expertos, ProGen3 mostró un "Flujo de la Profundidad" menos pronunciado en comparación con los modelos densos. La separación entre capas de alto y bajo efecto no se agudizó tanto con la escalabilidad, lo que sugiere que la dispersión puede mitigar la ineficiencia de profundidad.

5. Significado e Implicaciones

• Eficiencia Arquitectónica: Los resultados indican que los PLM profundos actuales son computacionalmente ineficientes. Una parte significativa de los costos de inferencia (FLOPs) se gasta en capas que contribuyen poco a la salida final.
• Futuras Direcciones: El artículo motiva el desarrollo de:
  • Cálculo Adaptativo: Mecanismos de Salida Temprana o políticas de profundidad dinámicas que salten capas redundantes durante la inferencia.
  • Arquitecturas Crecientes/Bucle: Modelos que expanden la profundidad durante el entrenamiento o reutilizan capas (bucle) para mejorar la utilización.
  • Modelos Dispersos: El éxito de ProGen3 sugiere que las arquitecturas MoE podrían ser una vía más eficiente para escalar modelos de proteínas.
• Diseño Multimodal: El comportamiento diferente de ESM3 subraya que, en entornos multimodales, la profundidad temprana puede consumirse por la alineación de modalidades en lugar del cálculo de tareas, lo que sugiere que los modelos futuros podrían beneficiarse de módulos especializados de "alineación" seguidos de pilas de "cálculo" específicos de la tarea.

En resumen, el "Flujo de la Profundidad" no es exclusivo de los LLM autoregresivos, sino una característica fundamental de los modernos modelos de proteínas profundos basados en transformadores, independientemente de su objetivo de entrenamiento o modalidad. Este hallazgo exige un cambio de paradigma hacia arquitecturas de IA de proteínas más eficientes en profundidad.