Using multi-task learning to determine gas-phase metallicity of star-forming galaxies

Aims . This study aimed to improve the estimation of the gas-phase metallicity of star-forming galaxies by using a multi-task learning approach that simultaneously performs gas-phase metallicity estimation and spectral classification of galaxies. Methods . We propose a multi-task learning model to perform simultaneous gas-phase metallicity estimation and spectral classification of galaxies (MTLforGalSpecZ). The architecture is composed of three main components: (1) a shared feature extraction module, (2) a channel attention mechanism, and (3) two task-specific output heads. Specifically, the shared feature extraction module consists of stacked convolutional blocks that process spectroscopic inputs to extract global spectral features. These features are then passed to a channel attention mechanism, which adjusts the importance of each spectral channel. Finally, these features are fed into two separate output heads: a regression head to estimate the gas-phase metallicity and a classification head to determine the spectral class. The model is optimised using a combined loss function that includes both classification and regression losses. A conditional masking strategy is applied to the regression loss to ensure that metallicity estimation is performed only for star-forming galaxies. Results . The model was trained on a dataset of approximately 100000 spectra, each labelled with a galaxy class, with gas-phase metallicity labels available for star-forming galaxies. On the test set, it achieves a prediction scatter of σ = 0.0374 for metallicity and a classification accuracy of 97.01%. Compared to running two independent single-task networks, MTLforGalSpecZ improves metallicity prediction performance by 30%, while also reducing total training time by 18.3% and inference time by 45.2%.