Resumen Técnico: Leyes de Escalamiento para Transformers de Reconstrucción Enmascarada en Transcriptómica de Célula Única

Planteamiento del Problema

Si bien las leyes de escalamiento neuronal —relaciones de ley de potencia entre la pérdida, el tamaño del modelo y el volumen de datos— están bien documentadas para los transformers de lenguaje y visión, su existencia y comportamiento en la genómica de célula única permanecen en gran medida inexplorados. La comunidad de célula única ha desarrollado modelos fundacionales de gran escala (p. ej., scVI, scGPT, Geneformer), pero estos se evalúan típicamente en tareas de seguimiento (anotación, corrección de lote) en lugar de en la relación fundamental entre la pérdida de preentrenamiento y la escala del modelo. Sin leyes de escalamiento establecidas, los profesionales no pueden determinar si los cuellos de botella en el rendimiento se deben a una capacidad de modelo insuficiente, al volumen de datos o al ruido intrínseco biológico/técnico. Además, el piso de pérdida asintótica de las leyes de escalamiento ofrece un posible límite superior empírico para el error reducible, que podría relacionarse con la entropía intrínseca de los datos de expresión génica.

Metodología

Los autores presentan el primer estudio sistemático del comportamiento de escalamiento para transformers de reconstrucción enmascarada entrenados con datos de secuenciación de ARN de célula única (scRNA-seq).

• Fuente de Datos: Los datos se extraen del CELLxGENE Census. El experimento principal utiliza un subconjunto estandarizado de 200,000 células y 512 genes altamente variables (HVGs). Los análisis adicionales utilizan subconjuntos de vocabulario fijo con cantidades de células variables ($D \in \{25k, 50k, 100k, 200k\}$).
• Preprocesamiento: El flujo de trabajo implica normalización del tamaño de la librería (suma objetivo de $10^4$), transformación log(1+x) y división estratificada en conjuntos de entrenamiento (90%), validación (5%) y prueba (5%).
• Arquitectura: Se emplea un codificador transformer invariante a la permutación, omitiendo codificaciones posicionales ya que el orden de los genes carece de significado biológico. Las entradas consisten en embeddings de identidad de genes, valores de expresión observados (o un token de máscara aprendido) y una máscara binaria. El modelo utiliza normalización de capa previa (Pre-LN) y una cabeza de predicción lineal para reconstruir los valores de los genes enmascarados.
• Objetivo de Entrenamiento: El modelo se entrena utilizando un objetivo de reconstrucción enmascarada, minimizando el Error Cuadrático Medio (MSE) entre los valores de expresión predichos y los reales en las posiciones enmascaradas (tasa de enmascaramiento del 15%).
• Diseño Experimental:
  • Barrido de Parámetros: Se entrenan seis tamaños de modelo (desde XXS hasta L) que abarcan desde 533 hasta $1.0 \times 10^8$ parámetros en el conjunto de datos estandarizado de $V=512, D=200k$. Cada tamaño se ejecuta con tres semillas aleatorias durante 60,000 pasos.
  • Barrido de Tamaño de Datos: Un análisis de tamaño de datos fijo (modelo de tamaño $M$) variando el tamaño del conjunto de datos ($D$) manteniendo constante el número total de pasos de entrenamiento.
  • Diagnósticos: Verificaciones de heterogeneidad de residuos entre genes y calibraciones directas de la Negativa Log-Verosimilitud (NLL) Gaussiana a pequeña escala.
• Formulación de la Ley de Escalamiento: La relación entre el MSE de validación ($L$) y el recuento de parámetros ($P$) se modela como $L(P) = aP^{-\alpha} + c$, donde $c$ representa el piso de pérdida irreducible. Los ajustes se realizan mediante una búsqueda de cuadrícula sobre $c$ y mínimos cuadrados ordinarios sobre la forma linealizada.

Contribuciones Clave

1. Primer Estudio Sistemático: Este trabajo proporciona la primera investigación sistemática de las leyes de escalamiento neuronal para el preentrenamiento de reconstrucción enmascarada en scRNA-seq.
2. Escalamiento Demostrado: Demuestra que emerge un escalamiento de ley de potencia en un barrido estandarizado de $V=512, D=200k$, con la evaluación en el conjunto de prueba siguiendo la misma tendencia cualitativa que la de validación.
3. Análisis Exhaustivo: El estudio incluye análisis adicionales de tamaño de datos con $V$ fijo, incluyendo un barrido de tamaño de datos, un diagnóstico de heterogeneidad de residuos entre genes y una calibración de NLL Gaussiana directa a pequeña escala.
4. Estimación Preliminar de Entropía: Proporciona una estimación preliminar del piso de incertidumbre asintótica, interpretándola en unidades de información teórica bajo supuestos específicos.
5. Hoja de Ruta para el Rigor: Identifica los ingredientes experimentales necesarios (barridos de $V/D$ emparejados, objetivos basados en verosimilitud, contabilidad de cómputo) requeridos para transicionar de curvas de pérdida vs. parámetros a estimaciones rigurosas de la entropía transcriptómica.

Resultados

• Escalamiento de Parámetros: En el barrido estandarizado, el MSE de validación disminuye constantemente desde los modelos XXS hasta los M. El modelo L ($\sim 10^8$ parámetros) no muestra mejoras sobre el modelo M ($\sim 1.9 \times 10^7$ parámetros), lo que indica que la curva de pérdida ha comenzado a aplanarse.
• Exponentes de Escalamiento: El exponente de la ley de potencia ajustado es $\alpha \approx 0.266$ (validación) y $\alpha \approx 0.268$ (prueba), con un $R^2$ de aproximadamente 0.86. Esto es sustancialmente mayor que los exponentes reportados para modelos de lenguaje autorregresivos ($\alpha \approx 0.076$), pero cualitativamente consistente con la existencia de un escalamiento de ley de potencia.
• Piso Irreducible: El piso de pérdida irreducible $c$ se estima en $\approx 1.44$ (MSE de validación) y $\approx 1.46$ (MSE de prueba).
• Efectos del Tamaño de los Datos: Un barrido de tamaño de datos con $V$ fijo (variando $D$ a pasos fijos) produce un ajuste de validación casi plano ($\beta \approx -0.002$), lo que sugiere que, con pasos de entrenamiento fijos, los conjuntos de datos más grandes se ven para menos épocas efectivas, impidiendo que emerja una ley $L(D)$ limpia en este montaje específico.
• Heterogeneidad de Residuos: Los errores residuales dependen fuertemente de los genes. La correlación entre la expresión media y el MSE por gen es de 0.892, y el coeficiente de variación entre genes es de 0.730, lo que indica que el supuesto Gaussiano homocedástico es imperfecto.
• Estimación de Entropía: Bajo una aproximación Gaussiana homocedástica, el piso asintótico corresponde a aproximadamente 2.30 bits de incertidumbre irreducible por posición de gen enmascarado. Esta estimación se deriva tanto del piso de MSE como de una conversión del piso de NLL Gaussiana, aunque los autores señalan que estos no son totalmente independientes debido a la derivación post-hoc.

Significación y Reivindicaciones

El artículo afirma establecer un marco cuantitativo para razonar sobre el comportamiento de escalamiento de los modelos fundacionales de célula única. Su principal significación radica en:

• Guía Práctica: Sugiere que para conjuntos de datos de $\sim 200,000$ células, los modelos en las decenas de millones de parámetros son suficientes para capturar la estructura aprendible, con rendimientos decrecientes más allá de $\sim 20M$ de parámetros en la configuración actual.
• Claridad Metodológica: Destaca que el volumen de datos, el presupuesto de optimización y la interpretación de la pérdida interactúan de forma compleja; un análisis limpio de la ley de datos requiere presupuestos de optimización emparejados entre tamaños de conjunto de datos, lo cual no se logró en el barrido de pasos fijos.
• Perspectiva Biológica: El piso de incertidumbre irreducible estimado ($\sim 2.3$ bits) se presenta no como una constante biológica universal, sino como una cantidad específica de la tarea que refleja los efectos combinados del ruido biológico, el ruido técnico y la pérdida de información del pipeline de preprocesamiento específico (transformación log, selección de HVG).
• Dirección Futura: El trabajo esboza que las estimaciones rigurosas de la entropía transcriptómica requieren trabajos futuros que involucren matrices experimentales de $V/D/C$ emparejadas, objetivos nativos basados en verosimilitud (predicción de $\mu$ y $\sigma^2$) y una contabilidad de cómputo exhaustiva.

Los autores mantienen un tono modesto, enfatizando que la cifra de 2.3 bits es un punto de referencia preliminar bajo supuestos específicos y que el estudio actual aún no desentraja definitivamente si los pisos están limitados por los datos o por el modelo.