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

Résumé technique : Le guide de la galaxie pour le tokenizer

Énoncé du problème
Bien que les modèles de fondation basés sur les transformeurs soient devenus un paradigme de premier plan pour l'extraction d'informations scientifiques à partir de données astronomiques, l'impact spécifique des stratégies de tokenisation sur les représentations apprises reste mal compris. Dans les domaines scientifiques comme l'astronomie, où les observations sont régies par des processus physiques connus et accompagnées de vérités de terrain mesurées indépendamment (ex. : décalage vers le rouge, masse stellaire), il est crucial d'évaluer comment la conversion des observations brutes en séquences de tokens façonne la capacité du modèle à réaliser l'inférence scientifique. La littérature actuelle se concentre souvent sur les choix architecturaux, laissant une lacune dans la compréhension de savoir si la fidélité de la reconstruction est corrélée à l'accessibilité de l'information physique au sein de l'espace latent.

Méthodologie
Les auteurs mènent une comparaison systématique de quatre stratégies de tokenisation distinctes au sein d'un cadre unifié et contrôlé utilisant l'architecture AstroPT, un transformeur de type décodeur seul entraîné de manière autorégressive sur des données astronomiques. L'étude utilise 640 000 images de galaxies issues du relevé DESI Legacy Survey (DR8), traitées sous forme de vignettes de 256×256 pixels.

Les quatre stratégies de tokenisation évaluées sont :

1. Affine : Une ligne de base minimale utilisant une simple projection linéaire des patchs d'images vers l'espace d'embedding.
2. AIM (basée sur MLP) : Remplace la projection linéaire par un perceptron multicouche (MLP) pour introduire une non-linéarité, suivant le cadre original d'AstroPT.
3. JetFormer : Une approche basée sur le flux utilisant un flot normalisant inversible pour produire des tokens continus. Il est entraîné de bout en bout pour maximiser la vraisemblance de l'image sans encodeur pré-entraîné distinct.
4. VQ-VAE : Une approche discrète qui quantifie les patchs d'images en un codebook de tokens visuels appris. L'encodeur est pré-entraîné sur des images de galaxies de la même distribution et est gelé lors de l'entraînement principal d'AstroPT.

Protocole d'évaluation
L'étude emploie deux axes d'évaluation principaux :

• Prédiction de propriétés physiques : Des sondes linéaires et MLP sont entraînées sur les embeddings extraits de la couche centrale des modèles AstroPT pour prédire les propriétés physiques de 167 000 galaxies mises de côté. Ces propriétés incluent les magnitudes photométriques, les couleurs, les décalages vers le rouge photométriques et spectroscopiques, le taux de formation stellaire spécifique (sSFR), la masse stellaire et les indicateurs morphologiques.
• Fidélité de la reconstruction : La qualité de la reconstruction d'image est évaluée à l'aide de l'SSIM et du PSNR sur 5 000 galaxies aléatoires de l'ensemble de test.
• Vérification du contenu informationnel : Pour distinguer la perte d'information de l'information inaccessible, les auteurs sondent les images reconstruites elles-mêmes en utilisant les caractéristiques d'un ResNet-50 pré-entraîné sur ImageNet, isolant ainsi le contenu informationnel des données de pixels de la géométrie de l'embedding.

Résultats clés
Les résultats révèlent un découplage significatif entre la fidélité de la reconstruction et la qualité des représentations pour l'inférence physique :

• JetFormer : Atteint la plus haute fidélité de reconstruction (PSNR moyen ~31,11 dB, SSIM ~0,76), préservant les détails spatiaux fins et les composantes de faible brillance de surface. Cependant, ses embeddings sont les moins accessibles linéairement pour la prédiction des propriétés physiques. Le sondage des images reconstruites confirme que JetFormer conserve la quasi-totalité de l'information originale, mais cette information est « éparpillée » dans l'espace latent d'une manière que les sondes simples ne peuvent facilement récupérer.
• VQ-VAE : Produit la meilleure performance pour la prédiction des propriétés physiques, particulièrement avec les sondes linéaires. Le goulot d'étranglement discret organise l'information physique de manière plus accessible linéairement. Cependant, il obtient une fidélité de reconstruction plus faible (PSNR moyen ~23,57 dB) et perd systématiquement les détails fins (ex. : morphologie, photométrie absolue), bien qu'il préserve les indices de couleur.
• Affine et AIM : Ces méthodes se comportent de manière presque identique pour presque toutes les propriétés, suggérant que la tête MLP dans AIM apporte un bénéfice négligeable par rapport à une simple projection linéaire lorsque le backbone du transformeur est suffisamment expressif. Elles préservent mieux l'information morphologique localisée que JetFormer, mais sont généralement moins performantes que VQ-VAE sur les tâches de régression.
• Compromis : L'étude identifie un compromis fondamental : les représentations douces, denses et complètes en informations (JetFormer) favorisent une haute fidélité de reconstruction, tandis que les stratégies de tokenisation par patchs rigides (VQ-VAE) favorisent des représentations plus directement interprétables et accessibles pour l'inférence physique.

Signification et affirmations
L'article soutient que la tokenisation n'est pas une étape de prétraitement neutre, mais un choix de conception critique qui dicte quelle information est apprise et comment elle est organisée. Les auteurs affirment que la qualité de la reconstruction n'est pas un indicateur fiable de l'utilité d'une représentation pour l'inférence scientifique.

En ancrant l'évaluation dans des quantités physiques mesurées indépendamment, l'étude sert de démonstration de la manière dont les données scientifiques peuvent fournir des benchmarks objectifs et non circulaires pour les modèles de fondation. Les auteurs postulent que le choix du tokenizer doit être dicté par la tâche en aval spécifique : le VQ-VAE est préférable pour l'inférence photométrique ou spectroscopique, tandis que le JetFormer est supérieur pour les tâches de génération ou de reconstruction d'images. Ce travail souligne le potentiel de l'astronomie pour construire des benchmarks interprétables, notant que des cadres similaires pourraient être précieux dans d'autres domaines où la vérité physique de terrain est disponible.

Limites
Les auteurs reconnaissent que la comparaison n'est pas entièrement ablative en raison de la nécessité de différents procédures d'entraînement pour chaque tokenizer (ex. : pré-entraînement du VQ-VAE vs entraînement de bout en bout pour JetFormer). De plus, le recours aux sondes linéaires et MLP peut sous-estimer l'information présente dans les représentations, bien que le sondage des images reconstruites aide à atténuer cela pour les cas extrêmes. L'étude n'explore pas non plus la mise à l'échelle avec la taille du jeu de données ou la capacité du modèle.