Compact neural networks for astronomy with optimal transport bias correction

TL;DR

WaveletMamba is a novel, efficient framework combining wavelet analysis and optimal transport bias correction, significantly improving astronomical image classification accuracy and robustness across resolutions with fewer parameters.

Contribution

The paper introduces WaveletMamba, a theory-driven approach that integrates wavelet decomposition, state-space modeling, and multi-level bias correction to enhance astronomical image analysis.

Findings

Achieves 81.72% classification accuracy at 64x64 resolution with 3.54M parameters.

Delivers high-resolution performance with 80.93% accuracy at 244x244 inputs, 9.7x more efficient.

Improves Log-MSE by 22.96% and reduces outliers by 26.10% through bias correction techniques.

Abstract

Astronomical imaging confronts an efficiency-resolution tradeoff that limits large-scale morphological classification and redshift prediction. We introduce WaveletMamba, a theory-driven framework integrating wavelet decomposition with state-space modeling, mathematical regularization, and multi-level bias correction. WaveletMamba achieves 81.72% +/- 0.53% classification accuracy at 64x64 resolution with only 3.54M parameters, delivering high-resolution performance (80.93% +/- 0.27% at 244x244) at low-resolution inputs with 9.7x computational efficiency gains. The framework exhibits Resolution Multistability, where models trained on low-resolution data achieve consistent accuracy across different input scales despite divergent internal representations. The framework's multi-level bias correction synergizes HK distance (distribution-level optimal transport) with Color-Aware Weighting (sample-level fine-tuning), achieving 22.96% Log-MSE improvement and 26.10% outlier reduction without explicit selection function modeling. Here, we show that mathematical rigor enables unprecedented efficiency and comprehensive bias correction in scientific AI, bridging computer vision and astrophysics to revolutionize interdisciplinary scientific discovery.