YOLO-CIANNA: Galaxy detection with deep learning in radio data: II. Winning the SKA SDC2 using a generalized 3D-YOLO network

1. O Problema

O Square Kilometre Array (SKA) e seus precursores gerarão volumes massivos de dados de rádio de alta dimensionalidade (cubos de dados hiperspectrais 3D). As pipelines de processamento de dados clássicas enfrentam dificuldades para escalar e lidar com a taxa de dados e a complexidade desses conjuntos, especialmente na detecção e caracterização de galáxias de emissão de hidrogênio neutro (HI) em meio a ruído térmico, artefatos instrumentais e interferência de radiofrequência (RFI).

O Desafio de Dados Científicos 2 (SDC2) do Observatório SKA foi criado para estimular o desenvolvimento de métodos eficientes para processar dados simulados que imitam as observações futuras do SKA. O desafio consiste em detectar e caracterizar fontes de HI em um cubo de dados principal (MAIN) de ~1 TB, contendo 233.245 fontes simuladas, com base em um catálogo de verdade (truth catalog) que foi inicialmente secreto.

2. Metodologia

Os autores generalizaram o método de aprendizado profundo YOLO-CIANNA, originalmente desenvolvido para imagens 2D (continuum), para a detecção de objetos 3D em cubos de dados hiperspectrais.

• Arquitetura da Rede (Backbone 3D CNN):

  • Foi construída uma rede neural convolucional (CNN) totalmente 3D, em vez de tratar a dimensão espectral como canais independentes 2D.
  • A arquitetura possui 23 camadas convolucionais e 3,35 milhões de pesos aprendíveis (menor que a versão 2D anterior, devido à menor diversidade morfológica das fontes HI simuladas).
  • Utiliza filtros anisotrópicos e configurações de stride para lidar com a distribuição diferente do sinal nas dimensões espaciais (céu) versus a dimensão espectral (frequência).
  • A rede processa subcubos de entrada de 64×64×256 pixels/canais.

• Generalização do Formalismo YOLO:

  • O método de regressão foi adaptado para prever caixas delimitadoras 3D (bounding boxes), incluindo coordenadas centrais (x, y, z), dimensões (largura, altura, profundidade) e um escore de "objetividade" (confiança).
  • Além da localização, a rede prevê parâmetros físicos das galáxias: fluxo de linha integrado, tamanho do HI, largura da linha (w20), ângulo de posição (PA) e inclinação (i).
  • A métrica de similaridade para associação de detecções foi atualizada para DIoU 3D (Distance-aware Intersection over Union), que considera a distância entre os centros e o volume das caixas.

• Pré-processamento e Treinamento:

  • Seleção de Dados de Treino: O catálogo de verdade do cubo LDEV (menor) foi filtrado usando uma função de seleção baseada em fluxo e brilho volumétrico para remover fontes indetectáveis, evitando ruído no treinamento.
  • Estratégia de Bootstrapping: Foi desenvolvido um processo iterativo onde um modelo inicial (BASE) é treinado, usado para detectar fontes no mesmo conjunto de treino, e essas detecções (incluindo fontes mais fracas que o critério original) são usadas para refinar o catálogo de treino para o próximo modelo (BT1, BT2, BT3). Isso permite que o modelo aprenda a detectar fontes que estavam inicialmente abaixo do limiar de seleção.
  • Aumento de Dados (Augmentation): Inclui flutuações de máscara nas bordas, rotações e flips (inclusive na dimensão de frequência, assumindo simetria morfológica), para aumentar a diversidade dos exemplos.

• Pipeline de Inferência:

  • O cubo MAIN é dividido em subcubos com sobreposição para garantir que fontes não sejam cortadas.
  • As detecções são filtradas usando Supressão Não-Máxima (NMS) 3D.
  • Um limiar de confiança (objetividade) é otimizado pós-processamento para maximizar a pontuação no desafio, considerando a densidade de fontes em diferentes faixas de frequência.

3. Principais Contribuições

• Generalização 3D: Adaptação bem-sucedida do YOLO para dados 3D hiperspectrais, demonstrando que redes CNN 3D podem capturar padrões complexos de rotação e estrutura de galáxias em cubos de dados de rádio.
• Vitória no SDC2: A equipe MINERVA, utilizando este método, venceu o desafio oficial. O artigo apresenta uma versão atualizada que supera o resultado original do desafio.
• Estratégia de Bootstrapping: Introdução de uma técnica de refinamento de dados de treino baseada na auto-avaliação do modelo, permitindo a inclusão de fontes mais fracas e difíceis que não eram inicialmente consideradas detectáveis.
• Eficiência Computacional: O método é extremamente rápido, processando o cubo de dados completo de ~1 TB em apenas 30 minutos em uma única GPU (RTX 6000 Ada).

4. Resultados

• Pontuação: A versão atualizada (modelo BT1) alcançou uma pontuação de 25.453 no SDC2, uma melhoria de 9,5% em relação à pontuação original da equipe no desafio e 13,2% acima do segundo colocado.
• Precisão e Pureza: O catálogo resultante exibe uma pureza de detecção de 92,3% e uma precisão de caracterização "best-in-class" (melhor da categoria).
• Sensibilidade: O método detectou 45% mais fontes confirmadas do que as ferramentas clássicas de detecção (como o SoFiA) usadas por outras equipes.
• Desempenho por Fluxo: O modelo consegue detectar fontes mais fracas do que o limite imposto pela função de seleção inicial de treino, graças à estratégia de bootstrapping.
• Comparação: Mesmo com limiares de pureza estritos (99%), o método supera significativamente a melhor abordagem não baseada em ML (Equipe SOFIA).

5. Significado e Conclusão

Este trabalho demonstra que os detectores baseados em CNN 3D são ferramentas state-of-the-art para o processamento de grandes cubos de dados hiperspectrais astronômicos.

• Preparação para o SKA: O sucesso no SDC2 valida a abordagem para a futura exploração científica do SKA, onde o volume de dados será imenso.
• Aplicabilidade Real: Os autores indicam que o método está sendo adaptado para dados observacionais de precursores como as pesquisas LADUMA (MeerKAT) e WALLABY (ASKAP), mostrando boa generalização após simples reescalonamento.
• Limitações e Futuro: O artigo reconhece que a dependência de dados simulados para treinamento é uma limitação, pois simulações podem não capturar todas as nuances dos dados reais. Sugere-se o uso de transfer learning e treinamento híbrido com dados reais rotulados para futuras aplicações.
• Impacto Ambiental: O custo computacional adicional foi estimado em 4.000 horas de GPU, resultando em uma pegada de carbono de aproximadamente 104 kg de CO2-e, um custo considerado baixo para o ganho científico obtido.

Em resumo, o YOLO-CIANNA 3D representa um avanço significativo na astronomia de rádio, oferecendo um equilíbrio superior entre velocidade, sensibilidade e precisão na extração de catálogos de galáxias de grandes volumes de dados.