Training and Inference on Any-Order Autoregressive Models the Right Way
Anonymous Author(s)
AffiliationAddress email
Abstract
Conditional inference on arbitrary subsets of variables is a core problem in probabilistic inference with important applications such as masked language modelingand image inpainting.
In recent years, the family of Any-Order Autoregressive
Models (AO-ARMs) – which includes popular models such as XLNet – has shownbreakthrough performance in arbitrary conditional tasks across a sweeping rangeof domains.
But, in spite of their success, in this paper we identify significantimprovements to be made to previous formulations of AO-ARMs.
First, we showthat AO-ARMs suffer from redundancy in their probabilistic model, i.e., they definethe same distribution in multiple different ways.
We alleviate this redundancy bytraining on a smaller set of univariate conditionals that still maintains support forefficient arbitrary conditional inference.
Second, we upweight the training loss forunivariate conditionals that are evaluated more frequently during inference.
Ourmethod leads to improved performance with no compromises on tractability, givingstate-of-the-art likelihoods in arbitrary conditional modeling on text (Text8), image(CIFAR10, ImageNet32), and continuous tabular data domains.
Introduction
Generative modeling has seen tremendous progress in building highly expressive models [2, 28], butrelatively little effort has been put into supporting efficient probabilistic inference.
Most existingmodels with deep neural architectures do not admit efficient inference on conditional queries ofthe form p ( x u | x v ) , where u and v are disjoint subsets of the variables of the joint distribution.
Such queries, however, have many important applications such as masked language modeling [38],image inpainting [39] and more.
For example, multi-modal models learn a joint distribution overall data modalities and at test time may only be presented with some subset (unknown in advance)of modalities [37, 18].
In robot shared-autonomy, an operator at test time may choose to provide asubset of the action inputs, leaving the model to fill in the remaining action dimensions [7, 20].
Evidently, expressive generative models that can support efficient conditional inference can bringprogress to many applications.
Towards this end, the family of Any-Order Autoregressive Models(AO-ARMs)
[33, 38, 32, 12] has shown surprising success.
AO-ARMs are built on the followinginsights.
An autoregressive model defines an ensemble of univariate conditionals p ( x t | x <t ) thatlead to efficient inference on certain conditional queries: ones whose variables form a prefix of theordering (Fig. 1a).
However, in standard autoregressive models, all other conditionals are definedonly 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
1, 1,2, 1,2,3,  ∅ p ( x 1 ) p ( x 2 | x 1 ) p ( x 3 | x 1 , x 2 ) p ( x 4 | x 1 , x 2 , x 3 ) p ( x 1 ) p ( x 1:2 ) p ( x 1:3 ) p ( x 1:4 ) (a) 1, 1,2, 1,2,3,  ∅ p ( x 1 ) p ( x 2 | x 1 ) p ( x 3 | x 1 , x 2 ) p ( x 4 | x 1 , x 2 , x 3 ) 3.
2,3,4 4 p ( x 4 ) p ( x 3 | x 4 ) p ( x 2 | x 3 , x 4 ) p ( x 1 | x 2 , x 3 , x 4 ) (b) set of univariate conditionals that can answer queries of either the prefix or the suffix of the ordering(Fig. 1b).
AO-ARMs take this insight to the extreme by training on all possible orderings, enablingefficient inference on all possible conditionals.
Remarkably, recent efforts in scaling up AO-ARMshave worked extremely well, with breakthrough performance in masked language modeling [38],image inpainting [12], and more [32].
However, in spite of their already strong empirical results, we identify significant improvements thatcan be made to AO-ARMs.
We begin by showing that previous formulations of AO-ARMs suffer fromredundancy in their probabilistic model.
Consider again the example of training two autoregressivemodels on the lexicographical and reverse ordering (Fig. 1b).
This already unnecessarily defines thesame joint distribution in two different ways, as there are two paths to the node p ( x 1:4 ) correspondingto two distinct factorizations.
By scaling up to all possible orderings, AO-ARMs inadvertently modelthe same distribution with a huge amount of redundancy.
Such redundancy makes training inefficientand, much more importantly, leads to worse asymptotic performance.
Due to finite capacity of themodel, 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 keyinsights for reducing redundancy and improving model performance.
The first insight is that arbitraryconditional inference can still be computed efficiently without relying on all possible univariateconditionals.
For example in Fig.
1b, omitting the edge p ( x 1 | x 2 , x 3 , x 4 ) still maintains efficientinference for all the nodes in the graph, since there is another path for computing p ( x 1:4 ) .
Basedon this observation, we reduce redundancy of AO-ARMs by training on a smaller set of univariateconditionals that still supports tractable arbitrary conditionals inference.
The second insight is thatsome univariate conditionals are evaluated much more often than others during inference.
Therefore,we upweight the training loss for the more frequently occurring univariate conditionals, therebyaligning the training and inference objectives more closely.
Combining these insights, MAC leads to a strictly-improved formulation of AO-ARMs that, amazingly, gives better arbitrary conditional and joint likelihoods with no compromises on tractability.
Weachieve state-of-the-art likelihoods for arbitrary conditional modeling across multiple domains suchas text (Text8), image (CIFAR10, ImageNet32), and continuous tabular data.
Finally, we conclude bydemonstrating a novel application of AO-ARMs to robot shared-autonomy.
Background
Problem definition Marginal and conditional inference are important queries for handling partialevidence on arbitrary subsets of variables, which is relevant for tasks such as masked languagemodeling or image inpainting.
We let a mask e be a subset of variables, and its cardinality | e | bethe number of elements in the subset.
A masked input x e is a partial instantiation on the variablesubset e using the values taken on by input x .
Marginal inference refers to queries of the formp ( x e ) =
R p ( x e x q ) d x q , where q = X \ e and x e x q denotes the union of the partial instantiations.
Conditional inference queries p ( x u | x e ) = p ( x u x e ) /p
( x e ) are similar and can be computed as aratio of two marginals, so we will refer to marginals and conditionals interchangeably.
To evaluatemodel performance on marginal inference queries over a test dataset X test , we assume a test maskdistribution M over the 2 N possible masks.
We sample masks e from M independently from data,and use the negative marginal log-likelihood of the masked input x e as the loss.
[Equation]
Computing arbitrary marginals and conditionals has been a long-standing challenge in probabilisticinference, with a rich history of techniques [25, 23, 13, 5, 36] ranging from belief propagation,variational inference, to MCMC.
Such queries are inherently difficult in high-dimensions: givenaccess to joint likelihoods p ( x ) , naïve computation of marginal likelihoods p ( x e ) requires integrationover the missing N − | e | dimensions.
As a result, most traditional approaches model the jointdistribution using non-neural architectures that admit efficient inference but are less expressive [16].
Autoregressive Models Autoregressive models are an influential family of generative models thatrepresents complex high-dimensional distributions by modeling a single dimension at a time.
Theyparameterize a joint distribution p ( x ) over N dimensions by factorizing it into univariate conditionals
Q Nt =1 p ( x σ ( t ) | x σ ( <t ) ) via chain rule (Fig. 1a), using an ordering σ of the N variables.
We write σ ( t )and σ ( < t ) to denote the masks corresponding to the t -th element and the first t − 1 elements of theordering, respectively.
The univariate conditionals can be learned via a weight-tied neural network, and composing themtogether leads to highly expressive architectures in practice, such as PixelCNN [34, 29] or Transformers [35].
Autoregressive models support efficient inference on the joint distribution, and on a specifictype of marginal query: ones where the mask forms a prefix of the variable order, i.e., σ ( ≤ | e | ) =
e(Fig. 1a).
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 arbitrarymarginal inference.
2.1 Any-Order Autoregressive Models (AO-ARMs)
AO-ARMs [33, 32, 12] learn a single model that can generate the joint distribution autoregressivelyusing any ordering of the N variables, by modeling
Q [Equation]
for all orderings σ .
Even though there are N ! different orderings, their chain-rule factorization is built from univariateconditionals, of which there are “only” N 2 N − 1 .
Therefore, an AO-ARM is a model that defines the
N 2 N − 1 distinct univariate conditionals p ( x j | x e \ j ) , where j ∈ e .
To answer a marginal inference query p ( x e ) with an AO-ARM, we choose an order σ that iscompatible with e , so that σ ( ≤ | e | )
=
e .
Then, we simply evaluate each of the univariate conditionalsp ( x e ) =
Q | e | t =1 p ( x σ ( t ) | x σ ( <t ) ; σ ) in the autoregressive factorization of x e .
Training AO-ARMs Architecturally, an AO-ARM models all univariate conditionals via a weighttied neural network, by using a special (so-called “absorbing-state” [1, 12]) token for variables thatare not present in the evidence set.
To enable parallel optimization, the AO-ARM architecture isdesigned so that given evidence x e , the model can predict the univariate conditionals p ( x i | x e ) for
alli ∈
X \ e at once.
Importantly, this parallelization works on non-causal architectures, opening thedoors to architectures similar in flexibility as diffusion models [30, 31, 11].
In previous works [33, 32, 12], AO-ARMs are trained to maximize the joint likelihood of a datapointx under the expectation over the uniform distribution U σ of orders.
[Equation]
This objective can be simplified by treating t as a random variable with a uniform distribution U t overcardinalities 1 to N .
We let M card-edge be the distribution over the tuple ( σ ( t ) , σ ( < t )) where σ ∼ U σand t ∼ U t , giving the following loss function for an AO-ARM parameterized by θ .
[Equation]
In practice, due to parallel optimization, we write the objective in the following equivalent form.
[Equation]
Compared to other arbitrary conditional models, AO-ARMs consistently give state-of-the-art performance on benchmarks across a range of domains [33, 32, 12].
Besides good empirical performance,they also serve as a conceptual unification between autoregressive models and diffusion models [12, 1],showing promise for both practical and theoretical advancements.
Improving Any-Order Autoregressive Models with MAC
Despite the success of AO-ARMs, in this section we identify two key deficiencies in previousformulations of AO-ARMs.
We propose solutions to them by re-interpreting AO-ARMs fromthe persepective of recursive decomposition over marginals.
Our resulting method MAC leads toconsistent improvements and state-of-the-art likelihoods across a range of domains in Sec. 4.
(A)
The first deficiency is that of redundancy.
While being order-agnostic is exactly whatenables arbitrary conditionals, learning the same distribution with multiple orderings makes
training inefficient.
More importantly, since the model has limited capacity (especially with theweight-tied architecture), learning orderings redundantly hurts asymptotic performance.
(B) The second deficiency stems from the AO-ARM training objective in Eq.
This objec -tive focuses only on the joint likelihood, leading to a mismatch with the marginal likelihoodevaluation objective in Eq.
1, even though computing marginals is the key feature of AO-ARMs.
When viewing AO-ARMs as autoregressive models that generate using all orders, it is not obvioushow 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, e.g., queries that are prefixes of the orders weremoved.
Similarly for (B) , it is unclear how to choose a distribution of orderings to best account forarbitrary conditional inference.
To solve these issues, we re-interpret AO-ARMs from the perspectiveof computing arbitrary conditional probabilities via recursive decomposition over masks.
3.1 Re-interpreting AO-ARMs as recursive decomposition on a binary lattice
A mask e is associated with the set of marginal inference queries of the form p ( x e ) .
We will refer tothis set of queries as the corresponding task for mask e .
We represent these masks/tasks as nodes in
Fig.
Supporting efficient computation of each task, therefore, implies efficient arbitrary marginalinference.
However, for masks with high cardinality, the corresponding task involves learning ahigh-dimensional distribution.
Fortunately, these masks do not actually form standalone tasks, but are subproblems of one another.
In particular, given non-empty mask e and a single variable j ∈ e , we can write p ( x e ) asp ( x e \ j ) p ( x j | x e \ j ) .
Then, we delegate the computation of p ( x e \ j ) to the task associated with maske \ j , and are only left with estimating the univariate conditional p ( x j | x e \ j ) .
From this perspective, we can view AO-ARMs as solving an ensemble of intertwined tasks, one foreach of the 2 N possible masks.
We visualize this as a binary lattice in Fig.
2a over N = 4 variables,where we have 2 4 = 16 nodes representing the possible masks/tasks, with an edge connecting nodes
1, 1, 1, 2, 2, 3, 1,2, 1,2, 2,3, 1,3, 1,2,3,    3 ∅ p ( x 3 | x 1 ) p ( x 3 | x 2 , x 4 ) p ( x 2 | x ∅ ) (a) 1, 1, 1, 2, 2, 3, 1,2, 1,2, 2,3, 1,3, 1,2,3,    3 ∅ (b) 1, 1, 1, 2, 2, 3, 1,2, 1,2, 2,3, 1,3, 1,2,3,4   ∅ (c) e ′ and e if e ′ =
e \ j for some singleton variable j ∈ e .
To answer a marginal query p ( x e ) , wechoose an element j ∈ e and compute p ( x e )
=
p ( x e \ j ) p ( x j | x e \ j ) , which corresponds to movingalong the edge from node e to node e \ j in the lattice.
We proceed recursively on p ( x e \ j ) until wereach p ( x ∅ )
= 1 .
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 probabilitythat different marginals and univariate conditionals are computed.
Definition 1 A decomposition protocol w defines, for each non-empty mask e , a probability distribution w e over the elements in e .
In other words, P j ∈ e w e ( j ) = 1 and w e ( j ) ≥ 0 for e ̸ = ∅ .
A decomposition protocol w specifies a process for answering marginal inference queries.
Given amarginal query at node e , we sample an element j ∼ w e and compute p ( x e ) =
p ( x e \ j ) p ( x j | x e \ j ) .
We proceed recursively on p ( x e \ j ) until we reach p ( x ∅ ) = 1 .
Definition
A mask distribution is a probability distribution over the 2 N possible marginals.
Anedge distribution is a probability distribution over the N 2 N − 1 possible univariate conditionals.
Mask and edge distributions are relevant for specifying the test mask distribution (Eq. 1) of theproblem domain, and the training edge distribution for the AO-ARM.
When viewed under this recursive decomposition interpretation, previous formulations of AO-ARMsuse a random decomposition protocol w -RND where w -RND e ( j )
= 1 / | e | .
Moreover, they use atraining edge distribution M card-edge (Eq. 3), which unfortunately neglects to take into account the testmask 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 acustom decomposition protocol and deriving a training edge distribution that aligns with the test maskdistribution M .
(A) Reducing edge redundancy To answer arbitrary marginals, we don’t need to learn to generateusing all possible orderings.
Rather, it suffices to decompose each task into any one of its childrentasks.
This constraint corresponds to learning a 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 obtainedby removing the largest element of the node’s mask.
This is precisely the decomposition protocolthat we use for MAC: w -MAC e ( j ) = 1 if j is the largest element in e .(B)
Matching edge distribution When presented with a marginal query e , the decompositionprotocol w traverses a path from node e to node ∅ , evaluating the univariate conditionals along theway.
Simulating w on masks e ∼ M leads to an induced edge distribution D M,w that representshow frequently different univariate conditionals are evaluated during the inference process, whichwe depict in Fig.
2c using red edges with different thicknesses.
We can see, for example, thatunder the protocol w -MAC, univariate conditionals such as p ( x 1 | x ∅ ) are evaluated highly frequentlybecause it appears as a subproblem to many marginal queries.
Therefore, we should tune our trainingdistribution on univariate conditionals to match this induced edge distribution D M,w given by thedownstream test mask distribution M and the decomposition protocol w -MAC.
Abstractly, let π ∼ w ( e ) denote the probability that w chooses a path of edges
π = { ( i t , e t ) } | e | t =
1when decomposing e .
[-In particular, e 1 = ∅ , and i t is a variable in e \ e t , and e t +
 = i t ∪ e t . Thenwe-]Then we can write the marginal inference objective as follows, where C is theconstant E i,e ∼ D M,w | e | .
[Equation]
Eq.
5 leads to the following algorithms for training MAC (Alg. 1).
We sample a marginal query p ( x e )by independently sampling a mask e ∼ M and an input x ∼ X .
Then, we simulate the path taken onthe lattice by the decomposition protocol w by iteratively sampling an element j ∼ w e and updatinge to be e \ j .
Throughout this process, we optimize our model on the edges (univariate conditionals)traversed on this path.
During testing (Alg. 2), we similarly step through the decomposition protocoland add together the univariate conditional likelihoods along the path.
Input: Test mask distribution M , decomposition protocol w , training data distribution X , model θ 1 while training do e ∼ M, x ∼ X while e ̸ = ∅ doj ∼ w eθ ← θ + ∇ θ log p θ ( x j | x e \ j ) e ← e \ j
Input: Mask e , decomposition protocol w , test data x , model θ
Output: Marginal log-likelihood log p θ ( x e ) 1 r ←2 while e ̸ = ∅ doj ∼ w er ← r + log p θ ( x j | x e \ j ) e ← e \ j Return: r
In practice, since the decompositition protocol w -MAC deterministically removes the largest element,there is an efficient way to directly sample from D M,w -MAC in batch (Eq. 6), without having tosimulate w -MAC as in Alg. 1.
We describe these details in the Appendix.
Finally, to conclude thissection, we discuss two important design choices and practical considerations for w and D M,w .
Parallel training As noted in Section 2, the architecture of AO-ARMs enables training on univariateconditionals of the form e in parallel.
Therefore, rather than dealing with a training distributionover univariate conditionals D M,w ( i, e ) , we instead work with one over masks D M,w ( e ) .
To a closeapproximation, we will simply let D M,w ( e ) ∝ P i D M,w ( i, e ) , visualized as shaded nodes in Fig. 2c.
Cardinality Reweighting (heuristic)
We present one additional technique for tuning the trainingmask distribution, with ablations and further discussion in later sections.
To foster generalizationof our weight-tied neural network, we reweigh the probability of a mask based on its cardinality e ,by using a final training distribution of D CR M,w ( e ) ∝ (1 + | e | ) D M,w ( e ) , which we sample using SIR(Sampling-Importance-Resampling).
With parallelism and cardinality reweighting, the final objective for MAC is the following:
[Equation]
Experiments
We evaluate MAC on high-dimensional language and image domains, and on a set of continuoustabular benchmarks.
We focus on two metrics: joint likelihood and marginal likelihood of the test set.
On both metrics, MAC shows state-of-the-art performance among arbitrary conditional models on themajority of benchmarks.
We conclude by demonstrating a novel application of arbitrary conditionalmodels on a simulated robot shared-autonomy task, where we learn a conditional policy in a 9-DoFaction space through implicit behavioral cloning with MAC.
Note on computation Each run was done on a single NVIDIA A40.
For the language and imageexperiments, we trained each model for approximately two weeks.
Although this was not enough tomatch the total number of epochs trained by the baseline ARDM 1 , we were still able to show state-ofthe-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.
[Table]
4.1 Language
We learn a character-level model on chunks of 250 characters using the Text8 dataset [21], whichconsists of 100M characters.
We follow the setup from previous work in modeling the chunkedtext segments without any additional context, and use the same 12-layer Transformer architecture as
ARDM and D3PM [12, 1].
We report the joint and marginal likelihoods in terms of bpd (lower isbetter), where the marginal bpd is computed w.r.t.
to the test mask distribution M card-mask ( e ) ∝ 1 / (cid:0) N | e | (cid:1) as used in [32] 2 .
In Tab. 1 we see that MAC outperforms all existing arbitrary conditional modelsin joint likelihood and comes close to a standard Transformer trained only to model in one variableorder, 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 to ARDM.
Ablations In Fig.
3 we present ablation studies on different choices of training mask distribution.
The baseline method (ARDM) trains on the M card-mask mask distribution (grey).
Using insight (B)from Sec. 3 with a random decomposition protocol leads to a slight improvement D M,w -RND (light blue).
Adding on insight (A) by using the w -MAC protocol leads to more improvement D M,w -MAC(light orange).
Lastly, the cardinality reweighing heuristic improves both methods even more (darkblue/orange).
The best performance is given by D CR M,w -MAC in dark orange, which is the setting weuse for MAC in all our other experiments.
4.2 Images
We evaluate MAC on CIFAR10 [17] and ImageNet32 [4, 6], both of which consist of images ofdimension 3072.
Again, we follow all experimental settings from previous work [1, 12, 15], using a UNet with 32 ResNet Blocks interleaved with attention layers, and use rotation/flip data augmentation.
In Tab.
2&3 we see that MAC consistently outperforms ARDM on joint bpd and marginal bpd,evaluated on the M card-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]
[Table]
[Table]
4.3 Continuous Tabular Data
Next
[Table]
we consider a tabular domain with variables taking on continuous values [32].
The mainchallenge with continuous values for AO-ARMs is the parameterization of the univariate conditionals,since 1-D Gaussians or even mixtures of Gaussians have limited expressiveness.
To this end, ACE [32]proposes to parameterize the 1-D conditionals as EBMs.
We build MAC on top of ACE, keeping allhyperparameters and experimental setup the same, modifying only the training mask distribution.
In Tab. 4&
[Table]
we see that MAC gives better marginal and conditional likelihoods than ACE on mostsettings
[Table]
and outperforms other arbitrary conditional models such as SPNs [27, 22] and ACFlow [19]
[Table]
[Table] [Table]
[Table]
4.4 Shared Autonomy on FrankaKitchen
Finally, we demonstrate an application of arbitrary conditional models on the task ofshared autonomy in robotic manipulation.
In shared autonomy settings, an operatoraims to control a policy in a complicated action space, but is allowed to delegatepartial control to an AI model for easier manipulation.
We focus on the simulated
FrankaKitchen environment [10, 9], where the goal is to control a 9-DoF roboticarm to move objects around a kitchen.
We train a base policy with behavioral cloning (BC) [26] and an arbitrary conditional policy with
MAC to solve the Kitchen-mixed0 task.
To simulate shared autonomy, we evaluate a hybrid policywhere at each state, the operator policy (BC) outputs 4 out of the 9 dimensions of the action space,and a conditional policy fills in the remaining 5 action dimensions.
In Tab. 6 we see that conditioningon MAC improves shared autonomy as compared to conditioning on an independent BC policy, andcomes close to the full autonomy performance [8].
Discussion
[Table]
Sharpen mask distribution
Intuitively, weshould choose a protocol w such that D M,whas low entropy so that our model has a greaterchance of revisiting similar masks.
To investigate the differences in choices of decompositionprotocol, in Table 7 we present Monte-Carlo simulations of the induced mask distributions withrespect to the test mask distribution M card-mask .
As expected, using w -MAC leads to lower empirical entropy than using w -RND.
Mask generalization (and why we do Cardinality Reweighting)
The success of AO-ARMsdepends on generalization to unseen masks, since there are exponentially many masks.
Training onmany 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, a decomposition protocol w that skews the mask distribution may have the unfortunateside-effect of hurting generalization.
In particular, the mask distribution D M,w -MAC puts muchmore weight on low-cardinality masks.
This is problematic since training on cardinality- 1 masksis not conducive to generalization: there are only N possible cardinality- 1 masks, whereas highercardinalities have exponentially more masks.
As such, we proposed D CR M,w -MAC (Table 7) with thehypothesis that focusing less on low-cardinality masks will lead to better generalization.
Our ablationstudy 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
D M,w ( e ) ∝ P i ∈
X \ e D M,w ( i, e ) .
However, each pass of D M,w ( e ) trains on all of { D M,w ( i, e ) : i ∈
X \ e } , which overtrains neighboring edges and leads to a minor distributional mismatch.
Aligningthis training distribution more closely may lead to improvements, but requires better techniques forparallelization and sampling of masks.
We leave this for future investigation.
Conclusion
We present MAC, an improved method of training an inference on AO-ARMs.
MAC trains on acarefully 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.
Ourmethod leads to better joint and arbitrary conditional likelihoods with no sacrifices on tractability.
We show state-of-the-art results in arbitrary conditional modeling on
, image, continuous
domains and present a novel application of arbitrary conditional models on robot shared-autonomy
[-• Batch size 128• Learning rate 1 e- 4 , with beta parameters ( 0 . 9 / 0 . 999 ), using Adam• Gradient clipped at 
-]