MAPA: Mode-Aware Prototype Allocation for Latent Dataset Distillation in 3D Medical Classification

원문 초록

Dataset distillation (DD) synthesizes a small training set from a large one, but applying it to 3D medical classification is constrained by the memory and optimization costs of synthesizing full volumetric samples. Under a storage-equivalent budget, distilling in the latent space of a frozen autoencoder reduces this cost and fits more synthetic samples (prototypes) than in voxel space. However, class-level matching, such as distribution matching, causes prototypes to concentrate in dense latent regions while underrepresenting sparser intra-class regions, a consistent imbalance that we term prototype concentration. To address this, we propose MAPA (Mode-Aware Prototype Allocation), a latent distillation method that spreads prototypes across intra-class modes while preserving class-level matching. MAPA marks these modes with fixed anchors estimated by clustering the real latents and distributes prototypes over them through a soft assignment, together with a lightweight anti-collapse regularizer and curriculum weighting that stabilize the optimization. Across five 3D medical classification tasks, MAPA improves the aggregate balanced accuracy by +4.2 points over the strongest latent baseline while using the same number of prototypes. An exploratory severity analysis on MosMedData suggests improved worst-stratum recall without loss of normal-case specificity.