# Scalable Set Encoding with Universal Mini-Batch Consistency and Unbiased Full Set Gradient Approximation

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

## Key facts

- Cited by 4 papers in the map (data through 14 Oct 2026).
- Builds on 24 of its 46 references that are in the map.
- First posted on arXiv 26 Aug 2022.

## Builds on (most cited first, 24 of 24)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [A Simple Framework for Contrastive Learning of Visual Representations](https://mlsanity.com/paper/2002.05709/) (2020)
- [Layer Normalization](https://mlsanity.com/paper/1607.06450/) (2016)
- [Auto-Encoding Variational Bayes](https://mlsanity.com/paper/1312.6114/) (2013)
- [Deep Learning Face Attributes in the Wild](https://mlsanity.com/paper/1411.7766/) (2014)
- [Longformer: The Long-Document Transformer](https://mlsanity.com/paper/2004.05150/) (2020)
- [Attention-based Deep Multiple Instance Learning](https://mlsanity.com/paper/1802.04712/) (2018)
- [Object-Centric Learning with Slot Attention](https://mlsanity.com/paper/2006.15055/) (2020)
- [Dual-stream Multiple Instance Learning Network for Whole Slide Image Classification with Self-supervised Contrastive Learning](https://mlsanity.com/paper/2011.08939/) (2020)
- [Conditional Neural Processes](https://mlsanity.com/paper/1807.01613/) (2018)
- [Large-Scale Multi-Label Text Classification on EU Legislation](https://mlsanity.com/paper/1906.02192/) (2019)
- [Hierarchical Transformers for Long Document Classification](https://mlsanity.com/paper/1910.10781/) (2019)
- [Janossy Pooling: Learning Deep Permutation-Invariant Functions for Variable-Size Inputs](https://mlsanity.com/paper/1811.01900/) (2018)
- [On the Almost Sure Convergence of Stochastic Gradient Descent in Non-Convex Problems](https://mlsanity.com/paper/2006.11144/) (2020)
- [Convergence rates for the stochastic gradient descent method for non-convex objective functions](https://mlsanity.com/paper/1904.01517/) (2019)
- [Universal Approximation of Functions on Sets](https://mlsanity.com/paper/2107.01959/) (2021)
- [FSPool: Learning Set Representations with Featurewise Sort Pooling](https://mlsanity.com/paper/1906.02795/) (2019)
- [The power of deeper networks for expressing natural functions](https://mlsanity.com/paper/1705.05502/) (2017)
- [Efficient Classification of Long Documents Using Transformers](https://mlsanity.com/paper/2203.11258/) (2022)
- [Understanding Dynamics of Nonlinear Representation Learning and Its Application](https://mlsanity.com/paper/2106.14836/) (2021)
- [Mini-Batch Consistent Slot Set Encoder for Scalable Set Encoding](https://mlsanity.com/paper/2103.01615/) (2021)

## Built on by (most cited first, 4 of 4)

- [Set-based Neural Network Encoding Without Weight Tying](https://mlsanity.com/paper/2305.16625/) (2023)
- [Advances in Set Function Learning: A Survey of Techniques and Applications](https://mlsanity.com/paper/2501.14991/) (2025)
- [Mergeable Model-Side Aggregation States for Long-Context Language Models](https://mlsanity.com/paper/2607.26448/) (2026)
- [Monotone and Separable Set Functions: Characterizations and Neural Models](https://mlsanity.com/paper/2510.23634/) (2025)

## Abstract (arXiv)

Recent work on mini-batch consistency (MBC) for set functions has brought attention to the need for sequentially processing and aggregating chunks of a partitioned set while guaranteeing the same output for all partitions. However, existing constraints on MBC architectures lead to models with limited expressive power. Additionally, prior work has not addressed how to deal with large sets during training when the full set gradient is required. To address these issues, we propose a Universally MBC (UMBC) class of set functions which can be used in conjunction with arbitrary non-MBC components while still satisfying MBC, enabling a wider range of function classes to be used in MBC settings. Furthermore, we propose an efficient MBC training algorithm which gives an unbiased approximation of the full set gradient and has a constant memory overhead for any set size for both train- and test-time. We conduct extensive experiments including image completion, text classification, unsupervised clustering, and cancer detection on high-resolution images to verify the efficiency and efficacy of our scalable set encoding framework. Our code is available at github.com/jeffwillette/umbc