# Distributionally Robust Set Representation Learning Under Inference-Time Element Corruption

arXiv: https://arxiv.org/abs/2605.30089 · HTML: https://mlsanity.com/paper/2605.30089/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 63 of its 93 references that are in the map.
- First posted on arXiv 28 May 2026.

## Builds on (most cited first, 30 of 63)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [PointNet: Deep Learning on Point Sets for 3D Classification and Segmentation](https://mlsanity.com/paper/1612.00593/) (2016)
- [Efficient Estimation of Word Representations in Vector Space](https://mlsanity.com/paper/1301.3781/) (2013)
- [Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distances](https://mlsanity.com/paper/1306.0895/) (2013)
- [Distributionally Robust Neural Networks for Group Shifts: On the Importance of Regularization for Worst-Case Generalization](https://mlsanity.com/paper/1911.08731/) (2019)
- [Perceiver: General Perception with Iterative Attention](https://mlsanity.com/paper/2103.03206/) (2021)
- [Remote Sensing Image Scene Classification: Benchmark and State of the Art](https://mlsanity.com/paper/1703.00121/) (2017)
- [Spectrally-normalized margin bounds for neural networks](https://mlsanity.com/paper/1706.08498/) (2017)
- [Network In Network](https://mlsanity.com/paper/1312.4400/) (2013)
- [Generalized Sliced Wasserstein Distances](https://mlsanity.com/paper/1902.00434/) (2019)
- [Understanding The Robustness in Vision Transformers](https://mlsanity.com/paper/2204.12451/) (2022)
- [CoEvoSkills: Self-Evolving Agent Skills via Co-Evolutionary Verification](https://mlsanity.com/paper/2604.01687/) (2026)
- [Towards a Neural Statistician](https://mlsanity.com/paper/1606.02185/) (2016)
- [Janossy Pooling: Learning Deep Permutation-Invariant Functions for Variable-Size Inputs](https://mlsanity.com/paper/1811.01900/) (2018)
- [AgentAuditor: Human-Level Safety and Security Evaluation for LLM Agents](https://mlsanity.com/paper/2506.00641/) (2025)
- [Distributionally Robust Optimization: A Review](https://mlsanity.com/paper/1908.05659/) (2019)
- [Distributionally Robust Optimization and Generalization in Kernel Methods](https://mlsanity.com/paper/1905.10943/) (2019)
- [FSPool: Learning Set Representations with Featurewise Sort Pooling](https://mlsanity.com/paper/1906.02795/) (2019)
- [Towards Robust Alignment of Language Models: Distributionally Robustifying Direct Preference Optimization](https://mlsanity.com/paper/2407.07880/) (2024)
- [Modeling Scale-free Graphs with Hyperbolic Geometry for Knowledge-aware Recommendation](https://mlsanity.com/paper/2108.06468/) (2021)
- [RecRanker: Instruction Tuning Large Language Model as Ranker for Top-k Recommendation](https://mlsanity.com/paper/2312.16018/) (2023)
- [Dynamically Expandable Graph Convolution for Streaming Recommendation](https://mlsanity.com/paper/2303.11700/) (2023)
- [Integrating Large Language Models into Recommendation via Mutual Augmentation and Adaptive Aggregation](https://mlsanity.com/paper/2401.13870/) (2024)
- [Rep the Set: Neural Networks for Learning Set Representations](https://mlsanity.com/paper/1904.01962/) (2019)
- [Learning Prototype-oriented Set Representations for Meta-Learning](https://mlsanity.com/paper/2110.09140/) (2021)
- [GAM: Hierarchical Graph-based Agentic Memory for LLM Agents](https://mlsanity.com/paper/2604.12285/) (2026)
- [DyG-Mamba: Continuous State Space Modeling on Dynamic Graphs](https://mlsanity.com/paper/2408.06966/) (2024)
- [Distributionally Robust Bayesian Optimization with varphi-divergences](https://mlsanity.com/paper/2203.02128/) (2022)
- [Learning Representations of Sets through Optimized Permutations](https://mlsanity.com/paper/1812.03928/) (2018)
- [Principled learning method for Wasserstein distributionally robust optimization with local perturbations](https://mlsanity.com/paper/2006.03333/) (2020)

## Abstract (arXiv)

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 set representation and degrade performance. We propose SW-DRSO, a distributionally robust optimization framework tailored for sets. Rather than minimizing loss solely on observed training data, SW-DRSO optimizes a tractable surrogate of the worst-case expected loss over a family of plausible inference-time variations. We introduce a barycentric adversary that approximates the intractable search over corrupted sets by a differentiable training-time optimization over simplex weights. Extensive experiments across four tasks demonstrate that SW-DRSO effectively enhances robustness against corruption while maintaining high overall performance.