Distributionally Robust Set Representation Learning Under Inference-Time Element Corruption
Abstract
Standard Set Representation Learning methods typically excel on curated data but often overlook the challenge of Inference-time Element Corruption. This refers to scenarios where deployed models encounter element-level degradations, such as outliers or missing components, that may distort the set representation and degrade performance. To address this, we propose SW-DRSO, a distributionally robust optimization framework tailored for sets. Rather than minimizing loss solely on the observed training data, SW-DRSO optimizes the worst-case expected loss over a family of plausible inference-time variations. We further introduce a barycentric adversary that transforms the intractable search for worst-case corrupted sets into a differentiable and efficient optimization process. Extensive experiments across four tasks demonstrate that SW-DRSO effectively enhances robustness against corruption while maintaining high overall performance.