The Galaxy's Guide to the Tokenizer: A Benchmark for Scientific Foundation Models

Resumo Técnico: O Guia Galáctico para o Tokenizador

Definição do Problema
Embora os modelos de fundação baseados em transformers tenham se tornado um paradigma dominante para a extração de insights científicos de dados astronômicos, o impacto específico das estratégias de tokenização nas representações aprendidas permanece pouco compreendido. Em domínios científicos como a astronomia, onde as observações são regidas por processos físicos conhecidos e acompanhadas por uma verdade fundamental (ground truth) medida independentemente (ex: redshift, massa estelar), há uma necessidade crítica de avaliar como a conversão de observações brutas em sequências de tokens molda a capacidade do modelo de realizar inferência científica. A literatura atual foca frequentemente em escolhas arquiteturais, deixando uma lacuna na compreensão de se a fidelidade de reconstrução correlaciona-se com a acessibilidade da informação física no espaço latente.

Metodologia
Os autores conduzem uma comparação sistemática de quatro estratégias distintas de tokenização dentro de um framework unificado e controlado utilizando o backbone AstroPT, um transformer decodificador (decoder-only) treinado de forma autorregressiva em dados astronômicos. O estudo utiliza 640.000 imagens de galáxias do DESI Legacy Survey (DR8), processadas em selos postais (postage stamps) de 256×256 pixels.

As quatro estratégias de tokenização avaliadas são:

1. Affine: Uma linha de base mínima que utiliza uma simples projeção linear de patches de imagem para o espaço de embedding.
2. AIM (baseada em MLP): Substitui a projeção linear por um Perceptron de Camada Multicamadas (MLP) para introduzir não-linearidade, seguindo o framework original do AstroPT.
3. JetFormer: Uma abordagem baseada em fluxo que utiliza um fluxo normalizador invertível para produzir tokens contínuos. É treinada de ponta a ponta para maximizar a verossimilhança da imagem sem um encoder pré-treinado separado.
4. VQ-VAE: Uma abordagem discreta que quantiza patches de imagem em um código de livros (codebook) de tokens visuais aprendidos. O encoder é pré-treinado em imagens de galáxias da mesma distribuição e congelado durante o treinamento principal do AstroPT.

Protocolo de Avaliação
O estudo emprega dois eixos primários de avaliação:

• Predição de Propriedades Físicas: Probes lineares e de MLP são treinados nos embeddings extraídos da camada central dos modelos AstroPT para predizer as propriedades físicas de 167.000 galáxias retidas (held-out). Estas propriedades incluem magnitudes fotométricas, cores, redshifts fotométricos e espectroscópicos, taxa de formação estelar específica (sSFR), massa estelar e indicadores morfológicos.
• Fidelidade de Reconstrução: A qualidade da reconstrução da imagem é avaliada usando SSIM e PSNR em 5.000 galáxias aleatórias do conjunto de teste.
• Verificação de Conteúdo de Informação: Para distinguir entre perda de informação e informação inacessível, os autores sondam as imagens reconstruídas utilizando features de uma ResNet-50 pré-treinada no ImageNet, isolando o conteúdo de informação dos pixels da geometria do embedding.

Principais Resultados
Os resultados revelam um desacoplamento significativo entre a fidelidade de reconstrução e a qualidade das representações para inferência física:

• JetFormer: Alcança a maior fidelidade de reconstrução (Média PSNR ~31,11 dB, SSIM ~0,76), preservando detalhes espaciais finos e componentes de baixo brilho superficial. No entanto, seus embeddings são os menos acessíveis linearmente para a predição de propriedades físicas. A sondagem das imagens reconstruídas confirma que o JetFormer retém quase toda a informação original, mas essa informação está "espalhada" no espaço latente de uma forma que sondagens simples não conseguem recuperar facilmente.
• VQ-VAE: Produz o desempenho mais forte para a predição de propriedades físicas, particularmente com probes lineares. O gargalo discreto organiza a informação física de uma maneira mais acessível linearmente. Contudo, alcança uma menor fidelidade de reconstrução (M ª PSNR ~23,57 dB) e perde sistematicamente detalhes granulares (ex: morfologia, fotometria absoluta), embora preserve índices de cor.
• Affine e AIM: Estes métodos performam de forma quase idêntica em quase todas as propriedades, sugerindo que o cabeçalho MLP no AIM adiciona um benefício negligenciável sobre uma simples projeção linear quando o backbone do transformer é suficientemente expressivo. Eles preservam melhor a informação morfológica localizada do que o JetFormer, mas geralmente performam abaixo do VQ-VAE em tarefas de regressão.
• Trade-offs: O estudo identifica um trade-off fundamental: representações suaves e densas e completas em informação (JetFormer) favorecem a alta fidelidade de reconstrução, enquanto estratégias de tokenização baseadas em patches rígidos e discretos (VQ-VAE) promovem representações que são mais diretamente interpretáveis e acessíveis para inferência física.

Significância e Alegações
O artigo argumenta que a tokenização não é um passo de pré-processamento neutro, mas uma escolha de design crítica que dita o que é aprendido e como a informação é organizada. Os autores afirmam que a qualidade da reconstrução não é um proxy confiável para a utilidade de uma representação para inferência científica.

Ao fundamentar a avaliação em quantidades físicas medidas independentemente, o estudo serve como uma demonstração de como dados científicos podem fornecer benchmarks objetivos e não circulares para modelos de fundação. Os autores postulam que a escolha do tokenizador deve ser impulsionada pela tarefa específica de jusante (downstream): o VQ-VAE é preferível para inferência fotométrica ou espectroscópica, enquanto o JetFormer é superior para tarefas de geração ou reconstrução de imagem. O trabalho destaca o potencial da astronomia como um cenário para a construção de benchmarks interpretáveis, observando que frameworks semelhantes poderiam ser valiosos em outros domínios onde a verdade física fundamental está disponível.

Limitações
Os autores reconhecem que a comparação não é inteiramente ablativa devido à necessidade de diferentes procedimentos de treinamento para cada tokenizador (ex: pré-treinar VQ-VAE vs. treinamento de ponta a ponta para JetFormer). Além disso, a dependência de probes lineares e de MLP pode subestimar a informação presente nas representações, embora a sondagem de imagens reconstruídas ajude a mitigar isso nos casos extremos. O estudo não explora o escalonamento com o tamanho do dataset ou a capacidade do modelo.