# AQuaMaM: An Autoregressive, Quaternion Manifold Model for Rapidly Estimating Complex SO(3) Distributions

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

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 29 of its 55 references that are in the map.
- First posted on arXiv 21 Jan 2023.

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

- [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)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Microsoft COCO: Common Objects in Context](https://mlsanity.com/paper/1405.0312/) (2014)
- [Neural Machine Translation by Jointly Learning to Align and Translate](https://mlsanity.com/paper/1409.0473/) (2014)
- [Zero-Shot Text-to-Image Generation](https://mlsanity.com/paper/2102.12092/) (2021)
- [Gaussian Error Linear Units (GELUs)](https://mlsanity.com/paper/1606.08415/) (2016)
- [Fourier Features Let Networks Learn High Frequency Functions in Low Dimensional Domains](https://mlsanity.com/paper/2006.10739/) (2020)
- [A Generalist Agent](https://mlsanity.com/paper/2205.06175/) (2022)
- [Image Transformer](https://mlsanity.com/paper/1802.05751/) (2018)
- [Neural Turing Machines](https://mlsanity.com/paper/1410.5401/) (2014)
- [Generating Sequences With Recurrent Neural Networks](https://mlsanity.com/paper/1308.0850/) (2013)
- [Memory Networks](https://mlsanity.com/paper/1410.3916/) (2014)
- [T-LESS: An RGB-D Dataset for 6D Pose Estimation of Texture-less Objects](https://mlsanity.com/paper/1701.05498/) (2017)
- [Strike (with) a Pose: Neural Networks Are Easily Fooled by Strange Poses of Familiar Objects](https://mlsanity.com/paper/1811.11553/) (2018)
- [Learning in an Uncertain World: Representing Ambiguity Through Multiple Hypotheses](https://mlsanity.com/paper/1612.00197/) (2016)
- [Overcoming Limitations of Mixture Density Networks: A Sampling and Fitting Framework for Multimodal Future Prediction](https://mlsanity.com/paper/1906.03631/) (2019)
- [Tabular Transformers for Modeling Multivariate Time Series](https://mlsanity.com/paper/2011.01843/) (2020)
- [RNADE: The real-valued neural autoregressive density-estimator](https://mlsanity.com/paper/1306.0186/) (2013)
- [Implicit-PDF: Non-Parametric Representation of Probability Distributions on the Rotation Manifold](https://mlsanity.com/paper/2106.05965/) (2021)
- [Deep Directional Statistics: Pose Estimation with Uncertainty Quantification](https://mlsanity.com/paper/1805.03430/) (2018)
- [Deep Bingham Networks: Dealing with Uncertainty and Ambiguity in Pose Estimation](https://mlsanity.com/paper/2012.11002/) (2020)
- [A Mixed Classification-Regression Framework for 3D Pose Estimation from 2D Images](https://mlsanity.com/paper/1805.03225/) (2018)
- [TraDE: Transformers for Density Estimation](https://mlsanity.com/paper/2004.02441/) (2020)
- [baller2vec: A Multi-Entity Transformer For Multi-Agent Spatiotemporal Modeling](https://mlsanity.com/paper/2102.03291/) (2021)
- [baller2vec++: A Look-Ahead Multi-Entity Transformer For Modeling Coordinated Agents](https://mlsanity.com/paper/2104.11980/) (2021)
- [The DEformer: An Order-Agnostic Distribution Estimating Transformer](https://mlsanity.com/paper/2106.06989/) (2021)

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

- [1 bit is all we need: binary normalized neural networks](https://mlsanity.com/paper/2509.07025/) (2025)

## Abstract (arXiv)

Accurately modeling complex, multimodal distributions for rotations in three-dimensions, i.e., the SO(3) group, is challenging due to the curvature of the rotation manifold. The recently described implicit-PDF (IPDF) is a simple, elegant, and effective approach for learning arbitrary distributions on SO(3) up to a given precision. However, inference with IPDF requires N forward passes through the network's final multilayer perceptron (where N places an upper bound on the likelihood that can be calculated by the model), which is prohibitively slow for those without the computational resources necessary to parallelize the queries. In this paper, I introduce AQuaMaM, a neural network capable of both learning complex distributions on the rotation manifold and calculating exact likelihoods for query rotations in a single forward pass. Specifically, AQuaMaM autoregressively models the projected components of unit quaternions as mixtures of uniform distributions that partition their geometrically-restricted domain of values. When trained on an "infinite" toy dataset with ambiguous viewpoints, AQuaMaM rapidly converges to a sampling distribution closely matching the true data distribution. In contrast, the sampling distribution for IPDF dramatically diverges from the true data distribution, despite IPDF approaching its theoretical minimum evaluation loss during training. When trained on a constructed dataset of 500,000 renders of a die in different rotations, AQuaMaM reaches a test log-likelihood 14% higher than IPDF. Further, compared to IPDF, AQuaMaM uses 24% fewer parameters, has a prediction throughput 52× faster on a single GPU, and converges in a similar amount of time during training.