Semantic search for 100M+ galaxy images using AI-generated captions

Resumo Técnico: Busca Semântica para mais de 100 Milhões de Imagens de Galáxias Usando Legendas Geradas por IA

Problema
A comunidade astronômica enfrenta um gargalo crítico na exploração dos bilhões de imagens de galáxias geradas por telescópios modernos. Os métodos tradicionais de descoberta dependem de campanhas lentas de rotulagem manual (ex: Galaxy Zoo) ou modelos de aprendizado profundo supervisionados, limitados a conjuntos de perguntas fixos. Essas abordagens não conseguem escalar para conjuntos de dados futuros (ex: Euclid, LSIST) e carecem da flexibilidade necessária para suportar consultas de vocabulário aberto. Além disso, os métodos existentes de busca por similaridade autossupervisionada recuperam imagens visualmente semelhantes em vez de imagens que contenham fenômenos físicos específicos, tornando-os insuficientes para a descoberta científica direcionada.

Metodologia
Os autores propõem o AION-Search, um pipeline projetado para criar um mecanismo de busca semântica a partir de dados de imagem completamente não rotulados. A metodologia procede em três estágios:

1. Geração de Legendas por VLM: Os autores testaram 14 Modelos de Linguagem-Visão (VLMs) contra anotações humanas do catálogo Galaxy Zoo-DECaLS. Eles avaliaram os modelos com base em sua capacidade de gerar descrições de texto livre que pudessem ser analisadas em respostas de árvore de decisão. O GPT-4.1-mini foi selecionado como o modelo ideal, oferecendo o melhor equilíbrio entre precisão (taxa de correspondência de 50,1% nas perguntas do Galaxy Zoo) e custo de API. Usando este modelo, a equipe gerou descrições detalhadas para 300.000 imagens de galáxias selecionadas aleatoriamente do Legacy Survey e do Hyper Suprime-Cam (HSC). Essas descrições foram posteriormente resumidas em sentenças únicas para reduzir o ruído.
2. Alinhamento Contrastivo: Para permitir uma busca escalável sem a necessidade de legendar cada imagem do catálogo, os autores treinaram um modelo de aprendizado contrastivo. Eles utilizaram o AION-1, um modelo de fundação astronômica pré-treinado em mais de 200 milhões de observações, como o codificador de imagem. As resumos de texto foram incorporados usando o text-embedding-3-large. Uma cabeça de projeção MLP residual de quatro camadas foi treinada para alinhar os embeddings de imagem e texto em um espaço semântico compartilhado de 1024 dimensões usando uma perda de entropia cruzada simétrica (InfoNCE). Os pesos subjacentes do AION-1 e do codificador de texto foram congelados durante este processo.
3. Reclassificação (Re-ranking) por VLM: Para melhorar a precisão da recuperação de fenômenos raros, os autores introduziram um estágio de reclassificação pós-busca. Após uma recuperação inicial dos 1.000 principais candidatos via similaridade de cosseno, um VLM (GPT-4.1) pontua cada candidato com base na consulta do usuário. Este passo aproveita o aumento de computação no tempo de inferência para refinar a ordenação dos resultados.

Principais Contribuições

• Busca Semântica Escalável: O sistema permite a busca semântica flexível e de vocabulário aberto para mais de 100 milhões de imagens de galáxias (especificamente galáxias do Legacy Survey DR10) sem exigir anotação humana para todo o conjunto de dados.
• Desempenho Zero-Shot: O modelo atinge um desempenho zero-shot de estado da arte na localização de fenômenos raros (galáxias espirais, fusões e lentes gravitacionais), apesar de ter sido treinado em imagens selecionadas aleatoriamente, sem curadoria deliberada para essas classes específicas.
• Estrutura de Reclassificação: A introdução de um método de reclassificação baseado em VLM quase dobra o recall para os alvos mais desafiadores (lentes gravitacionais) nos 100 principais resultados.
• Infraestrutura Pública: Os autores disponibilizaram o código, os pesos do modelo treinado, os embeddings pré-computados e uma aplicação web pública (hospedada no Hugging Face Spaces) que suporta modos de busca de texto, imagem e híbrida.

Resultados

• Desempenho de Recuperação: O AION-Search supera significativamente as linhas de base de similaridade (incluindo variantes do AIO-1, AstroCLIP e DINOv2). Na métrica nDCG@10, o AION-Search alcançou pontuações de 0,941 para espirais, 0,554 para fusões e 0,173 para lentes gravitacionais, comparado a bases que variam de 0,015 a 0,643.
• Eficácia da Reclassificação: A aplicação da reclassificação com GPT-4.1 com média de 5 amostras melhorou as pontuações nDCG@10 em todas as categorias. Para lentes gravitacionais, a reclassificação aumentou o número de lentes confirmadas encontradas nos 100 principais resultados de 7 para 13.
• Classificação Zero-Shot: No benchmark Galaxy10, o AION-Search alcançou 53,3% de precisão na classificação zero-shot, excedendo ligeiramente o desempenho do VLM (GPT-4.1-mini) que gerou seus dados de treinamento. Quando uma cabeça MLP supervisionada foi adicionada aos embeddings congelados do AION-Search, a precisão atingiu 84,1%, igualando o desempenho do modelo base AION-1, indicando que o alinhamento contrastivo preservou a informação morfológica.
• Descoberta Científica: Utilizando o sistema, os autores identificaram 36 novos candidatos a fluxos estelares extragalácticos através de consultas híbridas iterativas de imagem e texto. Esses candidatos foram confirmados como novos através do cruzamento com 11 catálogos existentes de características de maré.

Significância e Alegações
O artigo afirma que este trabalho fornece o primeiro exemplo na astrofísica do uso de legendas geradas por VLM para gerar conjuntos de treinamento para recuperação semântica. Os autores afirmam que sua abordagem demonstra que grandes arquivos de imagens não rotuladas podem se tornar semanticamente pesquisáveis sem supervisão humana. Eles posicionam isso como um framework generalizável aplicável a outros campos com grandes arquivos de imagens não rotuladas, como observação da Terra, microscopia e ciência dos materiais. O trabalho é apresentado como um passo em direção à visão de "um país de gênios em um datacenter", onde agentes de IA podem acelerar a descoberta científica ao analisar vastos conjuntos de dados que excedem a capacidade de inspeção humana. Os autores mantêm-se modestos quanto às limitações, reconhecendo que as descrições de VLM podem conter alucinações e que o modelo herda vieses do VLM de treinamento, mas argumentam que essas anotações imperfeitas ainda fornecem um sinal suficiente para um aprendizado contrastivo eficaz.