Flow Matching for Generating Weakly Labeled Bags of Foundation-Model Mammography Representations

Abstract

Annotated medical imaging data remain scarce, and labels are often weak and noisy: in mammography, an examination comprises several high-resolution views, yet the diagnostic outcome is recorded only at the breast level. Such problems are naturally cast as Multiple Instance Learning (MIL), where the model must infer instance-level structure from bag-level labels alone. Although contemporary foundation encoders supply strong general-purpose embeddings, augmenting MIL data in this representation space remains an open problem as established techniques act on one instance at a time and ignore the statistical dependencies binding a bag together. We address this with SetFlow, a generative model that learns the distribution of complete MIL bags directly in a frozen encoder’s embedding space. SetFlow couples flow-matching training with a Set Transformer-inspired backbone, making it invariant to instance ordering while modeling intra-bag relationships. Generation is conditioned jointly on class label and per-instance scale, yielding coherent, semantically faithful bags rather than isolated vectors. Evaluating on two large public mammography datasets and two encoders, we assess distributional fidelity, nearest-neighbor behavior, and downstream augmentation utility. We show that generated bags reproduce real-data statistics and improve classification in certain configuration, with performance gains varying on the amount of synthetic data. An architecture ablation confirms each design choice contributes to performance.