# Equilibrium Aggregation: Encoding Sets via Optimization

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

## Key facts

- Cited by 5 papers in the map (data through 14 Oct 2026).
- Builds on 44 of its 71 references that are in the map.
- First posted on arXiv 25 Feb 2022.

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

- [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)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Semi-Supervised Classification with Graph Convolutional Networks](https://mlsanity.com/paper/1609.02907/) (2016)
- [Graph Attention Networks](https://mlsanity.com/paper/1710.10903/) (2017)
- [Model-Agnostic Meta-Learning for Fast Adaptation of Deep Networks](https://mlsanity.com/paper/1703.03400/) (2017)
- [How Powerful are Graph Neural Networks?](https://mlsanity.com/paper/1810.00826/) (2018)
- [Generative Modeling by Estimating Gradients of the Data Distribution](https://mlsanity.com/paper/1907.05600/) (2019)
- [Neural Ordinary Differential Equations](https://mlsanity.com/paper/1806.07366/) (2018)
- [Open Graph Benchmark: Datasets for Machine Learning on Graphs](https://mlsanity.com/paper/2005.00687/) (2020)
- [Relational inductive biases, deep learning, and graph networks](https://mlsanity.com/paper/1806.01261/) (2018)
- [Convolutional Networks on Graphs for Learning Molecular Fingerprints](https://mlsanity.com/paper/1509.09292/) (2015)
- [Compressive Transformers for Long-Range Sequence Modelling](https://mlsanity.com/paper/1911.05507/) (2019)
- [Learning to learn by gradient descent by gradient descent](https://mlsanity.com/paper/1606.04474/) (2016)
- [Deep Equilibrium Models](https://mlsanity.com/paper/1909.01377/) (2019)
- [OptNet: Differentiable Optimization as a Layer in Neural Networks](https://mlsanity.com/paper/1703.00443/) (2017)
- [Measuring and Relieving the Over-smoothing Problem for Graph Neural Networks from the Topological View](https://mlsanity.com/paper/1909.03211/) (2019)
- [Meta-Learning with Implicit Gradients](https://mlsanity.com/paper/1909.04630/) (2019)
- [Principal Neighbourhood Aggregation for Graph Nets](https://mlsanity.com/paper/2004.05718/) (2020)
- [Input Convex Neural Networks](https://mlsanity.com/paper/1609.07152/) (2016)
- [3D Steerable CNNs: Learning Rotationally Equivariant Features in Volumetric Data](https://mlsanity.com/paper/1807.02547/) (2018)
- [Harmonic Networks: Deep Translation and Rotation Equivariance](https://mlsanity.com/paper/1612.04642/) (2016)
- [A Note on Over-Smoothing for Graph Neural Networks](https://mlsanity.com/paper/2006.13318/) (2020)
- [Order Matters: Sequence to sequence for sets](https://mlsanity.com/paper/1511.06391/) (2015)
- [Efficient and Modular Implicit Differentiation](https://mlsanity.com/paper/2105.15183/) (2021)
- [Differentiable MPC for End-to-end Planning and Control](https://mlsanity.com/paper/1810.13400/) (2018)
- [Towards a Neural Statistician](https://mlsanity.com/paper/1606.02185/) (2016)
- [Implicit Graph Neural Networks](https://mlsanity.com/paper/2009.06211/) (2020)
- [Janossy Pooling: Learning Deep Permutation-Invariant Functions for Variable-Size Inputs](https://mlsanity.com/paper/1811.01900/) (2018)
- [On the Limitations of Representing Functions on Sets](https://mlsanity.com/paper/1901.09006/) (2019)

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

- [PyG 2.0: Scalable Learning on Real World Graphs](https://mlsanity.com/paper/2507.16991/) (2025)
- [Learning multi-modal generative models with permutation-invariant encoders and tighter variational objectives](https://mlsanity.com/paper/2309.00380/) (2023)
- [Advances in Set Function Learning: A Survey of Techniques and Applications](https://mlsanity.com/paper/2501.14991/) (2025)
- [Exploiting Inferential Structure in Neural Processes](https://mlsanity.com/paper/2306.15169/) (2023)
- [Convolutional Set Transformer](https://mlsanity.com/paper/2509.22889/) (2025)

## Abstract (arXiv)

Processing sets or other unordered, potentially variable-sized inputs in neural networks is usually handled by aggregating a number of input tensors into a single representation. While a number of aggregation methods already exist from simple sum pooling to multi-head attention, they are limited in their representational power both from theoretical and empirical perspectives. On the search of a principally more powerful aggregation strategy, we propose an optimization-based method called Equilibrium Aggregation. We show that many existing aggregation methods can be recovered as special cases of Equilibrium Aggregation and that it is provably more efficient in some important cases. Equilibrium Aggregation can be used as a drop-in replacement in many existing architectures and applications. We validate its efficiency on three different tasks: median estimation, class counting, and molecular property prediction. In all experiments, Equilibrium Aggregation achieves higher performance than the other aggregation techniques we test.