Resumo Técnico: Leis de Escalonamento para Transformers de Reconstrução de Máscara em Transcriptômica de Célula Única

Definição do Problema

Embora as leis de escalonamento neural — relações de lei de potência entre perda, tamanho do modelo e volume de dados — sejam bem documentadas para transformers de linguagem e visão, sua existência e comportamento na genômica de célula única permanecem amplamente inexplorados. A comunidade de célula única desenvolveu grandes modelos de fundação (ex: scVI, scGPT, Geneformer), mas estes são tipicamente avaliados em tarefas downstream (anotação, correção de lote) em vez da relação fundamental entre a perda de pré-treinamento e a escala do modelo. Sem leis de escalonamento estabelecidas, os praticantes não podem determinar se os gargalos de desempenho devem-se à insuficiência de capacidade do modelo, ao volume de dados ou ao ruído intrínseco biológico/técnico. Além disso, o piso assintótico da perda nas leis de escalonamento oferece um potencial limite superior empírico para o erro redutível, o que poderia estar relacionado à entropia intrínseca dos dados de expressão gênica.

Metodologia

Os autores apresentam o primeiro estudo sistemático do comportamento de escalonamento para transformers de reconstrução de máscara treinados em dados de sequenciamento de RNA de célula única (scRNA-seq).

• Fonte de Dados: Os dados são extraídos do CELLxGENE Census. O experimento principal utiliza um subconjunto padronizado de 200.000 células e 512 genes altamente variáveis (HVGs). Análises adicionais utilizam subconjuntos de vocabulário fixo com contagens de células variando ($D \in \{25k, 50k, 100k, 200k\}$).
• Pré-processamento: O pipeline envolve normalização pelo tamanho da biblioteca (soma alvo de $10^4$), transformação log(1+x) e divisão estratificada em conjuntos de treinamento (90%), validação (5%) e teste (5%).
• Arquitetura: É empregado um codificador transformer permutação-invariante, omitindo codificações posicionais, uma vez que a ordem dos genes carece de significado biológico. As entradas consistem em embeddings de identidade gênica, valores de expressão observados (ou um token de máscara aprendido) e uma máscara binária. O modelo utiliza normalização pré-camada (Pre-LN) e uma cabeça de predição linear para reconstruir os valores de expressão mascarados.
• Objetivo de Treinamento: O modelo é treinado usando um objetivo de reconstrução de máscara, minimizando o Erro Quadrático Médio (MSE) entre os valores de expressão previstos e os reais nas posições mascaradas (taxa de mascaramento de 15%).
• Desenho Experimental:
  • Varredura de Parâmetros: Seis tamanhos de modelo (do XXS ao L) abrangendo de 533 a $1,0 \times 10^8$ parâmetros são treinados no dataset padronizado $V=512, D=200k$. Cada tamanho é executado com três sementes aleatórias por 60.000 passos.
  • Varredura de Tamanho de Dados: Uma análise de tamanho de dataset fixo (modelo de tamanho $M$) variando o tamanho do dataset ($D$) enquanto mantém o número total de passos de treinamento constante.
  • Diagnósticos: Verificações de heterogeneidade de resíduos entre genes e calibrações de Log-Verossimilhança Negativa (NLL) Gaussiana direta em pequena escala.
• Formulação da Lei de Escalonamento: A relação entre o MSE de validação ($L$) e a contagem de parâmetros ($P$) é modelada como $L(P) = aP^{-\alpha} + c$, onde $c$ representa o piso de perda irredutível. Os ajustes são realizados via busca em grade sobre $c$ e mínimos quadrados ordinários na forma linearizada.

Principais Contribuições

1. Primeiro Estudo Sistemático: Este trabalho fornece a primeira investigação sistemática das leis de escalonamento neural para pré-treinamento de reconstrução de máscara em scRNA-seq.
2. Escalonamento Demonstrado: Demonstra que o escalonamento de lei de potência emerge em uma varredura padronizada de $V=512, D=200k$, com a avaliação em dados não vistos seguindo a mesma tendência qualitativa da validação.
3. Análise Abrangente: O estudo inclui análises adicionais de $V$ fixo, incluindo uma varredza de tamanho de dados, um diagnóstico de heterogeneidade de resíduos entre genes e uma calibração de NLL Gaussiana direta em pequena escala.
4. Estimativa Preliminar de Entropia: Fornece uma estimativa preliminar do piso de incerteza assintótico, interpretando-o em unidades de informação teórica sob suposições específicas.
5. Roteiro para o Rigor: Identifica os ingredientes experimentais necessários (varreduras de $V/D$ correspondentes, objetivos baseados em verossimilhança, contabilidade de computação) para transitar de curvas de perda-vs-parâmetros para estimativas rigorosas de entropia transcriptômica.

Resultados

• Escalonamento de Parâmetros: Na varredura padronizada, o MSE de validação diminui constantemente do modelo XXS para o M. O modelo L ($\sim 10^8$ parâmetros) não mostra melhoria sobre o modelo M ($\sim 1,9 \times 10^7$ parâmetros), indicando que a curva de perda começou a achatar.
• Expoentes de Escalonamento: O expoente de lei de potência ajustado é $\alpha \approx 0,266$ (validação) e $\alpha \approx 0,268$ (teste), com um $R^2$ de aproximadamente 0,86. Este é substancialmente maior do que os expoentes reportados para modelos de linguagem autorregressivos ($\alpha \approx 0,076$), mas qualitativamente consistente com a existência de escalonamento de lei de potência.
• Piso Irredutível: O piso de perda irredutível $c$ é estimado em $\approx 1,44$ (MSE de validação) e $\approx 1,46$ (MSE de teste).
• Efeitos do Tamanho de Dados: Uma varredura de tamanho de dados com $V$ fixo (variando $D$ com passos fixos) produz um ajuste de validação quase plano ($\beta \approx -0,002$), sugerindo que, em passos de treinamento fixos, datasets maiores são vistos por menos épocas efetivas, impedindo que uma lei $L(D)$ limpa emerja neste setup específico.
• Heterogeneidade de Resíduos: Os erros residuais são fortemente dependentes dos genes. A correlação entre a expressão média e o MSE por gene é 0,892, e o coeficiente de variação entre os genes é 0,730, indicando que a suposição Gaussiana homocedástica é imperfeita.
• Estimativa de Entropia: Sob uma aproximação Gaussiana homocedástica, o piso assintótico corresponde a aproximadamente 2,30 bits de incerteza irredutível por posição de gene mascarado. Esta estimativa é derivada tanto do piso de MSE quanto de um NLL Gaussiano convertido, embora os autores notem que estes não são totalmente independentes devido à derivação post-hoc.

Significância e Alegações

O artigo afirma estabelecer um framework quantitativo para raciocinar sobre o comportamento de escalonamento de modelos de fundação de célula única. Sua principal significância reside em:

• Orientação Prática: Sugere que para datasets de $\sim 200.000$ células, modelos na casa de dezenas de milhões de parâmetros são suficientes para capturar a estrutura aprendível, com retornos decrescentes além de $\sim 20M$ de parâmetros no setup atual.
• Clareza Metodológica: Destaca que o volume de dados, o orçamento de otimização e a interpretação da perda interagem de forma complexa; uma análise de lei de dados limpa requer orçamentos de otimização correspondentes entre os tamanhos de dataset, o que não foi alcançado na varredura de passos fixos.
• Insight Biológico: O piso irredutível estimado ($\sim 2,3$ bits) é apresentado não como uma constante biológica universal, mas como uma quantidade específica da tarefa que reflete os efeitos combinados de ruído biológico, ruído técnico e perda de informação do pipeline de pré-processamento específico (transformação log, seleção de HVG).
• Direção Futura: O trabalho delineia que estimativas rigorosas de entropia transcriptômica exigem trabalhos futuros envolvendo matrizes experimentais de $V/D/C$ correspondentes, objetivos nativos baseados em verossimilhança (predizendo $\mu$ e $\sigma^2$) e contabilidade abrangente de computação.

Os autores mantêm um tom modesto, enfatizando que a figura de 2,3 bits é um benchmark preliminar sob suposições específicas e que o estudo atual ainda não separa definitivamente os pisos limitados por dados dos limitados por modelos.