Training and Inference on Any-Order Autoregressive Models the Right Way
Authors: Andy Shih
Dorsa Sadigh
Stefano Ermon
Venue: NeurIPS 2022 Accept (2022)
Keywords: any-order autoregressive models
tractable generative models
arbitrary marginal and conditional

Abstract: Conditional inference on arbitrary subsets of variables is a core problem in probabilistic inference with important applications such as masked language modeling and image inpainting. In recent years, the family of Any-Order Autoregressive Models (AO-ARMs) -- closely related to popular models such as BERT and XLNet -- has shown breakthrough performance in arbitrary conditional tasks across a sweeping range of domains. But, in spite of their success, in this paper we identify significant improvements to be made to previous formulations of AO-ARMs. First, we show that AO-ARMs suffer from redundancy in their probabilistic model, i.e., they define the same distribution in multiple different ways. We alleviate this redundancy by training on a smaller set of univariate conditionals that still maintains support for efficient arbitrary conditional inference. Second, we upweight the training loss for univariate conditionals that are evaluated more frequently during inference. Our method leads to improved performance with no compromises on tractability, giving state-of-the-art likelihoods in arbitrary conditional modeling on text (Text8), image (CIFAR10, ImageNet32), and continuous tabular data domains.

Training and Inference on Any-Order Autoregressive Models the Right Way
Andy Shih
Dorsa Sadigh Stanford University dorsa@cs. stanford. edu
Stefano Ermon Stanford University ermonOcs, stanford edu
Stanford University andyshih@cs stanford.
Abstract
Conditional inference arbitrary subsets of variables is problem in proba- bilistic inference with important applications such as masked language modeling and image inpainting. In recent years, the family of Any-Order Autoregressive Models (AO-ARMs) - closely related to popular models such as BERT and XI.- Net has shown breakthrough performance in arbitrary conditional tasks across sweeping range of domains. But, in spite of their success, in this paper we identify significant improvements to be made to previous formulations of AO-ARMs. First, we show that AO-ARMs suffer from redundancy in their probabilistic model, i.c., they define the same distribution in multiple different ways. We alleviate this re- dundancy by training on smaller set of univariate conditionals that still maintains support for efficient arbitrary conditional inference. Second, we upweight the training loss for univariate conditionals that are evaluated more frequently during inference. Our method leads improved performance with no compromises on tractability, giving stateoof-the-an likelihoods in arbitrary conditional modeling on text (Text8), image (CIFAR10, ImageNet32), and continuous tabular data domains.
1 Introduction
Generative modeling has seen tremendous progress in building highly expressive models [2, 31J, but relatively little effort has been put into supporting efficient probabilistic inference. Most existing models with deep neural architectures do not admit efficient inference on conditional queries of the form px=./ where and Such queries, however, have many important applications such as masked language modeling [42], image inpainting [43], and more. For example, multi-modal models learn joint distribution all data modalities, and at test time may only be presented with some subset (unknown in advance) of modalities [41, 20]. In robot shared-autonomy, an operator test time may choose to provide subset of the action inputs, leaving the model to fill in the remaining action dimensions [8, 23).
disjoint subsets of the variables of the joint distribution.
Evidently, expressive generative models that can support efficient conditional inference can bring progress to many applications. Towards this end, the family of Any-Order Autoregressive Models (AO-ARMs) [37, 42, 35, 14] has shown surprising success. AO-ARMs are built the following insights. An autoregressive model defines an ensemble of univariate conditionals p(z; (x<)) that leads to efficient inference certain conditional queries: ones whose variables form prefix of the ordering (Fig. la). However, in standard autoregressive models, all other conditionals are defined only implicitly through Bayes' rule, and hence do not admit efficient inference routines. To fix this, consider training another autoregressive model on the reverse ordering, which now defines a larger of univariate conditionals that can answer queries of either the prefix the suffix of the ordering (Fig. lb). AO-ARMs take this insight to the extreme by training on all possible orderings, enabling efficient inference on all possible conditionals. Remarkably, recent efforts in scaling up AO-ARMs
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

plg (4)u))
pla (x))
p(xy] (5.2))
p(tp
plx, (g
pla)
pl (90
1,233
12.3
p(x,)
p(aya)
pl
1234
plx
p(xa)
3.4
2,3,4
pig (5.
play 19.3)
(a)
(b)
Figure 1: (a) An autoregressive model defines N univariate conditionals p(z, (x<=)). allowing for efficient conditional inference when variable subsets form prefix of the ordering, e.g., p(21:3). (b) Learning an additional autoregressive model on the reverse ordering Tı, Z3, Z2, #1 enables efficient inference
queries: when variables form suffixes, e.g.. p( However, this leads to redundancy in the probabilistic model, as the joint distribution p((1144 1:44 is now defined in two different ways. This redundancy is exacerbated in AO-ARMs result of learning all possible orders.
have worked extremely well, with breakthrough performance in masked language modeling [42], image inpainting [14], and more [35].
However, spite of their already strong empirical results, we identify significant improvements that can be made to AO-ARMs. We begin by showing that previous formulations of AO-ARMs suffer from redundancy their probabilistic model. Consider again the example of training two autoregressive models on the lexicographical and reverse ordering (Fig. lb). This already unnecessarily defines the same joint distribution in two different ways, as there are two paths to the node p( 1111 corresponding to two distinct factorizations. By scaling up to all possible orderings, AO-ARMs inadvertently model the same distribution with a huge amount of redundancy. Such redundancy makes training inefficient and, much more importantly, leads to worse asymptotic performance. Due finite capacity of the model, attempting to learn too many univariate conditionals can lead to a poor fit on individual ones.
In this paper, we present MAC Mask-tuned Arbitrary Conditional Models which proposes two key insights for reducing redundancy and improving model performance. The first insight that arbitrary conditional inference can still be computed efficiently without relying conditionals. For example in Fig. 1b, omitting the edge p(01 [22 I3, I4) still maintains efficient inference for all the nodes in the graph, since there is another path for computing p(21-1). Based on this observation, we reduce redundancy of AO-ARMs by training on a smaller set of univariate conditionals that still supports tractable arbitrary conditional inference. The second insight is that some univariate conditionals are evaluated much more often than others during inference. Therefore, upweight the training loss for the more frequently occurring univariate conditionals, thereby aligning the training and inference objectives more closely.
all possible univariate
Combining these insights, MAC leads to strictly-improved formulation of AO-ARMs that, amaz- ingly, gives better arbitrary conditional and joint likelihoods with no compromises on tractability. We achieve tfateoof-the-ar likelihoods for arbitrary conditional modeling across multiple domains such as text (Text8), image (CIFAR1O, ImageNet32), and continuous tabular data. Finally, we conclude by demonstrating novel application of AO-ARMs to robot shared-autonomy.
2 Background
Problem definition Marginal and conditional inference are important queries for handling partial evidence on arbitrary subsets of variables, which is relevant for tasks such as masked language modeling or image inpainting. We let mask e be a subset of variables, and its cardinality |e| be the number of elements in the subset. A masked input I is a partial instantiation on the variable subset using the values taken on by input æ. Marginal inference refers to queries of the form p(xx.) J )( where Conditional inference queries y(ar,/a.) = ratio of two marginals, so we will refer to marginals and conditionals interchangeably. To evaluate model performance on marginal inference queries over a test dataset XLess, we assume a test mask distribution M over the N possible masks. We sample masks from M independently from data,
X and T.,.. denotes the union of the partial instantiations.
are similar and can be computed as a

---

and use the negative marginal log- oggliikelhood of the masked input ze as the loss.
LM =
Σ
EeeM log p(æc)
(I)
Computing arbitrary marginals and conditionals has been a long-standing challenge in probabilistic inference, with rich history of techniques [28, 26, 15, 5, 40] ranging from belief propagation, variational inference, to MCMC. Such queries are inherently difficult in high-dimensions: given access joint likelihoods p(æ), naïve computation of marginal likelihoods p(x:.) requires integration over the missing N - |el dimensions. As a result, most traditional approaches model the joint distribution using non-neural architectures that admit efficient inference but are less expressive [18].
Autoregressive Models Autoregressive models are an influential family of generative models that represents complex high-dimensionall distributions by modeling single dimension at time. They parameterize a joint distribution p(æ) over N dimensions by factorizing it into univariate conditionals Il Ih= (())(1)) [a(())) or) via chain rule (Fig. la), using an ordering of the N variables. We write (t) and (<< t) to denote the masks corresponding to the i-th element and the first t elements of the ordering, respectively.
The univariate conditionals can be learned via a weight-tied neural network, and composing them together leads to highly expressive architectures in practice, such as PixeICNN [38, 32] or Transform- ers [39]. Autoregressive models support efficient inference on the joint distribution, and on specific type of marginal query: ones where the mask forms prefix of the variable order, i.e., a [el) (Fig. la). We will call such a mask e and ordering compatible.
However, autoregressive models cannot support efficient arbitrary marginal inference in general. Next, we present AO-ARMs, which extend autoregressive models in a way that does support arbitrary marginal inference.
2.1 Any-Order Autoregressive Models (AO-ARMs)
AO-ARMs [37, 14] learn a single model that can generate the joint distribution autoregressively using any ordering of the N variables, by modeling II p(za(() [a(<)); o) for all orderings Even though there are N! different orderings, their chain-rule factorization is built from univariate conditionals, of which there are "only" N2 Therefore, an AO-ARM is model that defines the N2 distinct univariate conditionals p(r,|=0A)) where j e.e.
To answer marginal inference query p(æ,) with an AO-ARM, choose an order that is compatible with e, so that o(s [el) = e. Then, we simply evaluate each of the univariate conditionals p(xx)) I [ (=()); in the autoregressive factorization of Fe-
Training A0-ARMs Architecturally, an AO-ARM models all univariate conditionals via a weight- tied neural network, by using special (so-called "absorbing-state" [1, 14]) token for variables that are not present in the evidence set. enable parallel optimization, the AO-ARM architecture is designed so that given evidence 2.., the model can predict the univariate conditionals p(a: (..)) for all
i X \ e at once. Importantly, this parallelization works on non-causal architectures, opening the doors architectures similar in flexibility diffusion models [33, 34, 13].
In previous works [37, 35, 14], AO-ARMs are trained to maximize the joint likelihood of a datapoint under the expectation over the uniform distribution Ur of orders.
log)(#)) log E. Ep(ro(o) > Eg-il
Σ log (0)
(2)
This objective can be simplified by treating i as random variable with uniform distribution 24 over cardinalities to N. We let Meard-egge be the distribution over the tuple (a(t), a<<< t)) where Uh and ~ U4, giving the following loss function for an AO-ARM parameterized by 0.
L()) = - -Eg-ite Σ
log log (o(0))))))))( -N -E Muddot log:oolzizz) (p))
(3)
3

---

In practice, due to parallel optimization, we write the objective in the following equivalent form.
1
(()) -N Rewn-sttt
(1,
(4)
ce iE
Compared to other arbitrary conditional models, AO-ARMs consistently give state-of-the-art perfor- mance on benchmarks across range of domains [37, 35, 14]. Besides good empirical performance, they also serve conceptual unification between autoregressive models and diffusion models [14, 1], showing promise for both practical and theoretical advancements.
3 Improving Any-Order Autoregressive Models with MAC
Despite the success of AO-ARMs, in this section we identify two key deficiencies in previous formulations of AO-ARMs. We propose solutions to them by re-interpreting AO-ARMs from the perspective of recursive decomposition over marginals. Our resulting method MAC leads to consistent improvements and mate-of-the-ar likelihoods across range of domains in Sec. 4.
(A) The first deficiency is that of redundancy. While being order-agnostic is exactly what enables arbitrary conditionals, learning the same distribution with multiple orderings makes training inefficient. More importantly, since the model has limited capacity (especially with the weight-tied architecture), learning orderings redundantly hurts asymptotic performance.
(B) The second deficiency stems from the AO-ARM training objective Eq. 2. This objec- tive focuses only on the joint likelihood, leading to mismatch with the marginal likelihood evaluation objective in Eq. 1, even though computing marginals the key feature of AO-ARMs.
When viewing AO-ARMs as autoregressive models that generate using all orders, is not obvious how to fix the above two deficiencies. For example, if we try to omit certain orders to address (A), we may lose the ability to efficiently answer marginals, c.g., queries that are prefixes of the orders we removed. Similarly for (B), it is unclear how to choose distribution of orderings to best account for arbitrary conditional inference. To solve these issues, we re-interpret AO-ARMs from the perspective of computing arbitrary conditional probabilities via recursive decomposition over masks.
3.1 Re-interpreting A0-ARMs recursive decomposition a binary lattice
mask e is associated with the set of marginal inference queries of the form p(æ.). We will refer to this set of queries as the corresponding task for mask e. We represent these masks/tasks nodes in Fig. 1. Supporting efficient computation of each task, therefore, implies efficient arbitrary marginal inference. However, for masks with high cardinality, the corresponding task involves learning high-dimensional distribution.
Fortunately, these masks do not actually form standalone tasks, but are subproblems of one other. In particular, given non-empty mask and single variable j E e, we can write p(xe) ) as p(Z,)))(Z (eys). Then, delegate the computation of p(zeg) to the task associated with mask \.. and only left with estimating the univariate conditional (())))
From this perspective, we can view AO-ARMs as solving an ensemble of intertwined tasks, one for each of the N possible masks. We visualize this as binary lattice in Fig. 2a over N = variables, where we have 2 16 nodes representing the possible masks/tasks, with an edge connecting nodes e' and if e' = e \ j for some singleton variable j E e. To answer marginal query p(ze), we choose an element j E e and compute p(x:)) = p(æers)) )p(2)) (ae))), which corresponds to moving along the edge from node e to node in the lattice. We proceed recursively on p(æe),) until we reach p(xx)
Under this recursive decomposition framework, two key concepts are the decomposition protocol, which dictate how edges are chosen, and the mask/edge distributions, which specify the probability that different marginals and univariate conditionals are evaluated.
Definition A decomposition protocol w defines, for each non-empty mask e, a probability distri- bution we over the elements in In other words, Ejeo we we(i) 1 and we(j) 2 0.for + 0.

---

play (oq.a)) poss Ix? pis
12234
1M
3
2
2
(b)
Figure 2: (a) For N 4, marginal queries can take one of 2 16 possible masks, represented as square nodes. These masks are related to one another through univariate conditionals (edges), since p(are) = p( ()(r)) For example, given p(21) ), to compute p(211,)) we only need to learn the univariate conditional p(zz (m ). shown as the green edge. (b) We don't need to learn all univariate conditionals (edges) compute arbitrary marginals (nodes). suffices to learn one left-going edge for every node besides the 0 node. (c) Given a marginal query at node, we recursively decompose it by following red edges until we reach the @ node. Hence, some edges nodes will be traversed more often, so we should train on them more often.
decomposition protocol w specifies a process for answering marginal inference queries. Given marginal query at node we sample an element j ~ we and compute p(x.) = p(z,)) (ete))pp(()/eee)) We proceed recursively on p(æter)) until we reach p(ag) =
Definition A mask distribution is probability distribution over the 2 possible variable subsets for marginals. An edge distribution is a probability distribution over the N2N possible variable subsets for univariate conditionals.
Mask and edge distributions are relevant for specifying the test mask distribution (Eq. 1) of the problem domain, and the training edge distribution for the AO-ARM. When viewed under this recursive decomposition interpretation, previous formulations of AO-ARMs use a random decompo- sition protocol w-RND where w-RNND (j) 1/le]. Moreover, they training edge distribution Maae (Eq. 3), which unfortunately neglects to take into account the test mask distribution M.
3.2 MAC: Mask-tuned Arbitrary Conditional Model
We can now describe our method MAC. MAC improves upon previous AO-ARMs by choosing a custom decomposition protocol and deriving training edge distribution that aligns with the test mask distribution M..
(A) Reducing edge redundancy To answer arbitrary marginals, we don't 't need to learn to generate using all possible orderings. Rather, it suffices to decompose each task into any one of its children tasks. This constraint corresponds to learning single (red) incoming edge for each node besides the node. We show an example in Fig. 2b, where we connect each node to its leftward parent obtained by removing the greatest element of the node's mask. This is precisely the decomposition protocol that we use for MAC: w-MAC,
j is the greatest element in e.
(B) Matching edge distribution When presented with a marginal query the decomposition protocol w traverses path from node to node 9, evaluating the univariate conditionals along the way. Simulating w on masks ~ M leads to an induced edge distribution DM., that represents how frequently different univariate conditionals are evaluated during the inference process, which we depict in Fig. 2c using red edges with different thicknesses. We can see, for example, that under the protocol w-MAC, univariate conditionals such as p(x1 (arg) are evaluated highly frequently because it appears distribution on univariate conditionals to match this induced edge distribution D given by the downstream test mask distribution M and the decomposition protocol w-MAC.
subproblem to many marginal queries. Therefore, we should tune our training
5

---

Abstractly, let ~ w(ee denote the probability that w chooses a path of edges = {(i, |cl when decomposing e. In particular, é1 = 0. and i, is variable Et, and ct+1 = in U ex. Then we can write the marginal inference objective as follows, where C' is the constant E, Di,,-DDa,wle
i=l
le
Be [ogp((x))] = Bee B (ul())
log p(p, (me,))
(5)
= E.-.eeDDD.. Drs
log ((
(6)
Eq. 5 leads to the following algorithms for training MAC (Alg. 1). We sample a marginal query p( (e))) by independently sampling mask e ~~ M and an input ~ X. Then, simulate the path taken on the lattice by the decomposition protocol w by iteratively sampling an element j ~ We and updating to be \ j. Throughout this process, optimize our model on the edges (univariate conditionals)
traversed on this path. During testing (Alg. 2), we similarly step through the decomposition protocol and add together the univariate conditional likelihoods along the path.
Algorithm 1: Training MAC
Algorithm 2: Testing MAC
Input: Test mask distribution M, decomposition protocol w, training data distribution X, model 0
Input: Mask decomposition protocol test data æ, model
Output: Marginal log-likeli leglikelihhod log (xx.)
1 while training do
r
0-0 0
2
e ~ M,, x ~
while # 0 do
while * ( do
~ we 3
4
j ~ w/
r log po (zy)
0 с A
V Pologpo(z,jle (a())
e e
e e j
Return:
In practice, since the decomposition protocol w-MAC deterministically removes the greatest element, there is an efficient way to directly sample from MM, w-MAC in batch (Eq. 6), without having to simulate w-MAC as Alg. We describe these details the Appendix, Finally, to conclude this section, we discuss two important design choices and practical considerations for w and Das
Parallel training As noted in Section 2, the architecture AO-ARMs enables training on univariate conditionals of the form p(z, x.)) for i E X \e in parallel. Therefore, rather than dealing with training distribution over univariate conditionals aw (f, e), we instead work with one over masks (e). To close approximation, we will simply let (e) x Σ, DM. (i, e), visualized as shaded nodes in Fig.
Cardinality Reweighting (heuristic) We present one additional technique for tuning the training mask distribution, with ablations and further discussion later sections. To foster generalization of our weight-tied neural network, we reweigh the probability of a mask based on its cardinality by using a final training distribution of DCR MM. x (1 + (e) which we sample using SIR Inporegaccoeoeesaaaeeee
With parallelism and cardinality reweighting, the final objective for MAC is the following:
1
(M) == -EE MAC N lel
log lo,
teXVe
The decomposition protocol w-MAC recursively removes the greatest element of a mask, where comparison between two elements jı > j2 is determined by global canonical ordering (e.g. lexicographical ordering).
Experiments
We evaluate MAC on high-dimensional language and image domains, and on set of continuous tabular benchmarks. We focus on two metrics: joint likelihood and marginal likelihood of the test set. On both metrics, MAC shows state-of-the-an performance among arbitrary conditional models on the

---

majority of benchmarks. We conclude by demonstrating a novel application of arbitrary conditional models on a simulated robot shared- autonomy task, where we learn a conditional policy in 9-DoF action space through implicit behavioral cloning with MAC.
Note on computation Each run was done on a single NVIDIA A40. For the language and image experiments, we trained each model for approximately two weeks. Although this was not enough to match the total number of epochs trained by the baseline ARDM', we were still able to show state-of- the-art performance for arbitrary conditional models on 2 out of the 3 language/image benchmarks, and beat the baselines on all 3 benchmarks when compared under the same number of training epochs.
Character- modeling of the Text8 Table I:
Figure 3: Ablation on choices of training mask distribution. The best performing mask distribu- tion DM w-MMC uses the w-MAC decomposition protocol with cardinality reweighing.
dataset (without additional context), in bpd.
joint
marginal
Joint Models
1.40
from literature
D
R. WAND
Transformer [39]
1.35
1.35
Dat, MAC
Da, RND
Arb. Cond. Models
Mannh
Du. vailiatk
from literature
murgias
1.30
OA Transformer [14, 42]
1.64
D3PM [1]
1.47
1.25
ARDM [14] (14000 epochs)
1.43
1.20
experiments
ARDM (3000 epochs)
1.48
1.12
1.15
MAC (3000 epochs)
1.40
1.09
50
100
150
200
250
300
350
400
epoch
4.1 Language
We learn a character-level model on chunks of 250 characters using the Text8 dataset [24], which consists of 100M characters. We follow the setup from previous work in modeling the chunked text segments without any additional context, and use the same 12-layer Transformer architecture as ARDM and D3PM [14, 1]. We report the joint and marginal likelihoods in terms of bpd (lower is better), where the marginal bpd is computed w.r.t. to the test mask distribution Meant-mask( (e) x 1/ ((el N)
as used in [35] In Tab. we see that MAC outperforms all existing arbitrary conditional models in joint likelihood and comes close standard Transformer trained only model in one variable order, even when trained only on a fraction of the total number of epochs by baselines in literature. Similarly, MAC also gets better marginal bpd compared ARDM.
Ablations In Fig. 3 we present ablation studies on different choices of the training mask distribution. The baseline method (ARDM) trains the Meard-mask mask distribution (grey). Using only insight (A) from Sec. 3 with a deterministic decomposition protocol but no edge-weighting gives slight improvements M, an-allrcAA (light purple). Using only insight (B) from Sec. 3 with edge-weighting but random decomposition protocol also leads to slight improvements D DM. WRRND (light blue). Using both insights (A) and (B) gives the MMMAC protocol which leads to more improvements DM, wMC (light orange). Lastly, the cardinality reweighing heuristic improves both methods even more (dark bluc/orange). The best performance is given by DCR use for MAC in all other experiments.
M.W-MAC in dark orange, which is the setting we
4.2 Images
We evaluate MAC on CIFAR1O [19] and ImageNet32 [4, 6]. both of which consist of images of dimension 3072. Again, we follow all experimental settings from previous work [1, 14, 17], using U- Net with 32 ResNet Blocks interleaved with attention layers, and use rotation/flip data augmentation. In Tab. 2&3 we see that MAC outperforms ARDM on joint bpd and marginal bpd, evaluated on the
Previous works were unable to share checkpoints of baseline methods.
Maudowaat effectively samples a cardinality and then samples uniformly among masks with cardinality
7

---

Meard-mask mask distribution. Surprisingly, MAC even gives the best joint bpd on ImageNet32 when compared to joint models that do not support arbitrary conditionals.
Table 3: Pixel modeling of the CIFARIO dataset
(with rotation/flip data augmentation), in bpd.
Table 2:
Pixel modeling of the ImageNet32 dataset (with no data augmentation), in bpd.
joint
marginal
Joint Models
joint marginal
from literature
Joint Models
PixeICNN++ [32]
2.88
from literature
Sparse Transformer [3, 16]
2.56
Image Transformer [27]
3.77 3.72
VDM [17]
2.49
VDM [17]
Arb. Cond. Models
Arb. Cond. Models
from literature
our experiments
D3PM [1]
3.44
ARDM (16 epochs)
3.60
2.10
ARDM [14] (3000 epochs)
2.69
MAC (16 epochs)
3.58
2.08
our experiments
ARDM (1200 epochs)
2.86
1.84
MAC (1200 epochs)
2.81
1.81
4.3 Continuous Tabular Data
Next, we consider tabular domain with variables taking on continuous values [35]. The main challenge with continuous values for AO-ARMs is the parameterization of the univariate conditionals, since I-D Gaussians or even mixtures Gaussians have limited expressiveness. To this end, ACE [35] proposes to parameterize the 1-D conditionals EBMs. We build MAC top of ACE, keeping all hyperparameters and experimental setup the same, modifying only the training mask distribution. In Tab. 4&5, settings, and outperforms other arbitrary conditional models such as SPNs [30, 25] and ACFlow [22].
that MAC gives better marginal and conditional likelihoods than ACE most
Table 4: Marginal log- likelikoi 5 continuous tabular benchmarks (higher is better). The mask cardinality settings kept consistent with the ones reported [35].
power
gas
bepmaas
miniboone
bsds
Mask cardinality
6
5
3
5
10
3
10
SPHlow [25]
-4.42
ACFlow [22]
-0.57
1.34
0.42
0.78
3.01
10.13
-4.03
-6.19
-11.58
2.76
-531
-10.36
5.06
9.26
19.60
ACE [35]
-0.56
1.42
0.58
1.31
4.31
12.20
-4.00
-5.1
-10.72
2.13
3.80
-7.94
5.10
9.37
20.31
MAC
40.55
1.43
0.61
1.59
4.98
13,02
-4000
-5.90
10.69
2.12
3.76
-7.76
5.10
9.37
20.33
Table 5: Conditional log-likelihood on continuous tabular benchmarks.
Table 6: Shared autonomy on FrankaKitchen with BC operator policy. Full autonomy baseline (IBC) [10]: 2.15 + 0.06
power
gas
hepmass
miniboone
bsds
SPFlow [25]
-1.03
4.30
12.78
-18.34
-24.15
Conditional Policy
Reward
ACFlow [22]
0.56
8.09
-8.20
-0.97
81.83
ACE [35]
0.63
9.64
-3.86
0.31
86.70
BC
1.81 + 0.08
MAC
0.65
9.77
-3.05
0.07
86.05
MAC
2.00 + 0.05
4.4
Shared Autonomy on FrankaKitchen
Finally, we demonstrate an application of arbitrary conditional models on the task of shared autonomy in robotic manipulation. In shared autonomy settings, operator aims to control a policy in a complicated action space, but allowed to delegate partial control to Al model for easier manipulation. We focus the simulated FrankaKitchen environment [12, 11], where the goal is to control a 9-DoF robotic arm move objects around kitchen.
We train a base policy with behavioral cloning (BC) [29] and an arbitrary conditional policy with MAC to solve the Kitchen aecchen-miieee task. To simulate shared autonomy, evaluate

---

a hybrid policy where at each state, the operator policy (BC) outputs 4 out of the 9 dimensions of the action space, and conditional policy fills in the remaining 5 action dimensions. Tab. 6 we see that conditioning using MAC improves shared autonomy as compared to conditioning using an independent BC policy, and comes close to the full autonomy performance [10].
5
Discussion
Sharpen mask distribution Intuitively, we should choose a protocol w such that DM. has low entropy so that our model has a greater chance of revisiting similar masks. To investi- gate the differences in choices of decomposition protocol, in Tab. 7 we present Monte- Carlo sim- ulations of the induced mask distributions with respect to the test mask distribution Mank--mask. As expected, using w-MAC leads to lower em- pirical entropy than using w-RND.
Table 7: Monte- Carto simulation of induced mask distributions for N = 12, taking 1e6 samples. For brevity we only display a few selected rows.
Mask
M
DM w-ND DM DC
M:MMC
018111111110
1261
222
1827
100000000000
7037
14803
76381
41441
100000000101
407
521
0
100000001000
1288
1975
536
430
111111111110
6956
433
6935
21094
Entropy
7.21
5.74
4.44
5.51
Mask generalization (and why we do Cardi- nality Reweighting) The success of AO-ARMs depends on generalization to unseen masks, since there are exponentially many masks. Training on many masks improves generalization, and can be viewed as a form of data augmentation that allows AO-ARMs to sometimes even outperform standard autoregressive models on joint likelihoods.
In light of this, decomposition protocol w that skews the mask distribution may have the unfortunate side-effect of hurting generalization. In particular, the mask distribution Das.w-MAC puts much more weight low-cardinality masks. This is problematic since training on cardinality-1 masks is not conducive to generalization: there are only N possible cardinality-1 masks, whereas higher cardinalities have exponentially more masks. As such. we proposed DCR MM-C (Tab. 7) with the hypothesis that focusing less on low-cardinality masks will lead to better generalization. Our ablation study supports this hypothesis, as cardinality reweighting indeed improves performance (Fig. 3).
Parallel training To account for parallel training, we proposed to train on the mask distribution (e)) x Eiex,e DM. (i, e). However, each pass of (e) trains on all of {p(r;|(x.) : E X \ e), which overtrains neighboring edges and leads to a minor distributional mismatch. Aligning this training distribution more closely may lead to improvements, but requires better techniques for parallelization and sampling of masks. We leave this for future investigation.
5.1 Related Work
The idea of training AO-ARMs autoregressive models on all possible orderings was first introduced in NADE [37, 36), and has been seen more recently in models such as ARDM [14] for image / text / audio and ACE [35] for continuous domains. AO-ARMs are also closely related to non-autoregressive language models such as BERT [7] and XLNet [42]. Our method, MAC, proposes improved training and inference techniques for AO-ARMs.
Other techniques have also been proposed for non-autoregressive generation. For example, IN- TRUS [9] models sequences as consecutive insertion operations, as opposed to AO-ARMs which model sequences as consecutive unmasking operations. This also enables arbitrary order of genera- tion, but prevents INTRUS from computing marginal likelihoods on subsets of variables. Another work [44] casts the problem of non-autorrgressivv generation as a GFlowNet, taking on a reinforce- ment learning framework with states, actions, and rewards. However, this approach currently does not scale as well as AO-ARMs.
Lastly, some works have studied the problem of learning a single generation order [21]. Although learning orderings is not applicable to vanilla AO-ARMs (since they train on all possible orders uniformly), it is applicable to MAC, which uses decomposition protocols that can be learned. For example, we can learn different canonical orderings of the variable indices in the binary lattice. Though we did not explore this direction thoroughly, we noticed that strided canonical ordering (0, 32, 64, 1, 33, ...) often did better than the standard lexicographical ordering for images, suggesting promise for learning even better canonical orderings, or learning better protocols in general.

---

5.2 Limitations
In practice, likelihood evaluations of standard AO-ARMs are done by sampling a finite number of orders (typically just a single order), leading to unbiased but approximate estimates of joint likelihoods. Moreover, evaluating marginal p(ze) and conditional likelihoods p( (x)) is trickier and has additional tradeoffs. Since AO-ARMs evaluate marginal and conditional likelihoods using only compatible orderings, they are implicitly discarding away all the models in the N! ensemble with incompatible orderings. This can give biased estimates for marginal and conditional likelihoods.
For conditional likelihoods specifically, there two potential methods for evaluation. First, evaluate log p( (z.) as log p(x,, log p(x.) and estimate two marginal likelihoods. However, since the marginal likelihoods are approximate, this can lead invalid conditional probability estimates where p( (xe) > (for discrete domains). Second, we can decompose Su directly and evaluate log p( (x))) as E/u,, log p( [2,,,, ze), where the ordering variables within æ., is sampled randomly. This corresponds to evaluating random path that goes from node to node U in the binary lattice. Though this leads to valid probability estimates, the estimates may be more biased because paths that go to node t U without passing through node are not considered.
MAC shares similar limitations standard AO-ARMs, where marginal likelihoods estimates may be biased. Evaluating conditional likelihoods as two marginal likelihoods can lead to invalid (greater than 1) estimates, and evaluating them by tracing paths through the lattice from node node U may be more biased.
Nonetheless, both marginal and conditional likelihood estimates (using Method 2) for MAC and standard AO-ARMs are valid in the sense that Ez, p p( = I and En, p( (z.)) = for any u and Training MAC to optimize for marginal likelihoods (Tab. 1,2,3) or for both marginal/conditionall likelihoods (Tab, 4,5) still gives good empirical performance despite these limitations.
Conclusion
We present MAC, an improved method of training and inference on AO-ARMs. MAC trains on carefully designed distribution of univariate conditionals that (A) reduces modeling redundancy of AO-ARMs and (B) aligns the training objective of AO-ARMs with arbitrary conditional queries. Our method leads to better joint and arbitrary conditional likelihoods with no sacrifices on tractability. We show state-of-the-as results in arbitrary conditional modeling on text, image, continuous data domains, and present a novel application of arbitrary conditional models to robot shared-autonomy
7 Acknowledgments
We thank Rui Shu, Jiaming Song, Jesse Mu, and anonymous reviewers for their constructive feedback. This research was supported by NSF(#1651565). AFOSR (FA9550-19-110024)) ARO (W91 INF-21- 1-0125), ONR, DOE, CZ Biohub, and Sloan Fellowship.
10

---

References
[1] Jacob Austin, Daniel D. Johnson, Jonathan Ho, Daniel Tarlow, and Rianne van den Berg. Struc- tured denoising diffusion models in discrete state-spaces. In Advances in Neural Information Processing Systems 34, 2021.
[2] Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert- Hersert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. In Advances in Neural Information Processing Systems 33, 2020.
[3] Rewon Child, Scott Gray, Alec Radford, and llya Sutskever. Generating long sequences with sparse transformers. CoRR, abs/1904. 10509, 2019.
[4] Patryk Chrabaszez, Iyy Loshchilov, and Frank Hutter. A downsampled variant of imagenet as an alternative to the CIFAR datasets. CoRR, abs/1707 0700000819, 2017.
[5] Adnan Darwiche. differential approach inference in bayesian networks. J. ACM, 50(3):280- 305, 2003.
[6] Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In IEEE Computer Society Conference on Computer Vision and Pattern Recognition CVPR, 2009.
[7] Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: pre-training of deep bidirectional transformers for language understanding. Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, NAACL-HLT 2019, 2019.
[8] Anca Dragan and Siddhartha S. Srinivasa. A beiicc- formalism for shared control. The International Journal of Robotics Research, 2013.
[9] Dmitrii Emelianenko, Elena Voita, and Pavel Serdyukov. Sequence modeling with unconstrained generation order. Advances in Neural Information Processing Systems, 32, 2019.
[10] Pete Florence, Corey Lynch, Andy Zeng, Oscar A. Ramirez, Ayzaan Wahid, Laura Downs, Adrian Wong, Johnny Lee, Igor Mordatch, and Jonathan Tompson. Implicit behavioral cloning, In Conference on Robot Learning, 2021.
[11] Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4rl: Datasets for deep data-driven reinforcement learning, 2020.
[12] Abhishek Gupta, Vikash Kumar, Corey Lynch, Sergey Levine, and Karol Hausman. Relay policy learning: Solving long-horizon tasks via imitation and reinforcement learning. In 3rd Annual Conference on Robot Learning, 2019.
[13] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In Advances in Neural Information Processing Systems 33, 2020.
[14] Emiel Hoogeboom, Alexey A. Gritsenko, Jasmijn Bastings, Ben Poole, Rianne van den Berg, and Tim Salimans. Autoregressive diffusion models. 10th International Conference on Learning Representations, 2022.
[15] Michael I. Jordan, Zoubin Ghahramani, Tommi S. Jaakkola, and Lawrence K. Saul. An introduction to variational methods for graphical models. Machine Learning. 1999.
[16] Heewoo Jun, Rewon Child, Mark Chen, John Schulman, Aditya Ramesh, Alec Radford, and Ilya Sutskever. Distribution augmentation for generative modeling. In Proceedings of the 37th International Conference on Machine Learning, ICML, 2020.
[17] Diederik P Kingma, Tim Salimans, Ben Poole, and Jonathan Ho. Variational diffusion models. In Advances in Neural Information Processing Systems 34, 2021.

---

[18] Daphne Koller and Nir Friedman. Probabilistic graphical models: principles and techniques. MIT press, 2009,
[19] Alex Krizhevsky. Learning multiple layers of features from tiny images. 2009.
[20] Alexander Ku, Peter Anderson, Roma Patel, Eugene Ie, and Jason Baldridge. Room-across- room: Multilingual usioo-andd-anguage navigation with dense spatiotemporal grounding. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing, EMNLP, 2020.
[21] Xuanlin Li, Brandon Trabucco, Dong Huk Park, Michael Luo, Sheng Shen, Trevor Darrell, and Yang Gao. Discovering non-monotonic autoregressive orderings with variational inference. In 9th International Conference on Learning Representations, ICLR 2021, 2021.
[22] Yang Li, Shoaib Akbar, and Junier Oliva. Acflow: Flow models for arbitrary conditional likelihoods. In Proceedings of the 37th International Conference on Machine Learning ICML, 2020.
[23] Dylan P. Losey, Krishnan Srinivasan, Ajay Mandlekar, Animesh Garg, and Dorsa Sadigh. Controlling assistive robots with learned latent actions. In IEEE International Conference Robotics and Automation, ICRA, 2020.
[24] Matt Mahoney, Large text compression benchmark, 2011.
[25] Alejandro Molina, Antonio Vergari, Karl Stelzner, Robert Peharz, Pranav Subramani, Nicola Di Mauro, Pascal Poupart, and Kristian Kersting, Spflow: An easy and extensible library for deep probabilistic learning using sum-product networks. CoRR, abs/1901 0037044, 2019.
[26] Kevin P. Murphy, Yair Weiss, and Michael I. Jordan. Loopy belief propagation for approximate inference: An empirical study. In Proceedings of the Fifteenth Conference on Uncertainty in Artificial Intelligence UAI, 1999.
[27] Niki Parmar, Ashish Vaswani, Jakob Uszkoreit, Lukasz Kaiser, Noam Shazeer, Alexander Ku, and Dustin Tran. Image transformer. In Proceedings of the 35th International Conference on Machine Learning, ICML, 2018.
[28] Judea Pearl. Reverend bayes on inference engines: A distributed hierarchical approach. In Proceedings of the National Conference Artificial Intelligence, 1982.
[29] Dean Pomerleau. ALVINN: an autonomous land vehicle in a neural network. In Advances in Neural Information Processing Systems 1, NIPS, 1988.
[30] Hoifung Poon and Pedro M. Domingos. Sum-product networks: A new deep architecture. In Proceedings of the Twenty-Seventh Conference on Uncertainty in Artificial Intelligncce UAI, 2011.
[31] Aditya Ramesh, Mikhail Pavlov, Gabriel Goh, Scott Gray, Chelsea Voss, Alec Radford, Mark Chen, and Ilya Sutskever. Zero-shot text-to-image generation. In Proceedings of the 38th International Conference on Machine Learning ICML, 2021.
[32] Tim Salimans, Andrej Karpathy, Xi Chen, and Diederik P. Kingma. Pixelcnn++: Improving the pixelenn with discretized logistic mixture likelihood and other modifications. 5th International Conference Learning Representations, ICLR, 2017.
[33] Jascha Sohl-Dickstein, Eric A. Weiss, Niru Maheswaranathan. and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In Proceedings of the 32nd International Conference on Machine Learning. ICML. 2015.
[34] Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. In Advances in Neural Information Processing Systems 32, 2019.
[35] Ryan R. Strauss and Junier B. Oliva. Arbitrary conditional distributions with energy. In Advances in Neural Information Processing Systems 34, 2021.
12

---

[36] Benigno Uria, Marc-Alexandrr Côté, Karol Gregor, lain Murray, and Hugo Larochelle. Neural autoregressive distribution estimation. The Journal of Machine Learning Research, 17(1):7184- 7220, 2016.
[37] Benigno Uria, lain Murray, and Hugo Larochelle. A deep and tractable density estimator, In Proceedings of the 31th International Conference on Machine Learning. 2014.
[38] Aäron van den Oord, Nal Kalchbrenner, Lasse Espeholt, Koray Kavukcuoglu, Oriol Vinyals, and Alex Graves. Conditional image generation with pixelenn decoders. In Advances in Neural Information Processing Systems 29, 2016.
[39] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in Neural Informa- tion Processing Systems 30, 2017.
[40] Martin J. Wainwright and Michael I. Jordan. Graphical models, exponential families, and variational inference. Foundations and Trends in Machine Learning. 2008
[41] Mike Wu and Noah D. Goodman, Multimodal generative models for scalable seakly-supervised learning. In Advances in Neural Information Processing Systems 31, 2018.
[42] Zhilin Yang. Zihang Dai, Yiming Yang, Jaime G. Carbonell, Ruslan Salakhutdinov, and Quoc V. Le. XLNet: Generalized autoregressive pretraining for language understanding. In Advances in Neural Information Processing Systems 32, 2019.
[43] Raymond A. Yeh, Chen Chen, Teck- Yian Lim, Alexander Schwing, Mark Hasegawa-Johnson and Minh N. Do. Semantic image inpainting with deep generative models. In IEEE Conference on Computer Vision and Pattern Recognition, CVPR, 2017.
[44] Dinghuai Zhang, Nikolay Malkin, Zhen Liu, Alexandra Volokhova, Aaron C. Courville, and Yoshua Bengio. Generative flow networks for discrete probabilistic modeling. In International Conference on Machine Learning, ICML 2022, 2022.
13

---

Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's
contributions and scope? [Yes]
(b) Did you descrihe the limitations of your work? [Yes]
(c) Did you discuss any potential negative societal impacts of your work? [No]
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
2. If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [N/A]
(b) Did you include complete proofs of all theoretical results? [N/A]
3. If you experiments..
(a) Did you include the code, data, and instructions needed to reproduce the main ex- perimental results (either in the supplemental material or as URL)? [Yes] In the supplemental material.
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [Yes]
(c) Did you report error bars (e.g., with respect to the random seed after running ex- periments multiple times)? [No] Due computational limitations we ran on one seed.
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)? [Yes]
4. If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) your work uses existing assets, did you eite the creators? [Yes]
(b) Did you mention the license of the assets? [Yes] All datasets are publicly available.
(c) Did you include any new assets either the supplemental material or as URL? [No]
(d) Did you discuss whether and how consent was obtained from people whose data you're 'r
using/curating? [N/A]
(e) Did you discuss whether the data you are using/curating contains personally identifiable information offensive content? [N/A]
5. If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, applicable? [N/A]
(b) Did you describe any potential participant risks, with links Institutional Review Board (IRB) approvals, if applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation? [N/A]
14


Reviews:
[{"review_id": "uINKE7DIEJA", "reviewer": "Reviewer_fwZi", "summary": "The paper improves Any-order ARO models by  making them more efficient reducing redundancy in computations, while maintaining tractability and by better aligning the training and inference objectives.\n\nThe key idea is to reformulate the objectives in AO-ARO such that redundancies in computing the univariate conditionals can be fully exploited. Further, by tuning the training based on the distribution that represents the importance of univariate conditionals, the training is better suited to compute marginals during inference.\n\nExperiments are performed on several tasks and comparison studies show the improvements the proposed approach offers compared to state of the art", "questions": "None", "limitations": "They are not explicitly addressed. It would be nice to have a discussion on these in the paper.", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "1. The paper is written very well clearly defining the problem, the background work and its own novel contribution.\n2. The work is meaningful and significant considering the generality of computing marginal probabilities (by integrating from high-dimensional data)\n3. The experiments seem to cover many areas, standard benchmarks from computer vision, language models, robotics, etc. The variety of experiments clearly shows the strengths of this work.\n\nWeakness\n\nWhile this may not be a weakness, maybe limitations of the proposed work can be mentioned or discussed.\n\nOverall, I thought the paper proposed a significant problem, clearly outlined a good solution and conducted an excellent empirical study.\n", "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": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "beSdMKgjpHY", "reviewer": "Reviewer_HNCy", "summary": "The authors present the Mask-tuned Arbitrary Conditional Models - a method for training Any-order Autoregressive Models.\nThey argue that the paper presents two insights hurting conventional AO-ARM training which they address with their method. \n\nThe authors investigate their method with an ablation study and present it on a variety of experiments. \n\n", "questions": "I'm not sure what your \"sharpen mask distribution\" section is concluding. Is it just that your selected mask distribution has lower entropy (by design) than uniform? You're not showing that low entropy mask distributions actually help, except, possibly, in your earlier experiments. \n\nI think there is a small typo in the caption of Figure 1a: \"An autoregressive (model?) defines ...\"\n\n", "limitations": "A small discussion on the mask distributional mismatch is included. Generally, limitations, or the lack thereof, are not discussed. ", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths: \n\nThe explanations in the paper are well done and overall I found the paper clear and well-written. \n\nAblation study well isolates the components of your method and the contribution of each of them.\n\nExperiments are clear and well presented. A good variety of experiments were presented.\n\nThe authors state that it is not arbitrary joint distribution decomposition ordering that we really care about but, instead, arbitrary conditioning. They show that an ordering over every marginal can be specified, or a distribution over orderings, to address redundancy in training.\n\nThe authors show that tuning the mask distribution for training can lead to better results - an important idea that can be investigated further. Although the idea is straight-forward, I view this as a strength of the method as it is a simple, effective modification to the standard training techniques that yields better results with, as the authors put it, no compromises.\n\nWeaknesses: \nA lack of error measurements in most of the experiments makes it hard to determine if your results are significant. \n\nThe idea that we want to weight masks more often used (which are the lower carnality ones) and the cardinality reweighting seem contradictory so I'm not sure how much I believe your intuitive explanations of each idea. The ablation study shows the effect of using the insights in the order of B->A but I wonder how well A only would do. Did you try to use a uniform distribution over the masks remaining after insight A (pictured in figure 2c)?\n\nIf I understand correctly, your algorithm box doesn't actually reflect how you train the model (in parallel) using the objective in eq 7. There is a lot of machinery presented (although interesting) to arrive at the conclusion that the MAC objective (eq 7) is the standard AO-ARM objective (eq 4) with a different mask sampling distribution.\n\nA subjective point: I found the notation confusing in some places. You use several different types of notation to refer to the same objects - you refer to the conditional distribution of interest as p(x_sigma(t)| x_sigma(<t)), p(x_i_t|x_e_t), p(x_i|x_e), p(x_j| x_e/j). I was particularly confused by referring to edges as (i, e) or (i_t, e_t). What is i in this case? I assume from context it's a variable not contained in the mask e due to the summation on line 207 but I don't think it was defined. \n\n", "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": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "QpTB0GDiuD3", "reviewer": "Reviewer_9KgX", "summary": "A method for training nonautoregressive generative models of high-dimensional (e.g., sequence) data is proposed. The main algorithm requires the choice of a decomposition protocol, which is a choice of generation order defined by a deterministic Markovian reverse generation process. The training objective is equivalent to a weighted log-likelihood of each generation step in the chosen order leading to a randomly masked training sample. Strong results are shown on a variety of domains.", "questions": "- Would the results change if \"largest\" is replaced by \"smallest\" (or largest w.r.t. some random but fixed permutation of the dimensions) in the definition of the decomposition protocol (L179)? The definition relies upon a choice of natural ordering in the data space, so it may be important.\n   - In particular, how could the algorithm be combined with the learning of a (possibly probabilistic) decomposition protocol?\n- Relation to past work.\n   - A few papers that have addressed the question of non-autoregressive modeling have not been mentioned. For instance, [A1] is a relevant citation, while [A3] suggests an alternative approach that uses a continuous relaxation to learn generation order.\n   - (Also related to the first question:) There are several other algorithms that feature a learned generation order, albeit in slightly different settings. In [A2], non-uniform generation order emerges implicitly from a greedy optimization procedure. In [A4], which uses a similar framing of chains on a lattice, the learned \"backward policy\" is a stochastic, data-dependent decomposition protocol. What is the relationship of MAC with such approaches?\n\n[A1] B.Uria et al., Neural autoregressive distribution estimation (JMLR, 2016). arXiv:1605.02226\n\n[A2] D.Emelianenko et al., Sequence modeling with unconstrained generation order (NeurIPS 2019). arXiv:1911.00176\n\n[A3] X.Li et al., Discovering non-monotonic autoregressive orderings with variational inference (ICLR 2021). arXiv:2110.15797\n\n[A4] D.Zhang et al., Generative flow networks for discrete probabilistic modeling (ICML 2022). arXiv:2202.01361", "limitations": "Please see the second point of Weaknesses above. ", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 2, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n- The proposed algorithm is simple and intuitive and the derivation is sound.\n- Strong results on a diverse set of problems, including language, vision, tabular data, and robot motion planning. Evaluations and comparisons are solid.\n\nWeaknesses: (I am willing to raise the score if the authors discuss the limitations and relationship with past work in their response and commit to adding them to the paper.) \n- The definition of the decomposition protocol used is hidden in the text (L179) and not illustrated on real data. This makes it hard for a reader to check their understanding and gain intuition. \n- The use of a canonical ordering of data dimensions seems like a limitation of the approach, but it is not discussed. For example, even in the image domain, it is not clear that the standard ordering of pixels in the \"reading\" order is optimal for modeling. In other domains, such as tabular data and other settings where incremental generation in arbitrary order is possible (e.g., generic fixed- but high-dimensional data or data of variable size, such as graphs, point clouds, etc.), there is no natural ordering. Please see the questions below.\n- The discussion of related work could be improved. Please see the second question below.", "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": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "G-Tq2y9WKCV", "reviewer": "Reviewer_2pLm", "summary": "The paper argued that one current limitation of Any-Order Autoregressive Models (AO-ARMs) is that they suffer from redundancy in their probabilistic model, i.e., they define the same distribution in multiple different ways. Especially, to my understanding, the authors mean that a sequence of length $n$ can be defined by $n!$ orders, which causes the model to under-fitting any-order data distribution as observed in XLNet or ARDMs. A particular order (i.e., $w$-MAC, Mask-tuned Arbitrary Conditional Model) is introduced to remove such redundancy by sampling the order from a non-uniform decomposition protocol $w$. During the inference (which only contains the likelihood estimation task in this paper), the order is also sampled from $w$-MAC, rather than uniform distribution, and thus achieves better likelihood estimation (somewhere between left-to-right autoregressive model and real AO-ARMs).", "questions": "1. Clearly, a biased order distribution can produce better perplexity, as it is trained more on a specific order. And therefore, in principle, it should not perform better than real AO-ARMs, but this is not reflected in the experimental results of Table 1, which caused my confusion. My question is on the evaluation of marginal likelihood. Let's again take Figure 2 (b) as an example, I'm wondering how proposed method can compute the marginal likelihood of $p(x_1, x_2, x_3, x_4 | x_2, x_4)$ if there is no edges between (2, 4) and (1, 2, 3, 4)?", "limitations": "The work does not have an obvious potential negative societal impact.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\n1. The main strength of the paper is that it finds a (maybe principled) way to tweak the order distribution in AO-ARMs to get a better likelihood number (which also might not be true, as I argued in the Weaknesses). This could potentially improve the quality of tasks that require sampling from an AO-ARMs, and in the meantime, can produce a valid probability given any observed and unobserved variables.\n\n2. A heuristic cardinality reweighting trick is introduced in Section 3.2 and explained in Section 5. This may inspire future research in this direction.\n\n3. Despite the high complexity, the model can still be trained in parallel.\n\nWeaknesses:\n\nI mainly have two concerns.\n\n1. Using Figure 2 (b) as an example, in the proposed method, the following equation might hold $p(x_2, x_3, x_4) > p(x_2, x_4)$, which means $\\mathbb{E}_{\\sigma} p(x_3 | x_2, x_4) > 1$. Based on this, I think the mathematical foundation in the proposed model is invalid.\n\n2. In the evaluation, all the previous work samples a random order $\\sigma$ from the uniform distribution because of the following assumption:\n   1.  \"In other words, it does not matter which step $t + k$ the model predicts, in expectation, these all have the same associated likelihood.\" from [1]\n   2.  \"We note that a model which perfectly captures the true densities would give consistent likelihoods for all possible orderings (thus evaluating only one ordering would suffice).\" AND Algorithm 1 from [2]\n\nHowever, in this work, $\\sigma$ is no longer from a uniform distribution. Therefore, in order to get a precise likelihood estimation, the paper might need to get more order samples to reduce the variance in likelihood estimation. Still, no such information or discussion is provided in the paper.\n\n[1] Autoregressive Diffusion Models\n[2] Arbitrary Conditional Distributions with Energy\n", "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": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept