# Relational Reasoning via Set Transformers: Provable Efficiency and Applications to MARL

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

## Key facts

- Cited by 12 papers in the map (data through 14 Oct 2026).
- Builds on 37 of its 64 references that are in the map.
- First posted on arXiv 20 Sep 2022.

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

- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Neural Machine Translation by Jointly Learning to Align and Translate](https://mlsanity.com/paper/1409.0473/) (2014)
- [Playing Atari with Deep Reinforcement Learning](https://mlsanity.com/paper/1312.5602/) (2013)
- [Multi-Agent Actor-Critic for Mixed Cooperative-Competitive Environments](https://mlsanity.com/paper/1706.02275/) (2017)
- [Deep Sets](https://mlsanity.com/paper/1703.06114/) (2017)
- [Geometric Deep Learning: Grids, Groups, Graphs, Geodesics, and Gauges](https://mlsanity.com/paper/2104.13478/) (2021)
- [Tokens-to-Token ViT: Training Vision Transformers from Scratch on ImageNet](https://mlsanity.com/paper/2101.11986/) (2021)
- [Spectrally-normalized margin bounds for neural networks](https://mlsanity.com/paper/1706.08498/) (2017)
- [MINE: Mutual Information Neural Estimation](https://mlsanity.com/paper/1801.04062/) (2018)
- [Intriguing Properties of Vision Transformers](https://mlsanity.com/paper/2105.10497/) (2021)
- [Is Pessimism Provably Efficient for Offline RL?](https://mlsanity.com/paper/2012.15085/) (2020)
- [A PAC-Bayesian Approach to Spectrally-Normalized Margin Bounds for Neural Networks](https://mlsanity.com/paper/1707.09564/) (2017)
- [Information-Theoretic Considerations in Batch Reinforcement Learning](https://mlsanity.com/paper/1905.00360/) (2019)
- [Bellman-consistent Pessimism for Offline Reinforcement Learning](https://mlsanity.com/paper/2106.06926/) (2021)
- [Emergence of Grounded Compositional Language in Multi-Agent Populations](https://mlsanity.com/paper/1703.04908/) (2017)
- [FLAMBE: Structural Complexity and Representation Learning of Low Rank MDPs](https://mlsanity.com/paper/2006.10814/) (2020)
- [AlgaeDICE: Policy Gradient from Arbitrary Experience](https://mlsanity.com/paper/1912.02074/) (2019)
- [Pessimistic Model-based Offline Reinforcement Learning under Partial Coverage](https://mlsanity.com/paper/2107.06226/) (2021)
- [Minimax Weight and Q-Function Learning for Off-Policy Evaluation](https://mlsanity.com/paper/1910.12809/) (2019)
- [Inductive Biases and Variable Creation in Self-Attention Mechanisms](https://mlsanity.com/paper/2110.10090/) (2021)
- [Toward the Fundamental Limits of Imitation Learning](https://mlsanity.com/paper/2009.05990/) (2020)
- [Representation Benefits of Deep Feedforward Networks](https://mlsanity.com/paper/1509.08101/) (2015)
- [SO(2)-Equivariant Reinforcement Learning](https://mlsanity.com/paper/2203.04439/) (2022)
- [Near-optimal Offline Reinforcement Learning with Linear Representation: Leveraging Variance Information with Pessimism](https://mlsanity.com/paper/2203.05804/) (2022)
- [Neural Network with Unbounded Activation Functions is Universal Approximator](https://mlsanity.com/paper/1505.03654/) (2015)
- [Equivariant Q Learning in Spatial Action Spaces](https://mlsanity.com/paper/2110.15443/) (2021)
- [The Sensory Neuron as a Transformer: Permutation-Invariant Neural Networks for Reinforcement Learning](https://mlsanity.com/paper/2109.02869/) (2021)
- [Generalization Error in Deep Learning](https://mlsanity.com/paper/1808.01174/) (2018)
- [PIC: Permutation Invariant Critic for Multi-Agent Deep Reinforcement Learning](https://mlsanity.com/paper/1911.00025/) (2019)
- [Generalization Error of Invariant Classifiers](https://mlsanity.com/paper/1610.04574/) (2016)

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

- [What and How does In-Context Learning Learn? Bayesian Model Averaging, Parameterization, and Generalization](https://mlsanity.com/paper/2305.19420/) (2023)
- [Chain of Preference Optimization: Improving Chain-of-Thought Reasoning in LLMs](https://mlsanity.com/paper/2406.09136/) (2024)
- [LongSpec: Long-Context Lossless Speculative Decoding with Efficient Drafting and Verification](https://mlsanity.com/paper/2502.17421/) (2025)
- [An Analysis of Attention via the Lens of Exchangeability and Latent Variable Models](https://mlsanity.com/paper/2212.14852/) (2022)
- [Transformers in Reinforcement Learning: A Survey](https://mlsanity.com/paper/2307.05979/) (2023)
- [LightTransfer: Your Long-Context LLM is Secretly a Hybrid Model with Effortless Adaptation](https://mlsanity.com/paper/2410.13846/) (2024)
- [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)
- [Multi-level Advantage Credit Assignment for Cooperative Multi-Agent Reinforcement Learning](https://mlsanity.com/paper/2508.06836/) (2025)
- [Towards Better Understanding of In-Context Learning Ability from In-Context Uncertainty Quantification](https://mlsanity.com/paper/2405.15115/) (2024)
- [Puzzle it Out: Local-to-Global World Model for Offline Multi-Agent Reinforcement Learning](https://mlsanity.com/paper/2601.07463/) (2026)
- [Topology Enhanced MARL for Multi-Agent Cooperative Decision-Making of CAVs](https://mlsanity.com/paper/2507.12110/) (2025)

## Abstract (arXiv)

The cooperative Multi-A gent R einforcement Learning (MARL) with permutation invariant agents framework has achieved tremendous empirical successes in real-world applications. Unfortunately, the theoretical understanding of this MARL problem is lacking due to the curse of many agents and the limited exploration of the relational reasoning in existing works. In this paper, we verify that the transformer implements complex relational reasoning, and we propose and analyze model-free and model-based offline MARL algorithms with the transformer approximators. We prove that the suboptimality gaps of the model-free and model-based algorithms are independent of and logarithmic in the number of agents respectively, which mitigates the curse of many agents. These results are consequences of a novel generalization error bound of the transformer and a novel analysis of the Maximum Likelihood Estimate (MLE) of the system dynamics with the transformer. Our model-based algorithm is the first provably efficient MARL algorithm that explicitly exploits the permutation invariance of the agents. Our improved generalization bound may be of independent interest and is applicable to other regression problems related to the transformer beyond MARL.