Attention-based Preprocessing Framework for Improving Rare Transient Classification

Abstract With large numbers of transients discovered by current and future imaging surveys, machine learning is increasingly applied to light-curve and host galaxy properties to select events for follow-up. However, finding rare types of transients remains difficult due to extreme class imbalances in training sets, and extracting features from host images is complicated by the presence of bright foreground sources, particularly if the true host is faint or distant. Here, we present a data augmentation pipeline for images and light curves that mitigates these issues, and apply this to improve the classification of Type I superluminous supernovae (SLSNe-I) and tidal disruption events (TDEs) with our existing NEural Engine for Discovering Luminous Events (NEEDLE) code. The method uses a similarity index to remove image artifacts, and a masking procedure that removes unrelated sources while preserving the transient and its host. This focuses classifier attention on the relevant pixels and enables arbitrary rotations for class upsampling. We also fit observed multiband light curves with a two-dimensional Gaussian process and generate data-driven synthetic samples by resampling and redshifting these models, crossmatching with galaxy images in the same class to produce unique but realistic new examples for training. Models trained with the augmented dataset achieve substantially higher purity: for classifications with a confidence of 0.8 or higher, we achieve 75% (43%) purity and 75% (66%) completeness for SLSNe-I (TDEs). The updated NEEDLE 2.0 is publicly available at https://github.com/XinyueSheng2019/needle .