AQuaMaM: An Autoregressive, Quaternion Manifold Model for Rapidly Estimating Complex SO(3) Distributions
Authors: Michael A. Alcorn
Venue: Submitted to ICLR 2023 (2023)

Abstract: Accurately modeling complex, multimodal distributions is necessary for optimal decision-making, but doing so 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 a mixture of uniform distributions that partition their geometrically-restricted domain of values. 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. On a constructed dataset of 500,000 renders of a die in different rotations, an AQuaMaM model trained from scratch reaches a log-likelihood 14% higher than an IPDF model using a pretrained ResNet-50. Further, compared to IPDF, AQuaMaM uses 24% fewer parameters, has a prediction throughput 52$\times$ faster on a single GPU, and converges in a similar amount of time during training.

Under review as a conference paper at ICLR 2023
AQUAMAM: AN AUTOREGRESSIVE, QUATERNION MANIFOLD MODEL FOR RAPIDLY ESTIMATING COM- PLEX SO (3) DISTRIBUTIONS
Anonymous authors Paper under double-blind review
ABSTRACT
Accurately modeling complex, multimodal distributions is necessary for optimal decision-making. but doing so 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 ference with IPDF requires N forward passes through the network's final multi- layer serceptron-where N places upper bound on the likelihood that be calculated by the model-which is prohibitively slow for those without the com- putational resources necessary to parallelize the queries. In this paper, introduce AQuaMaM a neural network capable of both learning complex distributions on the rotation manifold and calculating exact likelihoods for query rotations in sin- gle forward pass. Specifically, AQuaMaM autoregresssively models the projected components of unit quaternions as mixtures of uniform distributions that partition their geometrically-rrstriiee domain of values. On an "infinite" toy dataset with ambiguous viewpoints, AQuaMaM rapidly converges to sampling distribution closely matching the true data distribution. In contrast, the sampling distribu- tion for IPDF dramatically diverges from the true data distribution, despite IPDF approaching its theoretical minimum evaluation loss during training. On con- structed dataset of 500,000 renders of a die in different rotations, an AQuaMaM model trained from scratch reaches model using a pretrained ResNet-50. Further, compared to IPDF, AQuaMaM uses 24% fewer parameters, has a prediction throughput 52x faster on single GPU, and converges in a similar amount of time during training.
SO(3) up to given precision. However, in-
logeliiel 14% higher than an IPDF
INTRODUCTION AND RELATED WORK
In many robotics applications, e.g., robotic weed control (Wu et al. 2020), the ability to accurately estimate the poses of objects is prerequisite for successful deployment. However, compared to other automation tasks, which primarily involve either classification or regression in IR", pose esti- mation is particularly challenging because the 3D rotation group SO [F a curved manifold. As a result, standard probability distributions (e.g., the multivariate Gaussian) are not well-suited for modeling elements of the SO(3) set. Further, because the steps for interacting with object in the "mean" pose between two possible poses (Figure [T] will often fail when applied to the object when it is in of the non-mean poses, accounting for multimodality in the context of rotations is essential.
2021 is simple, The recently described implicit-PDF (IPDF) Murphy al.
'Pronounced "aqua ma' na'am".
2 All code to generate the datasets, train and evaluate the models, and generate the figures can be found at: anonymized for review>.
'SO(3) stands for "special orthogonal group in three dimensions". " with the "special" referring to the fact that all rotation matrices have determinant of one. See: bttps: /ttes:///lii /wwggs..oo//.///////.o.oo.oooo.. sclentit wwweentiiiimmmiciaaaemmmammm com/ root -of-fuunn ty few- of my- favor itborite spaces 3/ for popular science introduc-
tion S0(3).

---

Under review as a conference paper at ICLR 2023
elegant, and effective approach for modeling distributions on SO(3) that both respects the curvature of the rotation manifold and is in- herently multimodal. The IPDF model f is trained through negative sampling where, for each ground truth image/rotation pair (X. R1), set of NN 1 negative rotations {Ra J2 are sampled and a score is assigned to each rotation f(x, R.). The final density pl Ri X) is then approximated as pR IX) R softmax ()) [1]/ where 8 is a vector contain- ing the scores with s[1] = 81, and V = n²/N, i... the volume of (3)3) split into N pieces. While effective, the N hyperparameter induces a non-trivial trade-off between model precision (in terms of the maximum likelihood that can be assigned to rotation) and inference speed in environments that do not have the computational resources necessary to extensively parallelize the task. Further, again due to the precision/speed trade-off, IPDF is trained with Numin < Nee which can make it difficult reason about how the model will behave in the wild.
Ro
R
R2
R
Figure 1: When minimizing the unimodal Bing- ham loss for the two rotations R1 and Rg. the maximum likelihood estimate R a rotation that was never observed in the dataset. Note, the die images are for demonstration purposes only, iee. no images were used during optimization. Ro is the identity rotation.
matrix as s;
The alternative to implicitly modeling distributions on SO(3) is explicitly modeling them. Here, I briefly describe the three baselines used in Murphy et al. (202I), which are representative of explicit distribution modeling approaches. Notably, IPDF outperformed each of these models by wide margin in distribution modeling task (see Table 1 in Murphy et al. (202T)) First, Prokudin et al.) (2118) described biternion network (Beyer al.] 2015) trained to output Euler angles by opti- mizing a loss derived from the von Mises distribution. The two multimodal variants of the model consist of: (1) a variant that outputs mixture components for von Mises mixture distribution, and (2) an "infinite mixture model" variant, which is implemented as a conditional (Sohn et al., 2015)) variational autoencoder (Kingma & Welling 2014). Second, Giiitschenssii et al. (2020) described procedure for training network to directly model distribution of rotations by optimizing loss derived from the Bingham distribution (Bingham 1974), with the multimodal variant outputting the mixture components for Bingham mixture distribution. Lastly, Deng et al.] (2020) extended the work of Gilitschenski et al. (2020) by optimizing a "winner takes all" loss (Guzmán-rivera et al., 2012 Rupprecht et al.) (017) in an attempt to overcome the difficulties (Makansi el al. 2019) com-
monly encountered when training Mixture Density Networks (Bishop 1994).
One additional approach that needs be introduced is the direct classification of an individual ro- tation from set of rotations, which, in some sense, is the explicit version of IPDF. Unfortunately, the computational complexity of this strategy quickly becomes prohibitive as more precision is re- quired. For example, Murphy et al. (202I) used an evaluation grid of ~224 million equivolumetric cells (Yershova et al. 2010), which would not only require an extraordinary number of parame- ters in the final classification layer of a model, but would also require an extremely large dataset to reasonably "fill in" the grid due to the curse of dimensionality
In this paper, I introduce AQuaMaM, an Autoregressive, Quaternion Manifold Model that learns arbitrary distributions on (3(3) with a high level of precision. Specifically, AQuaMaM models the projected components of unit quaternions domain of values. Architecturally, AQuaMaM is a Transformer (Vaswani eet al. 2017). Although Transformers were originally motivated by language tasks, their fiexibility and expressivity have allowed them to be successfully applied to a wide range of data types, includ- ing: images (Parmar et al.) 2018 Dosovitskiy al., (02T), tabular data (Padhi et al. 2021 Fakoor lt al. |220 Alcom & Nguyen 2021c), multi-agen trajectories (Alcorn & Nguyen 2021a), and proteins (Rao et al.] 2021). Recently, multimodal Transformers (Ramesh et al. 2021 Reed et al. 2022). i.c., Transformers that jointly process different modalities of input (c.g., text and images),
mixtures of uniform distributions that partition their
geometrically-restrices
have revealed additional fascinating capabilities of this class of neural networks. AQuaMaM mul-
timodal in both senses of the word- it processes inputs from different modalities (images and unit quaternions) to model distributions on the rotation manifold with multiple modes.
*In Murphy et al. (2021) Numin 4,096 and Nus = 2,359,296, i.c., Nunin/iuss = 0.2%. Notably, Mahendran al. 2018 used maximum of 200 rotations in their classification approach.
2

---

Under review as a conference paper at ICLR 2023
To summarize my contributions, find that:
AQuaMaM is highly effective modeling arbitrary distributions on the rotation manifold. On an "infinite" toy dataset with ambiguous viewpoints, AQuaMaM rapidly converges to a sampling distribution closely matching the true data distribution (Section 4.1). In contrast, the sampling distribution for IPDF dramatically diverges from the true data distribution, despite IPDF approaching its theoretical minimum evaluation loss during training. Ad- ditionally, on a constructed dataset of 500,000 renders of die in different rotations, an M model trained from scratch reaches log-likelihood 14% higher than an IPDF model using pretrained ResNer-50 (Section
AQuaMaM is fast, reaching a prediction throughput 52 faster than IPDF on single GPU.
THEORY
1.0
0.30
3
125
0.25
0.8
0
0.5
0.20
75
0.15
1.0
0.10
1.0
0.5
0.0
0055
25
0.05
0.0
0.5
0.00
1.0
-1.0
-1.00
-0.5
400
Sample
600
1000
(a)
(b)
(c)
Figure 2: Because there is a bijective mapping between the unit disk 2== =v E R |o|| < 1}
and the unit hemisphere 52 (vr E R3 | = 1, > 0), the challenging task of estimating distribution the curved 22 manifold be simplified to estimating distribution on the non- curved B², (a) Here, the true distribution on $2 uniform distribution, i.c., each point has density of 2 (because 5 has a surface area of 2T)). (b) Points that are uniformly sampled from 52 and then projected onto B² are more concentrated towards the edges of B due to the curvature of S2. If we model the distribution of (x: y) coordinates on B² as a mixture of uniform distributions, we can calculate p(x, y, 2) by dividing p(r,, y) by the area of the parallelogram defined by the Jacobian located at (æ, 2) on the hemisphere. (c) The p(æ, y, calculated through this procedure are generally quite close to the expected density. The mean density μ of the 1,000 points shown in (c) is 0. 154 (compared to 0.159 for the true density). A similar procedure is used by AQuaMaM to obtain the probability of a unit quaternion p(q) while only modeling the first three components of q: Ors ly. and 4=- See Section A.] for additional details.
2.1 THE QUATERNION FORMALISM OF ROTATIONS
Several different formalisms exists for rotations in three dimensions, likely the most widely en- countered of which is the rotation matrix 3 3 orthonormal matrix that has determinant of one. Another formalism, particularly popular in computer graphics software, is the unit quater- nion. Quaternions (typically denoted as the algebra H) are extensions of the complex numbers, where the quaternion is defined as q qa + 9y. + k + 4w with Jz, Qy: Ie- Iw R and 2 j 2 ijk = 11 Remarkably, the axis and angle of every 3D rotation can be encoded 1), i.e., qw cos (0/2), ds e, sin(0/2),
as a unit quaternion q Hi, {q H ||qll by by sin(0/2), and (== ez (0/2) where e = (e.,, Eg. 6z] is unit vector indicating the axis of rotation and the angle 0 describes the magnitude of the rotation about the axis
2.2 MODELING DISTRIBUTION ON A "HYPER-HEMMIPHEES"
Due to the curvature of the 3-sphere 33 {u R |/|| 1}. directly modeling distributions on HI, is difficult. However, two facts about H, present an opportunity for tackling this challenging task. First, the unit quaternions g and q encode the same rotation, i.e., HI, double covers S0(3).
3

---

Under review as a conference paper at ICLR 2023
As result, we can narrow focus to the subset of unit quaternions with gu > 0. ie., H1 {q E H1 : 9v > 0} (practically speaking, we can convert any q in a dataset with 9u < 0 to -q). Second, because all q E H, have unit norm, qr- 9y, and 9₂ must all be contained within the unit 3-ball 1² {u R ||>|| < 1}. which is not curved. Therefore, p(ga: Qy 9s) can be estimated using standard probability distributions. Importantly. Vqq E Hi,, qu is fully determined by 9.. qy and 9a because of the unit norm constraint, i.e., there is a bijective mapping f B + Hi such that fqz: 9y 9x) = [1 9y gz qw] where qu = / - o3 o Together, these facts mean we can calculate p(q) by simply applying the appropriate density transformation to p(ge, By 4z)
Intuitively, thinking about the probability density of point its infinitesimal probability mass divided by its infinitesimal volume, the density p(za, Qy: 9=) needs to be "diluted" by however much its volume expands when transported to Hi Specifically. the diluting factor is 1/5g where Sq is the magnitude of the wedge product of the columns of the Jacobian matrix J for f. Similar to how the magnitude of the cross product of two vectors equals the area of the parallelogram defined by the vectors, the magnitude of the wedge product for three vectors is the volume of the parallelepiped
defined by the vectors an analogous density transformation for the unit disk and unit hemisphere.
/ In this case, $q = 1/9u. i.e., the diluting factor is qu
[ Figure 2
visualizes
2.3 MODELING SO(3) DISTRIBUTIONS AS MIXTURES OF UNIFORM DISTRIBUTIONS ON PROJECTED UNIT QUATERNIONS
Ix = 0.7
min( ( 0.77 (z) for By E (a., b,)
4y
S V
0.72
TT
Mig ai by)
otherwise
0.7
0.7
0.8
0.72
Figure 3: When modeling the conditional distribution p(qulg=) as a mixture of uniform distributions, the geometric constraints of the unit quaternion are easily enforced. Here, focus on non-negative bins for clarity, i.e., intervals (a,, by) where 0 S < b S 1, but the same logic applies negative bins. Given 4r = 0.7, we know that 1gyl S v1 0.72 because q has a unit norm. As result, the mixture proportion π; for any bin where v1 - 0.72 < a; must be zero. AQuaMaM enforces this constraint by assigning value of -oc to the output scores for "strictly illegal bins" during training For the remaining bins, the corresponding uniform distribution is (gqy ag bx) where b, min( 0.7 b.). i.e., the upper bound of the uniform distribution for the partially legal bin reduced to v1 0.72.
Using the chain rule of probability, we can factor the joint probability p(q= qy: 9:) as p(qz- qy: 9o) = p(gz ( 9x). To model this distribution, desiderata for each factor include: (1) the flexibility to handle complex multimodal distributions and (2) the ability to satisfy the unit norm constraint of E H4, By partitioning [-1, into N bins, each with width of 2 we can model the distributions for quv qy. and 4s as mixtures of uniform distributions where N controls the maximum precision of the model. Recall that the density for the uniform distribution Uja,)) with -o0 < a <
1
for z [a,b) Therefore, the density for qu is: otherwise
b < 0o is
b-a 0
N
1
N
p(g=)= = Σ
Tillyai (i b:) = mxlllgzi 0x, be)
bb
Th
2/N
Tr
2
(I)
'See the lecture notes here: sottsspaffttrame
pdf for a survey of probability density transformations.
"See the lecture notes here: htttps:/ssheeLganaee
com/ spring2013 wtt.mma/mmamaaaaeeeme
cross_ pwwpucrrs.ppp...... ppf for additional connections between the cross product and the wedge product.
"See Section A.I for the exact recipe.
"See Section A.2 for a subtlety regarding "strictly illegal bins".
4

---

Under review as a conference paper at ICLR 2023
where ny is the mixture proportion for the mixture component/bin indexed by i, and is the index for the bin that contains 92-
For qy and dz. the conditional densities similar to the density for 4a. but the unit norm constraint on presents opportunities for introducing inductive biases into model (Figure 3 For example, consider the bins that are fully contained within R2 i.e., where a, b 2 For any such bin where
/1 L as, we know that π; 0 for 9y because || Further, again due the unit norm constraint, the conditional density for dy is in fact:
p(gylTa) = Wh min( /1 Lh. by) ax
i... when as < VI- 2h < bu. the upper bound of the bin's uniform distribution is reduced. Similarly, p(z:110 9x)
(x) = mt/(min( / o2 ag. bb) ax). The full density for p(gx, ty 2=) is
thus:
N
1
plg2. Iy 42) = Tux 2
T4=
TTz Tq T. 2wa Wa=
(3)
Way
Wg.
where Ty. is the mixture proportion for the bin containing component ge and Wg. is the potentially reduced width of the bin for qe used in Equation 2
2.4 OPTIMIZING p(q) AS "QUATERNION LANGUAGE MODEL"
After incorporating the diluting factor from Section 2.2 the final density for q E Hi is:
Nqw
p(q) = Agg Ta, Tgo 2way Wgz
(4)
As previously demonstrated by Alcorm & Nguyen (02Tb), models using mixtures of uniform dis- tributions that partition some domain can be optimized solely using "language model loss". Here, the negative log-likelihood (NLL) for a dataset X' is (where is the index for a sample):
N qd,
L= -= In Tea. + In Tag.y + In Aga,s + In
- d=1
(5)
2w0g Waa.
Notice that the last term is constant for a given dataset, i.e., it can be ignored during optimization. The final loss then 2 = E In Tga.. + In Tag.y + In Tga which is exactly the loss of three token autoregressive language model.
As a final point, it is worth noting that the last term in Equation is bounded below such that Nqw N i.e., for given Tg. Ta, Tg. the likelihood increases least cubically with the 2n0e Wez
number of bins. For 50,257 bins (i.e., the size of GPT-2/3's vocabulary (Radford et al., 2019) Brown al. 2020)), 3 1.26 1014 In contrast, the likelihood for IPDF only scales linearly with the number of equivolumetric grid cells (the equivalent user-defined precision hyperparameter).
3 AQUAMAM
3.1 ARCHITECTURE
Here, describe AQuaMaM (Figure 4) a Transformer that performs pose estimation using the Specifically, AQuaMaM consists of
autoregressive quaternion framework described in Section 2
See Scction A.3 for brief introduction Transformers.
5

---

Under review as a conference paper at ICLR 2023
20 [START]
4g = Bin(g=)
9a. (z)
4. Bin Bin((y)
Lax lan ba.
Zg. Dū, (9))
l. Bin((z)
X
R)
Transtormer
Attention Mask
740 2g, Zu
Patch Embeddings
Figure 4: An overview of the AQuaMaM architecture. Given image/rotation matrix pair (X, R). the image is first converted into a sequence of P patch embeddings while the rotation matrix is con- verted into its unit quaternion representation q = [ge: ay Iz: 9w By restricting the unit quaternions to those with positive real components (which is allowed because and -q encode the same rota- tion), qu becomes fully determined and does not need to be modeled. Next, each of the first two components 9e of is mapped to an embedding zq. by separate subnetwork 9g.- The full input to the Transformer is thus a sequence consisting of the P patch embeddings, a special [START] embedding 2o. and the two unit quaternion embeddings. The labels 4,_- La., and l, are generated by assigning qr. qy. and qs to one of N labels through a binning function Bin. Using a partially causal attention mask, AQuaMaM models the conditional distribution p(ue: By Q2 XX autoregressivelly i.e., p(gz. qy. 92 X) = p(g= X) X))p( gy 9.. X)p(q= 44. qy X) where each component is modeled as a mixture of uniform distributions that partition the component's s geometrically constrained domain. Because minimizing the loss of a mixture of uniform distributions is equivalent (up to a constant) to mini- mizing the classification loss over the bins, AQuaMaM is trained "quaternion language model".
Vision Transformer (ViT) (Dosovitskiy et al. 2021) with few modifications. As with ViT, AQua- MaM tokenizes an input image X by splitting it into P patches that are then passed through linear layer to obtain dimdte-dimeniional embeddings. Three additional embeddings are then appended to this sequence: (1) a [START] embedding 20 (similar to the [START] embedding used in lan- guage models) and (2) two embeddings Zq_ and kq, corresponding to the first two components of q. Each cq. is obtained by passing ge through subnetwork 9g. consisting of NeRF-like positional encoding function (Vaswani et al. 2017 Mildenhall et al. 2020; Tancik et al. 2020) followed by a multilayer perceptron (MLP), Le., the size of the input for the MLPs is 1 2 x L where L is the number of positional encoding functions.
Once positional embeddings have been added to the input sequence, the full sequence is then fed into a Transformer along with a partially causal attention mask (Alcorn & Nguyen 2021a b). which allows AQuaMaM to autoregressively model the conditional distribution p(gx: ly Ia XX using the density from Equation 4 The transformed zo. Zyz, and Zgy embeddings- -20. &q.. and 2q. -are used assign probability to the labels for Qas qy, and qs. e.g., Tqp h ((Ez)) [4 where is the classification head and l4- Bin(gz) where Bin is the binning function |1 Notably, AQuaMaM somewhat sidesteps the curse of dimensionality by effectively learning three different models with shared parameters: p(q.. X). p(qylgz, X), and p(q= 19g. qu X), each of which only needs classify N bins, opposed learning single model that must classify the order of N bins.
3.2 GENERATING SINGLE PREDICTION
While AQuaMaM could generate single prediction by optimizing q under the full density (Equa- tion 44) similar to what was done in Murphy et al. (022iipreiiction obtained via optimiza- "quaternion language model"
tion slow. An alternative strategy is embrace AQuaMaM and use a greedy search to return a "quaternion sentence" as a prediction. While sentences cor- responding regions where qu a sufficiently large N, the range of rotations can be kept reasonably small. For example, when N 500 (as in the AQuaMaM die experiment; see Section 4.2) if qu E [0.996, 1], then qhe E [0, VI 0.996"), and the geodesic distance between the rotations corresponding to the quater- nions q: = [0.996. 0, 0, v1 0.9962) and q2 = [1, 0, 0, 0] is 10.3*. Using the analogous numbers
will contain wider range of compatible rotations, with
1 Py Torch pseudocode for AQuaMaM's forward pass can he found in Listing 4 in Scction A.6

---

Under review as a conference paper at ICLR 2023
for N 50,257 (as in the AQuaMaM toy experiment; see Section 4.1 the geodesic distance be- tween 91 and 42 is 1.0%
Generating quaternion sentence with AQuaMaM requires three passes through the network. How- ever, the last two passes can be considerably optimized by noting that the first P + 1 tokens do not depend on the last two tokens due to the partially causal attention mask. A naive decoding strategy where a full (padded) sequence is passed through AQuaMaM three times has an attention complexity of O + 3)²). By only passing the first P + 1 tokens through AQuaMaM and caching the attention outputs, the complexity of the attention operations when decoding reduces to O((P + 1)2 + (P + 2) + (P- + 3)). Empirically, I found that using this caching strategy increased the prediction throughput by approximately twofold.
4 EXPERIMENTS AND RESULTS
Compared to the standard datasets used for computer vision tasks like image classification (Deng et al., (009) and object detection (Lin et al. 2014)- which contain millions of labeled samples-mane many commonly used pose estimation datasets are relatively small, with PASCAL3D+ (Xiang et al. 2014) con- sisting of ~200000 training images and T-LESS (Hodan et al. 2017) consisting of ~38,000 training images. The SYMSOL I and II datasets created by Murphy et al. (202T) consisted of 100,000 renders for each of eight different objects, but the pub- licly released versions only contain 50,000 renders per object [3 In addition to their relatively small sizes, these different pose datasets have the additional complication of including multiple categories of objects. The "canonical" pose for a category of objects arbitrary, so is not entirely clear whether combining categories of objects makes pose estimation harder easier for a learning algorithm (a jointly trained model essentially needs to do both classification and pose estimation) Therefore, to in- vestigate the properties of AQuaMaM in moderately large data regimes where Transformers excel, trained AQuaMaM two constructed datasets for pose estimation, an "infinite" toy dataset (Section 4.1) and dataset of 500,000 renders of die (Section 4.2) Because Murphy et al. 2021) found that IPDF outperformed other multimodal approaches by wide margin in their distribution estimation task, only compare IPDF baseline here.
NILL Theeretica Minimum NL1
0.5
0.4
0.3
(Valid) IPDF (Train)
0.2
AQuaMaM
0.1
0.0
20000
40000
Time (s)
Figure On the infinite toy dataset, AQuaMaM rapidly reached its theoretical minimum (classification) average negative log-likelihooe (NLL). In con- trast, IPDF never reached its theoretical minimum validation NLL, despite converging its training theoretical minimum.
4.1 TOY EXPERIMENT
The infinite toy dataset consisted of six categories of viewpoints indexed = 0, 1,... 5. associated each viewpoint with 2' randomly sampled rotation matrices R; {R, } i.c., each viewpoint had a difffrent number of rotation modes reflecting its level of ambiguity. I defined the true data- generating process as a hierarchical model such that a sample was generated by first randomly selecting viewpoint from uniform distribution over i, and then randomly selecting a rotation from uniform distribution over R, Given this hierarchical model, calculating the average log- likeeiioloood (LL) in the limit of infinite evaluation test samples simply involves replicating each (i, R2) pair 25-1 times and then averaging the LLs for each sample, Importantly, because the infinite dataset has identical training and test distributions, is impossible for a model to overfit the dataset.
12 Using the standard filtering pipeline where occluded and truncated images are excluded.
The publicly released dataset can be found at: http::///www......oooooooooaa/a/ wwtp::/wwwwwwwwootoooeeeeoeeeeee
//t::///wwww www.w:wwww..o.oo.oo..ooo.ooo.o..
catal oo/ symmmtmmmttr c__0ll
Further confusing the issue is the fact that many authors augment the PASCAL3D+ dataset with synthetic data, which makes it difficult to separate the influence of the architecture vs. the data augmentations when evaluating model performance.

---

Under review as a conference paper at ICLR 2023
AQuaMaM
Category
Category
Category 2
Category
Category
Category
1.0
1,01
0.4
0.2
0.3-
0.1
0.3-
0.26
proportion
Proportion
0.0
0.0
0
1
0.0
1
0.0
0.00
0 Category 4
0.00
1 Category
Category 3
Category
Category
Category 3
0.050
0.5
0.02
0.05
0.025
0.25-
12 4
10
15
1015202530 30
012
10
0 10152025.00
Rotation
Uniform
Rotation
Uniform
(a)
(b)
Figure 6: (a) The proportions of sampled rotations from the AQuaMaM model trained on the in-
finite toy dataset closely approximate the expected uniform distributions. (b) In contrast, despite approaching its theoretical minimum log-likelihood during training (Figure (5). the proportions of sampled rotations from the IPDF model drastically diverge from the expected uniform distributions.
Using this toy dataset, I trained an AQuaMaM model with N - 50,257 15 Because I was only in- terested in the distribution learning aspects of the model for the toy dataset, the category for each training sample was mapped to a Pémadd-ddieentin quence of P -mdderdiimnnsonn patch embeddings. also trained an IPDF baseline using the same hyperparameters and learning rate schedule from Murphy et al. (2021) for the SYMSOL. I dataset, except that increased the number of iterations from 100,000 to 300,000 after observing that the model was still improving at the end of 100,000 updates. Like the AQuaMaM model, I replaced the image processing component IPDF, ResNet-50 (He et al. 2015), with dReeNet-50 soddmenssoo embeddings where dRe = 2.048. In total, the AQuaMaM model had 3,510, 5,510,609 parameters and the IPDF model had 748,545 parameters; however, 93% AQuaMaM's parameters were contained in the final classification layer, with only 243,904 being found in the rest of the model.
embedding that was converted into se-
The final average LL for AQuaMaM on the toy dataset was 27.12. In comparison, IPDF reached an average I.I of 12.32 using the 2.4 million equivolumetric grid from Murphy al. (00T)) Note that the theoretical maximum LL. for IPDF when using 2.4 million grid is 12.38, which AQuaMaM comfortably surpassed. In fact, IPDF would need use a grid of nearly six trillion cells for it to even be theoretically possible for it to match the LL of AQuaMaM.
Figure 6 shows the results of generating 40,000 samples from the hierarchical model while using AQuaMaM and IPDF to specify the rotation distributions for each sampled viewpoint 16] For each sampled rotation, the rotation was assigned the ground truth rotation mode from the sampled viewpoint with the smallest geodesic distance. For AQuaMaM, the average distance was 0.04, while the average distance for IPDF was 0.84%. Further, the sampled rotations for AQuaMaM closely matched the true uniform distributions for each viewpoint. In contrast, IPDF's sampled distribution drastically diverged from the true data distribution despite its seemingly strong evaluation LL relative to its theoretical maximum. Notably, of the 40,000 quaternion sentences generated by AQuaMaM, only 23 (0.06%) were not from the true distribution.
4.2 DIE EXPERIMENT
The die dataset consisted of 520,000 renders of die (the same seen in Figure 1 where the die was randomly rotated to generate each image, and the dataset was split into training/validationn validationntsttt set sizes of 00000000000111100000000000000000
Similar the SYMSOL II dataset from Murphy et al.) (2021), different viewpoints of the die have differing levels of ambiguity, which makes the die dataset ap- propriate for exploring how AQuaMaM handles uncertainty. Note that, while 500,000 samples is a moderately large training set, for N 500, there are over 65 million unique token sequences
Full training details for both experiments can be found in Section A.5
16See Section A.4 for additional details.
17ehe die renders were generated using the renderer from Alcorn al. 2019 and the 3D die model from:
https://g / ttwp://tttttuueec com rrraa/aa// 2/ cubes
8

---

Under review as a conference paper at ICLR 2023
15.0
13.0
8.5
2.0
4.5
11.0
11.0
473
24.0
24.0
30.5
50.5
-88
50.
(a)
(b)
Figure 7: AQuaMaM effectively models the uncertainty of different viewpoints for the die object. In the plots, each point corresponds to rotation vector 0 = θe, i..., all rotation vectors are contained within a ball of radius л. Both the color and the transparency of each point are a function the log density In(p(q)) with redder, more opaque points representing more likely rotations and bluer, more transparent points representing less likely rotations. 1,000 rotations are depicted each pnot-one rotation (labeled 0) corresponds to the ground truth, 499 rotations were obtained by sampling from the AQuaMaM distribution, and the remaining 500 rotations were randomly selected from the uni- form distribution on SO(3). Each image of die corresponds to one of the points in its respective plot. The number before the colon indicates the associated annotated point, and the number after the colon is the In(p(g)) for that pose. (a) Because the five side of the die is facing the camera in the die's default pose, all rotations showing only the five occur along the (0, 0, 1) axis, which is reflected by the locations of the red points in the plot. Additionally, there are four high probability regions along the axis corresponding to the four rotations that could produce an identical image (see Figure [S in Section A.8 for plot of rotations sampled from this density). (b) For this unambiguous viewpoint, almost all of the probability is concentrated at the correct rotation, which causes the other points to fade entirely from view (the maximum In(p(q)) among the random rotations was -74.7).
covering (3) 1 and only 135 of the 10,000 token sequences in the test set were also found in the training set, Iee., AQuaMaM must still be able to generalize. I trained a ~~2 million parameter AQuaMaM with N = 500 on the die dataset in addition to a ~66 million parameter IPDF baseline that was identical to the model trained on the SYMSOL dataset in Murphy et al. (202T).
The average I.I. for AQuaMaM on the die dataset was 14.01 compared to 12.29 for IPDF, i.e., IPDF would need to use grid of over 12 million cells for to even be theoretically possible for it to match the average LL of AQuaMaM. When generating single predictions, AQuaMaM had mean prediction error of 4.32 compared to 4.57% for IPDF, i.c., a 5.5% improvement In addition to reaching a higher average LL and a lower average prediction error, both the evaluation and prediction throughput for AQuaMaM were over 52 faster than IPDF on single GPU. Lastly, as can be seen in Figure 7 AQuaMaM effectively models the uncertainty of viewpoints with differing levels of ambiguity. Figure 7 uses a novel visualization method where low probability rotations are more transparent, which allows all of the dimensions describing the rotation be displayed, in contrast to the visualization method introduced by Murphy et al. 2021 20
5 CONCLUSION
In this paper, have shown that a Transformer trained using autoregressive "quaternion language model" framework is highly effective for learning complex distributions on the SO(3) manifold. Compared to IPDF, AQuaMaM makes more accurate predictions at much faster throughput and produces far more accurate sampling distributions. The autoregressive factoring of the quaternion could be applied to other multidimensional datasets where high levels of precision are desirable, e.g., modeling the trajectory of a baskethall as in Alcor & Nguyen (2021a), which was partial inspiration for this work.
$$500 x (Ar + 2 Ro 65,000,000, i.c., the fraction of grid cells in a cube with a side of two that fall inside the unit sphere.
'See Scction A.7 for additional details.
AAdiitional experiments for cylinder dataset and a mixture of Gaussians design can be found in Section
A.9

---

Under review as a conference paper at ICLR 2023
REFERENCES
Michael A Alcorn and Anh Nguyen. baller2vec: A multi-entity transformer for multi-agent spa- tiotemporal modeling. arXiv preprint arXiv:2102.03321, 2021a.
Michael A Alcorn and Anh Nguyen. baller2vec+4: look-ahead multi-entity transformer for mod- eling coordinated agents. arXiv preprini urXiv:2104. 1980, 2021b.
Michael A. Alcorn and Anh Nguyen. The deformer: An order-agnostic distribution estimating transformer. ICML Workshop on Invertible Neural Networks, Normalizing Flows, and Explicit Likelihood Models (INNF+), 2021c.
Michael A Alcorn Li Zhitao Gong, Chengfei Wang, Long Mai, Wei-Shinn Ku, and Anh Nguyen. Strike (with) pose: Neural networks are easily fooled by strange poses of familiar objects. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp-
4845 4854, 2019.
Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning to align and translate. In International Conference on Learning Representations, 2015.
Lucas Beyer, Alexander Hermans, and Bastian Leibe. Biternion nets: Continuous head pose regres- sion from discrete training labels. In German Conference on Pattern Recognition, pp- 157-168. Springer, 2015.
Christopher Bingham. An antipodally symmetric distribution on the sphere. The Annals of Statistics, pp. 1201-1225, 1974.
Christopher M Bishop. Mixture density networks. 1994.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhari- wal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agar- wal, Ariel Herbert- Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel Ziegler, Jeffrey Wu, Clemens Winter, Chris Hesse, Mark Chen, Eric Sigler, Ma- teusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCan- dlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models few-shot learners. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), Ad- vances in Neural Information Processing Systems, volume 33, pp. 1877-1901. Curran Asso- ciates, Inc., 2020. URL 1457c0d6bfcb49e 2t7.:o6666o.001212..o.aaa..22222
mww.eeediiinnnneeeeeeeeeeeee pdf
cc/ppaper/222//11
Haowen Deng, Mai Bui, Nassir Navab, Leonidas Guibas, Slobodan Ilic, and Tolga Birdal. Deep bingham networks: Dealing with uncertainty and ambiguity in pose estimation, 2020.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hier- archical image database. 2009 IEEE Conference on Computer Vision and Pattern Recognition, pp.
2009. doi: 10. 109/00/R.20009 2209006688....
5206848..
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszko- reit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recogni- tion at scale. In International Conference on Learning Representations, 2021. URL https: //openzeview. www.eee//ooooo http://wwwww.etttoooooo////ooaa. dNNTYY
Rasool Fakoor, Pratik Chaudhari. Jonas Mueller, and Alexander J Smola. Trade: Transformers
for density estimation. Invertible Neural Networks, Normalizing Flows, and Explicit Likelihood Models Workshop, 2020.
Igor Gilitschenski, Roshni Sahoo, Wilko Schwarting, Alexander Amini, Sertac Karaman, and Daniela Rus. Deep orientation uncertainty learning based on a bingham loss. In International Conference Learning Representations, 2020. URLhttps::/ wwtpr/vveeeeeeeeeeeee/ee net. forum?
[d=ryylogSSDDS
Alex Graves. Generating sequences with recurrent neural networks.
arXiv preprint
arXiv: arXiv:1308.000 0850, 2013.
10

---

Under review as a conference paper at ICLR 2023
Alex Graves, Greg Wayne, and Ivo Danihelka. Neural turing machines. arXiv preprint arXiv:1410.5401,200 2014.
Abner Guzmán-rivera, Dhruv Batra, and Pushmeet Kohli. Multiple choice learning: Learning to produce multiple structured outputs. In F. Pereira, C.J. Burges, L. Bottou, and K.Q. Wein- berger (eds.), Advances in Neural Information Processing Systems, volume 25. Curran Asso-
ciates, Inc., 2012. URL https: /tttis::/ppcceeeddggs.... neurips httpppps....ppppe/////////// pdf
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. arXiv preprint arXiv:1512.03338; 2015.
Dan Hendrycks and Kevin Gimpel. Gaussian error linear units (gelus). arXiv preprint urXiv:16066 084 2016.
Tomáš Hodan, Pavel Haluza, Stepán Obdržálek, Jiri Matas, Manolis Lourakis, and Xenophon Zab- ulis. T-less: An rgb-d dataset for 6d pose estimation of texture-less objects. In 2017 IEEE Winter Conference on Applications of Computer Vision (WACV), pp. 880-888. IEEE, 2017.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In International
Conference on Learning Representations, 2015.
Diederik P. Kingma and Max Welling, Auto-Encoding Variational Bayes. In 2nd International Conference on Learning Representations, ICLR 2014, Bansf, AB, Canada, April 14-16, 2014, Conference Track Proceedings, 2014.
Tsung-Yi Lin, Michael Maire, Serge Belongie, James Hays, Pietro Perona, Deva Ramanan, Piotr Dollár, and C. Lawrence Zitnick. Microsoft coco: Common objects in context. In David Fleet, Tomas Pajdla, Bernt Schiele, and Tinne Tuytelaars (eds.), Computer Vision ECCV 2014, pp- 740-755, Cham, 2014. Springer International Publishing. ISBN 978-3-319-10602-1.
Siddharth Mahendran, Haider Ali, and René Vidal. A mixed classification--egressii framework for 3d pose estimation from 2d images. In British Machine Vision Conference 2018, BMVC 2018, Newcastle, UK. September 3-6, 2018, pp. 72. BMVA Press, 2018. WRLnttp:: /bmwc2018.
contents paapttss/pppers//22222
pdf
Osama Makansi, Eddy Ilg, Ozgun Cicek, and Thomas Brox. Overcoming limitations of mixture den- sity networks: A sampling and fitting framework for multimodal future prediction. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 7144-7153, 2019.
Ben Mildenhall, Pratul P. Srinivasan, Matthew Tancik, Jonathan T. Barron, Ravi Ramamoorthi, and Ren Ng. Nerf: Representing scenes neural radiance fields for view synthesis. In ECCV, 2020.
Kieran A Murphy, Carlos Esteves, Varun Jampani, Srikumar Ramalingam, and Ameesh Makadia. Implicit-pdff Non- con-aarmmetric representation of probability distributions on the rotation manifold. In Proceedings of the 38th International Conference on Machine Learning, pp. 7882-7893, 2021.
Inkit Padhi, Yair Schiff, Igor Melnyk, Mattia Rigotti, Youssef Mrouch, Pierre Dognin, Jerret Ross, Ravi Nair, and Erik Altman. Tabular transformers for modeling multivariate time series. In ICASSP 2021-2021 IEEE International Conference Acoustics, Speech and Signal Process- ing (ICASSP), pp. 3565-3569 IEEE, 2021. URL https://Leexxlore. www.eeeeeeeegg.... document 94144
Niki Parmar, Ashish Vaswani, Jakob Uszkoreit, Lukasz Kaiser, Noam Shazeer, Alexander Ku, and Dustin Tran. Image transformer. In Jennifer Dy and Andreas Krause (eds.), Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pp. 4055 PMLR, 10-15 Jul 2018. URL httttt::/pprotto ngs.
m r. prw.s v80 ht/880ppaa.
rmar 8a. hamml
Sergey Prokudin, Peter Gehler, and Sebastian Nowozin. Deep directional statistics: Pose estimation with uncertainty quantification. In Proceedings of the European conference on computer vision (ECCV). pp. 534-551, 2018.
11

---

Under review as a conference paper at ICLR 2023
Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and llya Sutskever. Language models are unsupervised multitask learners. 2019.
Aditya Ramesh, Mikhail Pavlov, Gabriel Goh, Scott Gray, Chelsea Voss, Alec Radford, Mark Chen, and Ilya Sutskever. Zero-shot text-to-image generation. Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp. 8821-8831 PMLR, 18-24 Jul 2021. URL
https: https://proceedinge m wtt.:/ss.....pe./11313333.3 http99//aameh22aac..am
Roshan M Rao, Jason Liu, Robert Verkuil, Joshua Meier, John Canny, Pieter Abbeel, Tom Sercu, and Alexander Rives. Msa transformer. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp. 8844 8856. PMLR, 18-24 Jul 2021. URL https: http::tts///poooo http:::://prcced.nng..
m wtt.ppeessvvv/1ree..
html
Scott Reed, Konrad Zolna, Emilio Parisotto, Sergio Gomez Colmenarejo, Alexander Novikov, Gabriel Barth-Maron, Mai Gimenez, Yury Sulsky, Jackie Kay, Jost Tobias Springenberg, et al. A generalist agent. arXiv preprint arXiv:2205.06177 2022.
Christian Rupprecht, Iro Laina, Robert DiPietro, Maximilian Baust, Federico Tombari, Nassir Navab, and Gregory D Hager. Learning in an uncertain world: Representing ambiguity through multiple hypotheses. In Proceedings of the IEEE international conference on computer vision, pp. 3591-3600, 2017.
Kihyuk Sohn, Honglak Lee, and Xinchen Yan. Learning structured output representation using deep conditional generative models. C. Cortes, N. Lawrence, D. Lee, M. Sugiyama, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 28, Curran As- sociates, Inc., 2015. URL http//tttts://prrcedi.ggg neurips..eppppppp http://pwaaaaaooeeeeeeeeeeeeeppp (dd552249 oa55522226bba0007777777222222222
pdf
Matthew Tancik, Pratul Srinivasan, Ben Mildenhall, Sara Fridovich-Keii, Nithin Raghavan, Utkarsh Singhal, Ravi Ramamoorthi, Jonathan T. Barron, and Ren Ng. Fourier features let net- works learn high frequency functions in low dimensional domains. NeurlPS, 2020.
Benigno Uria, lain Murray, and Hugo Larochelle. Rnade: The real-valued neural autoregressive density-estimator In C.J. Burges, L. Bottou, M. Welling, Z. Ghahramani, and K.Q Wein- berger (eds.), Advances in Neural Information Processing Systems, volume 26. Curran Asso- ciates. Inc., 2013. URL https: http::://cpppcccpppeeeoe neurips cc/paaper/2013/fill .pdf
$3adaf494dcc8er 961111966666666666667666
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Infor- mation Processing Systems, pp. 5998-6008, 2017.
Jason Weston, Sumit Chopra, and Antoine Bordes. Memory networks. In International Conference Learning Representations, 2015.
Xiaolong Wu, Stéphanie Aravecchia, Philipp Lottes, Cyrill Stachniss, and Cédric Pradalier. Robotic weed control using automated weed and crop classification. Journal of Field Robotics, 37(2): 322-340, 2020.
Yu Xiang, Roozbeh Mottaghi, and Silvio Savarese. Beyond pascal: benchmark for 3d object detection in the wild. In IEEE winter conference applications of computer vision, pp. 75-82.
IEEE, 2014.
Anna Yershova, Swati Jain, Steven M Lavalle, and Julie C Mitchell. Generating uniform incremental grids on so (3) using the hopf fibration. The International journal robotics research, 29(7) 812. 2010.
12

---

Under review as a conference paper at ICLR 2023
A APPENDIX
A.1 TRANSFORMING A DENSITY ON THE UNIT DISK TO A DENSITY ON THE UNIT HEMISPHEREE
Here, briefly walk through the diluting procedure for 2D example that is directly analogous to HI and B Consider the task of approximating uniform distribution the unit hemisphere 8 {o R ||사|| and vs 0} by modeling the projected coordinates of 8 mixture of uniform distributions on the unit disk 8 {or E R² ||a|| < 1} (Figure Note that f(z.8) v) = Iz. y, and:
1 0
0
y². The columns of J define parallelogram tangent to the hemisphere at (1) with area equal to the square root of the sum of the squared 2 2 minors in J 21
where: 2 V x² -
i.c.
0
2
x
+ 1
x
1
22
22
22
corresponds to the top of the hemisphere where the tangent parallelogram is parallel to B and there is pansion/dilution. In contrast, as - O, the tangent parallelogram becomes increasingly "stretched" leading to greater dilution of p(r,y).
We can now calculate p(z, y, p(z, 2. p(a.u
p(z,y).. Intuitively,
The recipe for calculating 8y is nearly identical, with:
1
0
0
0
J
0
0
1
Iw
Tw
Tu
and 3 x 3 minors being used in place of 2 x minors. As result, 8n n.. i.e., the diluting factor for AQuaMaM is lw-
A.2 DEFINING "STRICTLY ILLEGAL BINS"
The definition of "strictly illegal bins" used by AQuaMaM is best motivated by an example. Con- sider model that was trained on dataset consisting of single rotation. Python pseudocode for a naive sampling algorithm can be found in Listing where the pet-probs .naiie viee function as- signs probability of zero to any bins where all of the values would violate the unit norm constraint given the components of q generated up to that point. This naive sampling procedure will generate sequences of bins that were not found in the training dataset. For example, assume the q for the single rotation had yz 0.42 and 4y = 0.901 and that N 20 for the model. Given globally
optimal model, åz will be sampled from [0.4, 0.5]; however, if d. > vi 0.92 0.436, then get.probs_naiiee will assign zero probability to the bin [0.9, 1.0).
def anive_sample_ge (AQuuMMa( (AQuaMaM)
q_cs []
for range (3):
bin_probs AquaMaMMgett gwwggeep____eeeeeeeeee (g_os)
gg categorical (bin_prebs)
This procedure resembles the "formal determinant" method for computing the product: https: //en. wik k pedía www.ar/wwww/...sss/e www.rooss__roooccetaaariii.c.ee
2he area of the parallelogram can also he computed by taking the square root of the determinant of the
Gram matrix G J.
13

---

Under review as a conference paper at ICLR 2023
(lower, upper) AQuaMaM.Bin edges q_
cs appaad hontredddooottttataeeen.u..momm (lower, upper, __s))
(e_x, g_,, __) = g_es
P sqrt - x y 9_2++2)
(q_x, Ly,, -,,, [_x)
Listing I: Python pseudocode for naively generating a unit quaternion sample with AQuaMaM.
To avoid this issue, AQuaMaM defines "strictly illegal bins" as those that would violate the unit norm constraint given the components of generated up to that point when using the minimum magnitudes of their corresponding bins. Effectively, during training AQuaMaM must learn to assign zero probability to bins bordering "edge bins" rather than being told that they have zero probability (which is the case for strictly illegal bins). Python pseudocode for the final sampling algorithm can be found in Listing 2 In summary, the sampling procedure involves first generating a sequence of bins, and then sampling from the region defined by the edges of those bins using rejection sampling. When viewed as generative model, AQuaMaM can be interpreted as modeling noisy measurements of unit quaternions.
def sample_ (AQuaMaM)
The q_c_hats are used for generating the bins, but are not the final g_es.
q_c_hats - []
bell_edges
for i in range
bin probs AQuaMaM. get Batpaaaaagee_ppbbb///eee (q___haatss
1_g_ categorical (bin_prcbs) (lower, upper) = @guaMaM.bin__dgee. [1_q_c] c_c_hats, append (uniform (lower, upper))
cell_edges append (lower, upper))
11
12
(e_x, yy,, 2 = seanttiionsmlmeeeellll q_ sqrt (1 - xx*22 yy++2 - q_2*+2)
13
return x qy, 2, x_x)
Listing 2: Python pseudocode for generating a unit quaternion sample with AQuaMaM.
A.3 BACKGROUND ON TRANSFORMERS
In this section. I briefly describe the Transformer Vaswann al. 2017) architecture at high level, and define the Transformer-specifie terms used in the main text. Readers who seek a deeper under- standing of Transformers encouraged explore the following excellent pedagogical materials:
The Illustrated Transformer by Jay Alammar: fitts://jaaanncen htt.jjaaaaaaaaaiiiiaa github.thh..... io/ i [ilsstrateddrranssteet
2. The Annotated Transformer by Sasha Rush, Austin Huang, Suraj Subramanian, Jonathan Sum, Khalid Almubarak, and Stella Biderman: http://nlp.seas. harvard.ed./
http:/ppp..aaa..aaaaaa.....d
3. Attention?
Attention! by Lilian Weng https://lHLianwenggiihhbb.
Transformers are class of neural networks that are capable of processing sequential data without being explicitly sequential. To understand what that means, it is instructive to contrast Transform- ers with recurrent neural networks (RNNs), which are explicitly sequential. Both Transformers and RNNs take as input sequence of N d-dimensional vectors and output a sequence of N d- dimensional vectors. During both training and testing, RNNs operate by iteratively applying the same neural network to each element of sequence and an evolving hidden state. In contrast, during training, Transformers process an entire sequence in parallel.
To accomplish this parallelism. Transformers use the attention mechanism (Graves 2013 Graves
al. 2014: Weston et al. 2015 Bahdanau et al. 2015
Intuitively, the attention mechanism is
Atteention is all [they] need"
Vaswani al. 2017
14

---

Under review as a conference paper at ICLR 2023
function that tells neural network how much to "focus" on the different elements of a sequence when processing a specific element of that sequence. The pure attention approach to sequence modeling provides two main benefits: (1) it can greatly speed up training because sequences processed in parallel through the network, and (2) long-term dependencies are much easier to model because distant elements of the sequence do not need communicate through many neural network layers (which is the case for RNNs)
While the attention mechanism is extremely powerful, it is also permutation invariant. As a re- sult, Transformers must encode positional information and conditional dependencies through other means, i.e., with position embeddings and attention masks. Consider the following sequences of word embeddings
W1 w [StArt] Whee Wiat Win: Wiee What
W2- w [StART] Wins Wher Weats Whe What
Because the attention mechanism is permutation invariant, W1 and W2 look identical to Trans- former. Therefore, to encode the order of the word embeddings, the Transformer adds a position embedding t; to each word embedding where i indicates the word embedding's position in the sequence, As a result, Wi becomes:
Wi = (START) + f Wiee + 22, Went + L3, Win + Las Wheh + L5, What + to)
and the attention mechanism would be applied to Wi and W2.
In order fully understand how Transformers model conditional dependencies, additional technical details about the attention mechanism are necessary. The basic attention mechanism in Transform- ers consists of architectural pieces: (I) a query/key/value function á and (2) a score function vb. The query/key/value function á independently maps each element of a (position encoded) se- quence Wij to query, key, and value vector, i.e.:
(q: k;, 0;} <(W())
The score function vú takes a query vector and a key vector k, as input and outputs a scalar score, ie.:
81, (1 k:))
The score matrix T is thus an N x N matrix where T [ = 8i.j. The final step of the attention mechanism involves calculating new vector for each element in the sequence:
W
Gi.j
j=1
where:
e #1
0i.j
km
i.e., the attention weights azj calculated by applying the softmax function to each row of L
Now consider the common natural language processing task of "language modeling". i... learning P(S) where S is sentence with N words. A common way factorize p(S) is
p(S) - ((0==((()))((()111)))
[STCP] |S[N 1I,..., 8(2),
" [STARC] is special embedding that appended to beginning sequences in both Transformers and RNNs in situations where the first element of the sequence is being modeled.
[STOP] ] is special token indicating the end of sentence.
15

---

Under review as a conference paper at ICLR 2023
which is using the chain rule of probability. To enforce such conditional dependencies, Transformers use an attention mask M, which is an N x matrix where M = if the Transformer is allowed to "look" at element W[]] of the sequence when processing element W[i], and Mi.j] = -oo if the "looking" is not allowed. Instead of applying the softmax function to each row of the raw score matrix V,, Transformers apply the softmax function to each row of masked score matrix L v M. As result, ai when the Transformer is not allowed to "look" at element W[i] of the sequence when processing element W[i] (because e e 0).
To model the chain rule factorization described earlier, M takes the form of strictly upper trian- gular matrix where all of the values above the diagonal are set to -30. Because a strictly upper triangular attention mask means each element in the sequence cannot be influenced by elements that occur later than that element in the sequence (i.e., the present is only affected by the past), this particular attention mask referred to as causal attention mask. Because AQuaMaM is model-
ing the distribution Hj conditioned on an image, the Transformer is allowed to "look" at all of the image patch embeddings when processing any single patch embedding, i... AQuaMaM uses partially causal attention mask, which depicted Figure
The full Transformer architecture consists of L blocks where each block (indexed by first applies the masked attention mechanism to the full sequence with de and the, and then passes each element of the resulting sequence through MLP fe, i.e.:
Wi s((W(i)
As previously mentioned, the pure attention approach to sequence modeling means Transformers can evaluate entire sequences in parallel, i.e., the likelihood of full sequence can be calculated in single forward pass. However, both making predictions and generating samples require N forward passes through the Transformer because the elements of the sequence can only be "filled in" one by one. The naive Transformer decoding algorithm thus starts off with a sequence of [NULL] embedding following the [START] embedding, e.g.
W W [SLART] W [NULL]: ; W [NUUL]] w (suli): W [NU_L]} w [NULL)]
and then iteratively generates the words of the sentence. As discussed in Section 3.2 the caching strategy used by AQuaMaM is motivated by the fact that this naive decoding algorithm computa- tionally wasteful with regards to the attention mechanism.
A.4 GENERATING SAMPLES
Generating samples with IPDF simply involves sampling from the probability distribution defined by IPDF over the equivolumetric grid. For AQuaMaM, the sampling procedure uses a conditional version of the algorithm found in Listing 2 Python pseudocode for the conditional AQuaMaM sampling procedure can found Listing
def sample_s __ (AQuaMaM, image):
The q_c_hats are used for generating the bins, but are not the final c_cs.
q_c_hats = []
cell_edges []
for in range
bin_probs
DonamaM.ge__prdatt (image, (_c_hats)
7
iqc categorical prcbs)
(lower, upper) = AQuaMaM.bin_edeee [odges/l_gg [_]]]]]]
e hatsa append (uniform (lower, upper))
10
cell_edges, append (lower, upper))
11
12
(g_x, g_y, 22) = w sgrt (1 - x+12 y++2 q_2+22)
For example, vector full of zeros, although the actual values irrelevant.
16

---

Under review as a conference paper at ICLR 2023
return X, q Yr 2, (.))
Listing 3: Python pseudocode for generating a unit quaternion sample conditioned on an image with AQuaMaM.
A.5 AQUAMAM TRAINING DETAILS
A.5.1 TOY EXPERIMENT
For the toy dataset, the hyperparameters for the AQuaMaM model dmodel 64, P 196
(i.e., the number 16 x 16 patches in a 224 x 224 image), eight attention heads, dr 256 (the dimension of the inner feedforward layers of the Transformer), three Transformer layers,
6,
9q. MLPs with 16. 32, and 64 output nodes and Gaussian Error Linear Unit (GELU) nonlinearities (Hendrycks & Gimpel 2016), and N 50,257. used the Adam optimizer (Kingma Ba 2015) with an initial learning rate of 10-4, 31 = 0.9, B3 = 0.999, and = 10 to update the model's parameters. The learning rate was reduced by half every time the validation loss did not improve for five consecutive epochs. The model was implemented in PyTorch and trained with a batch size of 128 on a single Tesla V V100 GPU for 2.5 hours (624 epochs) where each epuch consisted 40,000 training samples, and the evaluation loss was used for early stopping.
A.5.2 DIE EXPERIMENT
For the die dataset, the hyperparameters for the AQuaMaM model were dmodel 512, P = 196.
eight attention heads, dr 2,048, six Transformer layers, = 6. 9q. MLPs with 128, 256, and 512 output nodes and GELU nonlinearities, and N 500. In total, this AQuaMaM model had 20,000,75 parameters. The model was trained using an identical optimization strategy to the toy experiment, and trained for ~2 days (121 epochs).
17

---

Under review as a conference paper at ICLR 2023
A.6
THE AQUAMAM FORWARD PASS
Shape values come from Section A.5.
3 imgs has a shape of (128, 3, 224, 224).
qs has a shape of (128, 4)
self.z_0 has shape of (1, 1, 512).
self. self.possembed (the position embeddings) has shape of (1, 199, 512).
self.attn, mask has a shape self.h = nn. Sequentiii (nn. LayeLeeyNNerm552)) m..inear (512, 500)).
(199, 199).
def (erarddssee., imgs, qs):
patch_embeds has shape of (128, 196, 512).
11
patch_embeds self. gelf.patch__meeee (ings)
12
13
$ z_0 has a shape of (128, 1, 512).
14
2_0
(len
embeds), -1, -1)
15
r___x has shape of (128, 1, 512).
16
2_e_x selt.
01)
17
$ 2 has a shape of (128, 1, 512).
18
2_y = ( 11)
20
has a shape of (128, 199, 512), i.e., 199 x torch.cat (epatcaat(ppthhmmeeeeeeee
21
z_0, 2_g,,, y_ dim-1)
22
x = x + self.pos_ embed
# still has shape of (128, 199, 512).
self.transformer (x, self.attn, mask)
25
26
x has a shape of (128, 3, 500).
27
X = sclf.h(xli, -3:])
Section 2.3. still has shape of (128, 3, 500) x constrain_qs (gs, sel fiins, x)
30
31
32
return
Listing 4: PyTorch pseudocode for the AQuaMaM forward pass.
A.7 GENERATING PREDICTIONS
Generating prediction with IPDF simply involves taking the highest probability element from the probability distribution defined by IPDF over the equivolumetric grid. For AQuaMaM, generating deterministic prediction similar the steps for generating sample (See Listing 3 except the maximum probability bin is always selected and is set as the midpoint of the selected bin rather than sampled from it. Python pseudocode for the conditional AQuaMaM sampling procedure can be found in Listing
def predict_qe (AQuaMaM, (Acdict__ image)
The h_c_hats are used for selecting the bins, but not the final
c_cs.
__c_hats - []
for i in (3)
I AQuaMaM. get probs (image, hats) argmax ()
(lower, upper) AQuaMaM.bin edges q_a]
c.c. hachhat..pppen (lower upper) / 2)
if norm (q___hat))) <= 1: _cc h_t_hats
else:
12
c cs normaliza (hats
13
14
Q_w s1 xx*+2 g - q_2*+2)
return x qy, 2 x_x)
Listing 5: Python pseudocode for deterministically predicting a unit quaternion given an image with AQuaMaM.
18

---

Under review as a conference paper at ICLR 2023
A.8 SAMPLED ROTATIONS FOR FIGURE 10A
3
2
1
-1
1
20
8
-2
-3
Figure 8: 1,000 sampled rotations from the AQuaMaM density shown in Figure 10a equivalent modes is represented, and there are no points incorrect regions.
Each of the
19

---

Under review as a conference paper at ICLR 2023
A.9
ADDITIONAL EXPERIMENTS
A.9.1
CYLINDER EXPERIMENT
10.2002
9.6127
8.7908
122: 9.5357
15.0
2.0
44,
1
-11.0
21
122
in(pig))
-17.5
2
-24.0
-3
-30.5
-37.0
-2
80
-1
2+
-43.5
2
-5..0
3
3
Figure 9: AQuaMaM is also capable of learning distributions for objects with continuous symme- tries, as shown here for a viewpoint of the cylinder. 1,000 rotations are depicted in the plot-one rotation (labeled 0) corresponds to the ground truth, 499 rotations
obtained by sampling from the AQuaMaM distribution, and the remaining 500 rotations were randomly selected from the uni- form distribution SO(3). The cylinder's two axes of symmetry clearly visible high density curves in the rotation vector ball.
Following the same training setup as the die experiment (Section 4.2) trained an AQuaMaM and IPDF model cylinder") 27 dataset (notably, cylinder is of the objects in the SYMSOL I dataset). The average LL for AQuaMaM on the cylinder dataset was 7.24 compared to 5.94 for IPDE 28 Figure 9 visualizes the distribution learned by AQuaMaM using the same visualization technique described in Figure 7
A.9.2 MIXTURE OF GAUSSIANS EXPERIMENT
This section describes a mixture of Gaussians (MoG) variant of AQuaMaM, AQuaMaM-MoG, which is conceptually similar to how (non-curved) densities are modeled in RNADE Uria et al. 2013 As a preliminary, recall the change-of-variane technique 29
Let X be continuous random variable with generic probability density function f (ar) defined over the support 61 < C2- And, let Y u(X) be an invertible function of X with inverse function X .(Y). Then, using the change-of- variable technique, the probability density function of Y is:
27 The 3D cylinder model can found at: http:::///githh...coo/aaaaaaarrr
Codoe
blob/ master / Wavef ront. OBJ woader/Mcdells/// inder innde.ddrb.... .ob]
The reported average I.I. the cylinder dataset in Murphy al. (202)) 4.26, hut they used consid- erably smaller training dataset (100,000 vs. 500,000) updated the model considerably fewer times during training (100,000 vs. 300,000).
Taken verhatim from: https://on/ ine. str..stat.c.. . www.:e/tttttttteeeeeeeeeeeeeee
20

---

Under review as a conference paper at ICLR 2023
fr(v) = fx(o(y)) (v((0) defined over the support u((1) < < u(c2).
Let qe by + (Ug. l 1te where Ey. and Ug. define the lower and upper bounds for quaternion component de (i.e., la u...) ). and 84 is distributed according to MoG parameterized by AQuaMaM-MoG. In words, Sq. determines where falls in -ug Ugee through the logistic function. Solving for 8ge (to get the equivalent ely)) gives:
Ugo
2uge
+ e
21q
+
94
qe + ugg.
211ge
4e Ugg
2ug. - (ge + Uge
9e + Ugg
Ug. 9e
e 940
qe + +ug
8g.
In(ug - qa)) In In(1c + ugy))
In(g. + Uge)
In([uge
qr) 8q0
0.1
-0.5
0
0.5
45
0.5
(a)
(b)
Figure 10: Plots of the denominator for the [v'(y) term the change-of-variable formula when
using mixture of Gaussians approach with AQuaMaM. (a) = 1- o 1s and (b) 1- + 0.5%).
Diffcrentiating with respect to ge (to get the equivalent of v'()) gives:
1
- [ln(go + un.)) - In(ug. - 90)l
+
dg
4e + Ugg
Ug Ge
Uge De + qo + Uge
(1e + Ugg (ug. - y6)
20ho
u2 2z
Note that the minimum of juz q2l zero and occurs when de tug (see Figure 10 i.e..
(v(()) ge ug
21

---

Under review as a conference paper at ICLR 2023
The full AQuaMaM-MoG density is thus:
gd
211g=
p(g) = 7 Ruz
T Bag
Noas
Tu
(6)
L
4g
43
where Tay. is the density assigned to 8q. by the MoG parameterized by AQuaMaM-MoG. Because change-of- varriable terms are constant for a given dataset, the training loss for AQuaMaM-MoG
is:
X In Wd.5.g
In Td
In Td,
d=1
I trained an AQuaMaM-MoG model on the toy dataset described in Section 4.1 using 512 mixture
components and otherwise identical hyperparameters test set for AQuaMaM-MoG was 10.52, which is not only far worse than AQuaMaM (27.12), but is also considerably worse than IPDF (12.32). Further, the sampling distribution of AQuaMaM-MoG is far from the true data distribution, as can be seen in Figure
described in Section A.5.T The LL on the
22

---

Under review as a conference paper at ICLR 2023
3
2
1
0
-1
-2
3
3
2
1
-3
-2 -1
0
er
0
-1
1
-2
2
3
-3
Figure 11: Rotations sampled from AQuaMaM-MoG for the unimodal viewpoint in the toy dataset are often far from the true rotation (the red point).
23


Reviews:
[{"review_id": "fLmyHFwRP2", "reviewer": "Reviewer_oE5t", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper uses a quaternion model for capturing rotational distributions. The approach is branded as a new 'method' even if it appears more like a model construction. The efficiency over a baseline is demonstrated on simple data sets.\n", "summary_of_the_review": "This paper is interesting but appears to have multiple flaws.\n", "strength_and_weaknesses": "*Strengths*\n\n1. Quaternions are the standard way of modelling rotations in tracking, position estimation, and pose modelling, and thus most likely a sensible way of modelling them also for capturing complex rotational distributions in general.\n\n2. The paper is interesting, and could topic-wise be a good fit for the conference.\n\n*Weaknesses*\n\n3. The originality and novelty of the approach is questionable: This appears more like a standard model design choice for the particular application (you know that the phenomenon has rotational symmetry, and thus you incorporate this as prior knowledge in your model), and not as a new 'method'.\n\n4. The presentation could be improved. The text is wordy and could be stating things more directly and clearly. The abstract alone almost fills the first page.\n\n5. The experiments act as proof of concept. Additional experiments would have strengthened the paper.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "U9a2SOrrc7J", "reviewer": "Reviewer_WNVz", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper describes an approach to learning distributions on the SO(3) manifold via an autoregressive factorization of the distribution on quaternions, wherein each conditional and the first marginal are approximated by a mixture of uniform distributions.\n\nThe paper provides a comparison to the prior work of Murphy 2021, and shows some empirical advantages over that approach.", "summary_of_the_review": "The paper describes a new approach for modeling distributions on the SO3 manifold.  I recommend rejection because of little motivation for the choice of the method relative to other possibilities, and because of a lack of substantial positive empirical results to demonstrate that the method performs well.\n\nUpdate: I have changed my score from 3 to 5.", "strength_and_weaknesses": "Strengths:\n* The method is relatively simple and very computationally efficient as compared to the approach of Murphy.\n* The experimental results in figure 7 are illustrative.  This is a very nice choice of visual.\n\nWeaknesses:\n* The method is not very clearly motivated relative to other possibilities.  E.g. why use quaternions as the underlying representation on which to apply the autoregressive decomposition? And why use a mixture of uniforms rather than something else with a computable likelihood?\n* The empirical component of the validation is very limited.  The authors could strength the paper with comparisons to additional methods.\n* The experimental results in figure 7 A appear to be a negative result.  There is much probability mass spread throughout the space, not just on the four possible orientations given the top view of the “5”. \n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "7kP4vywTAcb", "reviewer": "Reviewer_KBHk", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper presents a novel method (AQuaMaM) for 3D quaternion orientation estimation from potentially ambiguous 2D images. It employs a Transformer architecture to learn sequences of quaternion parameters representing distributions over the SO(3) group of 3D rotations, treating them as language tokens. \n\nThe proposed architecture is able to efficiently learn multimodal distributions over SO(3), allowing to represent multiple legitimate candidate rotations corresponding to a given ambiguous 2D image.\n\nThe approach is validated on a toy dataset and a 3D die orientation estimation dataset, demonstrating higher accuracy and higher training and inference time efficiency with respect to a strong SoTA baseline (IPDF).", "summary_of_the_review": "Overall, in my view the paper makes a significant algorithmic contribution by presenting a novel Transformer-based approach to complex multi-modal distribution learning on SO(3). The proposed approach is demonstrated to be faster and more accurate and reliable than a strong baseline on two clear and controllable experimental setups.\n\nThe paper is generally well-written and pleasant to read. Some improvements would be required in terms of clarity, as detailed above.", "strength_and_weaknesses": "Strengths\n---------\n- To my knowledge, the proposed approach is novel. It seems to be the first successful application of a Transformer autoregressive model for multimodal distribution estimation on SO(3).\n- AQuaMaM is compared against a strong baseline (IPDF), demonstrating:\n   1) Significantly faster convergence time\n   2) Faster inference, since it requires few forward passes with respect to the $N$ required by IPDF\n   3) Higher accuracy\n   4) Higher reliability (see Fig. 6)\n- The paper is generally well-written. The introduction is particularly pleasant to read and positions the paper very well in the literature.\n\nWeaknesses\n------------\n- As detailed in the following, there are several points in which clarity shall be improved. This would especially help readers who may be familiar with the orientation estimation problem, yet not necessarily accustomed to language modeling and Transformers.\n- The die experiment is convincing and allows for detailed analysis in a controlled way. Still, an additional comparison on a more challenging dataset (e.g., the one employed in the original IPDF paper) would strengthen the paper by validating AQuaMaM on an additional, possibly harder problem.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject