Learning-Order Autoregressive Models with Application to Molecular Graph Generation
Authors: Zhe Wang
Jiaxin Shi
Nicolas Heess
Arthur Gretton
Michalis Titsias
Venue: ICML 2025 poster (2025)
Keywords: Generative Modeling
Variational Inference
Autoregressive Model
Graph Generation
Molecule Generation
Primary area: probabilistic_methods->variational_inference

Abstract: Autoregressive models (ARMs) have become the workhorse for sequence generation tasks, since many problems can be modeled as next-token prediction. While there appears to be a natural ordering for text (i.e., left-to-right), for many data types, such as graphs, the canonical ordering is less obvious. To address this problem, we introduce a variant of ARM that generates high-dimensional data using a probabilistic ordering that is sequentially inferred from data. This model incorporates a trainable probability distribution, referred to as an order-policy, that dynamically decides the autoregressive order in a state-dependent manner. To train the model, we introduce a variational lower bound on the exact log-likelihood, which we optimize with stochastic gradient estimation. We demonstrate experimentally that our method can learn meaningful autoregressive orderings in image and graph generation. On the challenging domain of molecular graph generation, we achieve state-of-the-art results on the QM9 and ZINC250k benchmarks, evaluated using the Fréchet ChemNet Distance (FCD), Synthetic Accessibility Score (SAS), Quantitative Estimate of Drug-likeness (QED).

Learning-Order Autoregressive Models with Application to Molecular Graph Generation
Zhe Wang 12
Jiaxin Shi
Nicolas Heess
Arthur Gretton 12 Michalis K. Titsias
Abstract
high-dimensional data distribution using the product-rule factorization of conditionals. They have become the stan- dard approach for generating text (Brown et al., 2020).
Autoregressive models (ARMs) have become the workhorse for sequence generation tasks, since many problems can be modeled as next-token pre- diction. While there appears to be natural order- ing for text (i.c., left-to-right), for many data types, such as graphs, the canonical ordering is less ob- vious. To address this problem, we introduce a variant of ARM that generates high-dimensional data using probabilistic ordering that is sequen- tially inferred from data. This model incorporates trainable probability distribution, referred to as an order-policy, that dynamically decides the au- toregressive order in a state-dependent manner. To train the model, we introduce a variational lower bound on the log-likelihood, which we op- timize with stochastic gradient estimation. We demonstrate experimentally that our method learn meaningful autoregressive orderings in im- age and graph generation. On the challenging do- main of molecular graph generation, we achieve state-of-the the---r results on the QM9 and ZINC250K benchmarks, evaluated across key metrics for dis- tribution similarity and drug-likeless
A key limitation of traditional ARMs is their reliance on a predefined, canonical ordering for factorizing the probability distribution over data dimensions. While text data possesses natural left-to-right ordering, many other data types, such as images and graphs, lack such an obvious or universally applicable ordering. Indeed, the optimal ordering may vary even between individual data points, and the best generation order is often context context-depente dependent. Therefore, it is desirable to derive extensions of ARMs that do not treat the ordering as fixed, but rather as a latent random variable that follows probability distribution that adapts to the evolving state of the generation process.
One basic way to extend ARMs with probabilistic orderings is any order ARMs (AO-ARMs, Uria et al., 2014), where the ordering follows a uniform distribution over all possible per- mutations of data dimensions. Such models can be thought of as order-agnostic, and also connect with the masked or absorbing discrete diffusion models (Austin et al., 2021; Hoogeboom et al., 2022; Shi et al., 2024; Sahoo et al., 2024; Ou et al., 2024). However, in practice such models are less effective in terms of likelihood scores compared to ARMs (Yang et al., 2019; Hoogeboom et al., 2022). A possible reason that these models try to solve an extremely chal- lenging training problem, which requires fitting neural network adequately match all possible conditional distri- butions over all orderings of the data dimensions, without learning any preference towards particular orderings.
1. Introduction
Artists draw paintings in their own styles stroke by stroke, and chemists synthesize stable molecules step by step fol- lowing the procedure chemical reactions. What is com- mon in both cases is that the next data component to generated is dynamically determined by the current state of the creation process. In machine learning, autoregressive models (ARMs) can encode such sequential dependencies between data dimensions. In general, ARMs represent
To address the limitations of fixed order and any order ARMs, we introduce an ARM variant which can ficxibly learn probabilistic orderings for generating the data dimen- sions. Our main contributions and findings include:
We introduce the Learning Caaning--Oddr Autoregressive Mod- els (LO-ARMs), novel generative model that learns context-dependent generation orders from data. Specif- ically, we extend AO-ARMs to incorporate a trainable probability distribution that dynamically decides the sampling order of the data dimensions. We derive a variational lower bound on the exact likelhooo for
'Google DeepMind University College London. Correspon- dence to: Zhe Wang <zhewang @ google.com>>, Jiaxin Shi <jiax- ins @ geggooll.com. Michalis K. Titsias <mtitsias @ google.com>>.
Proceedings of the 4nnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Learning- Autoregressive Models
18) Molecule at 1-1 (b) Unnaaske molecule a Stepeo
(c) Policy over masked positions at
[6] Entropy al each masked dimension
sampled dimereson
Value distrbutiion
Slep41
0.5
0.0
0.00
0.0
0.4
0.002
0.2
Buidd Skeleton
0.000
0.0
0.0
Stepre
Stept99
0.05
1.00
0.04
0.50
0.02
0.5
Proces
Praden-Faddee
drodes
#nadec+ dadgen
drgle couple tripe aromatic
Step#18
Step#19
0.04
0.03
04
0.1
At Step 19, the RDKE visualzeer rotates the right aromatic rig.
#hodes
Phoce
Prodee- Bedges
mage engle diaile tiple anmatic
0.0100
Stopa119
Stepa220
1.00
8.0075
0.75
90030
0.0025
0.25
Intil
$10005
4n10005
Pnopes Pooges
sengtee doodte pre aromeoc
Step11551
Stepf1122
3.00
0.75
0.75
0.25
0.25
Arcolea Aedggo
Artocla Audgo
ru.odge single dodale trole unomals
S1ep1/52
S1ep4153
0.05
1.0
0.73
0.04
0.25
Proces
Anodece mnodece Redoes
Bhoces
#nodece #nodessdeedees dedoes
N
0
Inii
Slepk1170
Stsp
1.00
1:00
0.3
0.75
0.5
0.50
0.25
0.1
0.00
0100
00
2500
Aodgos
Anodes
Polges
Figure 1. An example of molecule generation with our best LO-ARM model trained on ZINC250k. Our model generates molecules st-ep, commencing with all nodes and edges masked (in the figures masked nodes are labeled as x) and adding node edge at time. First, an order-policy selects which dimension (node or edge) to fill, and then a classifier determines its value. Each step is illustrated in the provided figures: Columns (a) and (b) illustrate the (partially) generated molecular structures generation steps. Columns (c), (d), and (e) provide detailed insights: (c) the order-policy's 's probability distribution over dimensions, (d) the classifier's entropy for each dimension, and (e) the classifier's 's prediction the selected dimension. Note that, the order-policy and classifier entropy are zeroed for unmasked dimensions. The generation proceeds through three phases: bond skeleton construction, "no-edge" state population (representing "imaginery" bonds in the dense adjacency matrix), and atom type assignment. particular, the order-poll policy learns to favor dimensions with high classifier certainty, demonstrated by an inverse relationship between policy probability and classifier entropy. For example, at Step I, dimensions with higher probabilities (c) correspond to lower entropies in (d). This trend is especially evident during the "no-edge" infilling phase, where masked edge dimensions at Step 20 and 152 exhibit high order-policg probabilities and near-zero classifier entropies, while masked node dimensions show the inverse. Importantly, the classifier assigns probabilities potential values, enabling valid molecular variations. For instance, at Step 171, as shown in (c), the final masked atom has a near-equal chance of being Carbon, Nitrogen, or Sulfur, resulting in three different yet valid molecules. Step 19 also demonstrates this, with each possible edge value offering different, yet feasible, molecular structure. We provide the full generation path in Appendix B.2.
successive
training the model using stochastic gradient estimation.
art results for both datasets, as measured by sample quality metrics such as Frèchet ChemNet Distance (FCD) (Preuer et al., 2018), Synthetic Accessibility Score (SAS), Quantitative Estimate of Drug-likeness
We apply LO-ARMs to two molecule generation tasks (i.e., QM9 and ZINC250K) and obtain state-of-the-
2

---

Learning-Order Autoregressive Models
(QED), and sample diversity.
point x (so training for each example xx() there is different o (())), the log-likelihood can be lower bounded as log (z) log E. pp(zlo)) 2 bogt)]logpie [)],
We find that LO-ARMs can learn consistent order- ings to generate new molecules with high validity and uniqueness.
which yields the training objective in Equation (3). Note that in practice this objective is optimized stochastically; see Uria et al. (2014) for more details.
We investigate a variety of architectural choices for I.O- ARMs and provide an ablation analysis in the context of molecular graph generation.
2.2. Autoregressive Generation as Unmasking Process
Here, we describe a convenient way to characterize the sequential generation process for discrete data which we will use for building our proposed method in Section 3.
2. Background
2.1. Problem Setup
Given that each discrete data dimension (or token) xiz, i {1, Lb} takes m categorical values we further augment the space with an extra auxiliary category or mask. Thus, we represent each r; as an m + 1-dimensional (rather than m- dimensional) one-hot vector where the final m + 1-th value indicates that 2; is masked. Then we can model the genera- tion process with ARMs as an unmasking process. Specifi- cally, starting with a fully masked state æ (Z1...., .t). where 3e represents the mask, each step, we choose dimension T; and "unmask it" which means to sample a categorical value 2; among the m categories. Then we re- peat this process until all dimensions are unmasked, which yields a final generated data point æ. Under this representa- tion our model also connects with recent masked discrete diffusion models (see, e.g., Shi et al., 2024) as we further discuss in Related Work. However, note that unlike discrete diffusson models we do not specify forward process in framework, but only the unmasking or backward process.
We assume data vectors æ = (I1,...,TI) where each di- mension z; takes values from a set X' which could be real or discrete. Without loss of generality, we assume a a vector of L categorical variables, so that X' is a discrete | categories. For example, for molecular graphs, X' VUE, where V and E are the atom and bond types respectively. An autoregressive model defines joint probability distribution over æ that factorizes as
set of
L po(x)) = II po(
where Ty denotes the i-th dimension of x; Tei (Z1...., denotes the first - elements of the vector a and (1, (e)))) is the conditional distribution with the = pe (21).
convention
Sampling from the model generates the data dimensions sequentially starting from 21 and ending at æT. As men- tioned in the introduction, having fixed or pre-specified order can often be disadvantageous, since it may introduce inappropriate inductive bias when modeling data without a natural ordering. AO-ARMs (Uria et al., 2014; Hooge- boom et al., 2022) address this problem by training a model that can generate the data dimensions under any random ordering drawn uniformly from L! permutations of the in- dices (1.....I) T.J. Given permutation a the model joint distribution factorizes as
3. Learning Order ARM
We replace the uniform prior distribution p(o) in AO-ARMs by with a learnable distribution over orderings. We call this distribution the order-policy since it dynamically decides the next dimension to generate by conditioning on the already generated data dimensions of x. Next, we focus on the gen- eral description of our method and the key design choices, and delay the technical details of specifying it to practical tasks, including images (Appendix A) and molecular graphs (Appendix E) generation. Specifically, this section unfolds as follows. We introduce the order-policy in Section 3.1, de- fine the learning objective in Section 3.2, and then describe options for the model parametrizations and neural network architectures in Section 3.3.
7. (o(0) = II pol (s
(2)
where dee denotes the indices of the i --1 first elements under the permutation If p(o) denotes the uniform distri- bution over the L! permutations, the parameters 0 are found by maximizing the expected log-likeliine likellhoooodd (per data point):
3.1. Model with Order-Policy
Ep(c) [log (x)
(3)
Sampling an ordering can be represented by set of L latent variables 2,, i = 1,..., L,, so that
As noted by Hoogeboom et al. (2022), a way to interpret the above objective is as a variational lower bound on the log- likelihood of a probabilistic latent variable model. Specif- ically, if a is the latent variable corresponding to the data
L.
p(z) = Il p(z (2)))
(4)
3

---

Learning- Autoregressive Models
where 21 p(z1)= p(=1|=<1) is categorical variable that takes L values from the set (1, (1,..., .,, I), and each sub- sequent 2; p(zi|zze takes L - + 1 values from the set 221 fl...., [} Rec,, where z<z= = (21,....24-)) 2-1). The probability distribution po(x)) can be written as
where the factors po Reb Fx and ( de- pend on parameters that we want to learn. Since the exact likelihood is intractable, we will maximize an evi- dence lower bound (ELBO) on the log- ligeliikdd We use an amortized variational distribution over z that conditions on the full data vector æ, and has the general form
L (a(z)== = p(z)po = III (e (Za))
qa (z ax)
II to (2 2=1
where the two conditionals p(2; =<z) and po (2z, (a fol- low side-by-side an autoregressive structure that unfolds from i = 1 to i = L. When each p(2 (<)) is uniform then p(z) is an autoregressive representation of the uniform dis- tribution over the L! permutations. In this case, the model reduces to standard AO-ARMs.
where the specific parametrized form of qa is given in Equa- tion (12) and Section 3.3. Structurally, each variational factor (0(2; a:) has similar form as the order-policy factor po (z; Z:i Tx ). but the difference is that the former is allowed condition on the full æ while the latter only z This amortized factorization is one of our key design considerations, and makes model training practically feasible while still allowing us to efficiently compute an unbiased estimate of the ELBO. We will discuss this point at the end of the section.
In our proposed method, named learning order ARM (LO-ARM). we use the latent variables but we replace p(zi|z<s) with more informed distribution. We call this distribution order-policy and we define it follows.
Definition 3.1. The order-policy is distribution over that follows the factorization
Using go we can lower bound the log likelihood as follows:
pe (2)
II Po (2 (zil Fzee),
log po (x) 2
Σ g0 log
(z x) qo (0(2(ar))
po Zec Ei )oo (Jiz (zec)
where each factor po (= ) is a parameterized cate- gorical distribution over the index z of the next data dimen- sion in the autoregressive order that conditions not only on the indices Zςi but also on the corresponding data dimen- sions Ezc. (Zzz....,, (...))
Σ (z(()))) log
q0 (2ci w)
Eapis Euotesl log D0(2ilzci, Pe D ppola 1
q0(zi|
The order-policy enables modeling of the sequential depen- dence structure of data dimensions that can exist in the true data distribution. Therefore, is desirable to flexibly learn this distribution from data. full description of specific parametrization of the order-policy and an overall model architecture given in Section 3.3. Before proceeding to this, discuss general training procedure in the next sec- tion, based on variational inference and stochastic gradient estimation.
E Eaplaccle) [F((ze))))]
(8)
To obtain the final expression, we performed the exact expec- tation over E Zzi = {I,.... over the set of all currently masked dimensions) and defined the function Fo 2) as shorthand for the long expecta- tion highlighted blue. We also analytically marginalized out all future latent variables z>: = [1..... ,, \ æ since the function Pg (2 x) does not depend on these. Comput- ing the full ELBO is too expensive, and thus in practice we construct an unbiased estimate by sampling one term in the sum Σ/= together with one z<i ~ qe(2 (Z)) Thus we obtain the stochastic estimate of our loss as
......... L} Z (i.e., taking values
3.2. Training with Variational Inference
To train the LO-ARM model from a set of training examples we want to maximize the log-likelihood
D = (x) N
N
((0) = -LF 2ςl go (z< x))
(9)
Σ log (og
Σ log
Σ po an! [n)
For the special case where each variational factor go [zz x)) and po(zi zei, tze.) are non-leaeeable and
For simplicity, we only consider one data point for now and drop index n. The joint distribution po(z, .) can then be factorized as
set to uniform distributions, the stochastic negative ELBO in Equation (9) reduces to
Po (2, x)-
L II
po 22p (
Eze,
(6)
(())
L
Σ 2E
log log ( (10)
1-ith 1-i+1
4

---

Learning-Oree Autoregressive Models
Algorithm 1 Training with LO-ARM
is a non-linear function that receives as input the data a. Based on these logits each variational factor has the form
Given a dataset D. (00((xxx))) pol- 2cci I Lza.). 90(-|z))
while training do
e (m)
Uniformly sample x (Z1,... LL) from D
qa (z. = k(zz,, Evers 200
(12)
Sample =² ~ 2o |2)) and 2² ~ qo (-(2r)
Sample i ~ Uniform[L,
L]
The logits combined in autoregressive way based on the factorization in Equation (7). Computing and sampling from the variational distribution is very fast, since it requires a single evaluation to obtain and store the vector of logits go (z). and sampling from qe in (7), with factors given by (12), has negligible additional cost. In fact full path z can be sampled at once in parallel using the Gumbel-top-k trick do in our implementation. The exact parametrization of go using the neural architecture is discussed next in Section 3.3. Moreover, similar to train- ing variational autoencoders (Kingma & Welling, 2014), we can obtain unbiased gradients of the ELBO by sampling the discrete latent variable z ~ qo, which allows for training go and the model pe jointly with a single optimizer.
Set z/4- = Get Lzz,, 222, by masking dimensions at 2/,00 z
=(_1) and 221 (e) 2)_1)
Compute RI.O unbiased gradient from (11) and up-
date 0 via SGD
end while
Algorithm 2 Unconditional sampling from LO-ARM
(Kool et al., 2019b).
Initialize a fully masked state क = (ZI,.... fr))
as the indices to be sampled for step i = 1, 2,. I do Sample data index 2, ~ Po (f(Z2ci Sample data dimension value #zz; ~ po((|azzel
(1,...,.) } = Initialize & =
Unmask TT, with a..; E +- #) \ Tzs U (&t)
Remove 2; from ž: E <= z \ 2z
Once training is completed, we can generate new samples from the model using Algorithm 2. Specifically, we can generate new sample o = (f1,.... .t) autoregressively starting from a fully masked state ã = this, at step E {I....: ......... LJ, we unmask T.; to u... which is sampled with the categorical distribution. Note that we only need the model order-policy over 2; ~ po(1| Fze.) and the categorical distribution a.. po( (zze.) for inference.
end for Return xC
To do
which is precisely the stochastic objective used for training AO-ARMs; see Equation (12) in Uria et al. (2014). In other words, the objective in Equation (9) generalizes these
previous objectives.
Unlike AO-ARMs, however, the stochastic objective in Equation (9) is more complex, which therefore would have higher variance, In order to propagate gradients through the sampling step from the discrete distribution go and reduce variance, we also need REINFORCE gradient estimation (see, e.g., Shi et al., 2022). To keep it simple, we use the REINFORCE leave-one-out (RIOO) estimator (Salimans & Knowles, 2014; Kool et al., 2019a) with two samples. First we select single random index i E {1,...,L}, Then given this fixed i we draw two sample paths of previous indices z and 22 by sampling from qo (z))). The unbiased gradient over 0 is then written as
3.3. Parametrization of the Distributions
Here, we discuss specific parametric forms for the model and variational distributions. Since Zz, is discrete, parameterize the model conditionals po( for k = 1,..., L as L classifiers with Neural Network (NN) having I. heads,
po po( softmax(fe.($*za)))
(13)
where Ez,, is the state with the z<; dimensions being un- masked and the rest masked, while (0 1 are the logits of the k-th softmax head over m classes. The spe- cific architecture of these classifiers could be task dependent. For instance, we employ UNet (Ronneberger et al., 2015) for image generation (Appendix A) and Graph Transformer (Vi- gnac et al., 2023) for graph generation (see Appendix E).
L
{ (V
V Volog AF
+ VoFo x) (o
(11)
where AF Fo(2 x) - F Φ).
Next, we parameterize the model order-policy factors Po = 2/zz,, xze.)) using L functions h-, E R, k 1,... L so that
Algorithm 1 outlines the whole training procedure. In partic- ular, = (2), ................................ 2 is full permutation of (1, (1...., L} and fast sampling of 2, e.g., by requiring only single NN forward pass, is critical to obtain scalable training. To en- able this, we construct the full Oo as an amortized (by data vector x) Plackett- Luce model (Plackett, 1975). More pre- cisely, we assume the vector of logits go (a:) E R.L, which
pa(z; = kkzze, Everz, Co
(14)
We explore two options to define he,(()) i) An entropy- based parametrization, where these functions are scaled
5

---

Learning-Orde Autoregressive Models
entropies of pe (22)) (computed from logits fe. dis- cussed ahove). This can favor sampling dimensions 2; with higher certainty (about the value of Tx2)) early on. ii) A shared-torso parametrization, where we add an extra final linear layer to the so outputring I, scalars providing the h-)x(-) values.
following the new node. Bu al. (2023) reformulated the node ordering problem dimensionality reduction. These existing works mostly follow paradigm of incrementally adding new nodes and connecting them to existing nodes. By contrast, our work provides a general order learning framework for ARMs, which enables more flexible ordering relationships between nodes and edges.
Finally, we have already described the general form of the variational factors q0(2;= = k|z<i, x) in Equation (12), so what remains is to specify the vector of logits go(2)). We explore two parametrizations: 1) A shared-torso param- eterization where we add another final linear layer to fe (similarly to (ii) above) outputting L values for go(2)). 2) A separate NN with outputs to model (z().. The motiva- tion of this latter option is to more freely capture the form of the variational distribution.
Discrete diffuusion and application to graph generation. Our method also relates to discrete diffusion models based absorbing or masked diffusion (Austin al., 2021; Lou et al., 2024; Shi et al., 2024; Sahoo al., 2024; Ou et al., 2024). Similar to masked diffusion formulations, our spe- cific neural architecture for discrete data assumes that all not-yer-generrted data dimensions are assigned the mask category. One important difference with masked diffusion is that we learn a non-uniform and data-dependent order of the data dimensions a flexible way. using a NN parametrized order-policy a confidence- based (using entropic uncertain- ties) policy. In contrast, masked diffusion methods operate similarly to AO-ARMs (Hoogeboom et al., 2022), i.e. they generate discrete tokens using a completely random order. Further, unlike diffusion-basee methods, our approach does not explicitly specify a forward noising process. Instead, we only define a backward generative model which generates samples from a fully masked state, and our model learns a variational order distribution (ge) from data to encode the to- ken unmasking sequence. Analogous to discrete Variational Autoencoders ( (VAEs), qa is optimized jointly with the gen- erative model using a REINFORCE-styL. gradient estimator. This contrasts with diffusion models like DiGress (Vignac et al., 2023), which rely on a fixed forward process instead of a learned variational distribution.
We provide further details on the above options in Ap- pendix F. In the experiments we present an ablation for these options in the context of molecule graph generation.
4. Related Work
Learning Non-Monotonic Autoregressive Orderings has been studied extensively in recent literature (e.g., Li et al., 2021; Gu et al., 2019; Welleck et al., 2019), and is challenged by the need to find an optimal permutation from a factorial (L!) search space, where L is the sequence length. Some methods reduce this space with domain as- sumptions (Welleck et al., 2019; Gu et al., 2019). Specif- ically, Welleck et al. (2019) proposes use a tree-based recursive generation method to learn abitrary generation orders, and Gu et al. (2019) combines 1) pretaining with prescribed base orderings and 2) fine-tuning those order- ings with Searched Adaptive Order (SAO). More relevant to our work, Variational Order Inference (VOI) (Li et al., 2021) learns orderings directly from data using a policy gradients procedure, but it needs to optimize a complex variational ordering distribution which has an intractable normalizing constant and requires a Bethe-type approxi- mation. In contrast, the variational distribution in our I.O- ARM is fully tractable, and it allows us to use fast and exact unbiased gradient- based optimization of the ELBO using REINFORCE leave-one-out.
5. Evaluation and Analysis
We are interested in two questions: 1) whether LO-ARM can learn meaningful orderings to generate data, and 2) whether the learned orderings can yield better generation performance. To answer these, we apply LO-ARMs two tasks: 1) toy image generation task on the MNIST dataset to showcase the ordering the model learns, and 2) molecular graph generation on the QM9 and ZINC250K datasets (Xu et al., 2019; Dwivedi et al., 2023). In both tasks, the data does not exhibit a "canonical"" ordering. For the MNIST task, we provide a qualitative illustration that LO-ARM prefers to sample first the border pixel values and then the "digit" pixels in the centered foreground. Presumably, the reason is that the model learns to prefer order where the "easier" border pixels are generated first, which would be of higher certainty. We provide an illustrative figure together with a more detailed description in Appendix A. In the present section, we focus on the quantitative evaluation of LO-ARM on molecular graph generation tasks.
Orders of autogressive graph generation. Unlike image and text domains, the generation order of ARMs has been a central focus in the graph domain. Early work on graph gen- erative models (Li et al., 2018) mentioned the difficulty of learning the ordering and eventually settled on using either uniform or fixed ordering. You et al. (2018) introduced a fixed BFS ordering scheme. Chen et al. (2021); ; Kong et al. (2023) extended AO-ARMs to enable a partially learnable ordering over nodes, where the edges connecting a newly added node to existing nodes are generated immediately
6

---

Learning- Autoregressive Models
Table Molecule generation on QM9. We directly eite the results of other methods. The negative loglikelihoods (NI.ss) of our methods are evaluated against the test set, and other metrics are calculated with the generated samples with the corresponding methods.
Method
NLL
Validity %
Uniqueness% %
FCD,
Ground truth (test data)
99.3
100
0.005
GDSS
95.7
98.50
2,900
DiGress
69.6
99.00
96.20
GraphARM
90.25
95.62
1.22
CatFlow
99.81
99.95
0.441
Our Results
P8
AO-ARM
< 24.66
98.88
99.11
0.671
LO-ARM-st-st
shared torso
shared torso separate
X 22.38
99.05
98.59
0.437
LO-ARM-st-sep
shared torso
21.42
99.85
98.85
0.240
We evaluate LO-ARM on two widely-adopted molecular graph generation benchmasks: QM9 and ZINC250K (Xu et al., 2019; Dwivedi et al., 2023). We follow the standard setup e.g., in Eichelsbacher & Reinert (2008); Vignac et al. (2023); Jo et al. (2022), including data preprocess- ing, network parametrization, and evaluation metrics. We represent molecules as graphs with dense adjacency ma- trices (see Appendix E). For each model variant, we gen- erate 16384 samples and evaluate them according to two aspects: 1) validity and uniqueness of individual molecules and 2) Frèchet ChemNet Distance (FCD), which evaluates the distance between the distributions of true and gener- ated molecules using the activations of ChemNet (Preuer et al., 2018). Note that, these two classes of metries are not explicitly correlated, since invalid molecules may have latent activations close valid ones. As we can see in Ta- ble 2, the validity of methods with similar FCDs may vary in big range (e.g., GraphAF, GraphARM, EDP-GNN and SPECTRE). On the other hand, if the true dataset has high quality, we could still expect molecule validity to be im- proved when just optimizing FCD-related objectives, c.g., the lower bound of log-likelihoods. this study, we focus more on the evaluation with FCD show the generative capability of LO-ARM, and we will show how validity and uniqueness correlate with FCD.
chitectures, are described in Appendix E. The results are summarised in Table and Table 2, and samples generated with our best models are shown in Appendix B.
First, on QM9 (Table 1), we can see that all variants of LO-ARM consistently outperform AO-ARM. With simi- lar model capacity, LO-ARM-st-st (both pe and qo share torso with the classifier) achieves similar performance as the previous best method CatFlow. Moreover, with increased capacity for the variational order policy, LO-ARM-st-sep (shared torso for pe and separate NN for qe) achieves new state-of-the -fie--ff-th-arttt results in FCD, as well as competitive va- lidity and uniqueness. Note that the cost of inference with LO-ARM-st-scp remains close to other Graph Transformer- based methods, because only the posterior model order pol- icy (pa) and the classifier are used during inference.
Next, evaluate LO-ARM with ZINC250K (Table 2). which is much more challenging dataset made of larger drug-like molecules, as data dimensions increase in (nn) when representing n-atom molecules as graphs. Compar- ing LO-ARM with our own baselines (i.e., AO-ARM and Biased- AO-ARM), we see that introducing order policies improves the performance with a large margin of gain on all metrics (especially validity). Furthermore, our model demonstrates substantial FCD improvement, indicating higher sample quality than previous molecule generative models in terms of resembling the true data distribution.
To see the effect of the order policy, we introduce two base- lines, i.e., 1) AO-ARM in which both the variational (ga) and model (po) order policies are set to uniform, 2) Biased- AO-ARM (in Table 2), which the order policies always uniformly samples edges first and then uniformly samples nodes after all edges are unmasked. In addition, we fix the parameterization of the classifier Se,z, to be the same as two recent -fate-of-the-arts DiGress (Vignac et al., 2023) and CatFlow (Eijkelboom et al., 2024), and ablate different combinations of the parametrizations of model he.* and variational policy distributions. These parameterizations are described in Section 3.3 and in Appendix F with more details. The training configurations, including network ar-
We observe that there is generally a gap between the valid- ity of ARMs and that of the best method on this task. We hypothesize that this is caused by the errors accumulated during autoregressive sampling, as we unmask new di- mension at a time without refining the previously unmasked. Still, as demonstrated by the gain in validity compared to AO-ARM and GraphARM, LO-ARM has largely fixed this problem by training an order-policy that maximizes the ELBO. Moreover, this accumulated error can be further re- duced by employing Top-p sampling (Holtzman et al.), popular method used in autoregressive language models to
7

---

Learning- Cerningg-Orderrar Antoregressive Models
Table Molecule generation on ZINC250K. particular, the NI..s of the L.-ARMs with the Top-p sampling omitted the sampler is only used in inference time and does not change the training-time time metrics.
Method
NLL.
Validity $
Uniqueness %
FCD4
Ground truth (test data)
100.00
99.88
0.005
GraphDF
90.61
99.63
33.55
MoFlow
63.11
99.99
20,93
SPECTRE
90.20
67.05
18.44
EDP-GNN
82.97
99.79
16.74
GraphARM
88.23
99,46
16.26
GraphAF
68.47
98.64 68.64
16.02
GDSS
97.01
99.64
14,66
CatFlow
99.21
100.00
13.21
Our Results
po
qo
Top-p
AO-ARM
1.0
80.24
32.93
100.00
6.541
Glased-A--RMM
1.0
< 77.94
34.16
100.00
5,026
LO-ARM-st-st
shared torso
shared torso
1.0
& 70.03
90.84
100.00
4.087
I.O-ARM-st-st-oopp
shared torso
shared torso
0.9
96.02
100.00
9.237
L.--ARM-t-seep
shared torso
separate
1.0
S 68.26
96.26
100.00
3.229
LO-ARM-st-scp-topp
shared torso
scparate
0.3
96.70
100.00
3.859
st-sep on validity and uniqueness. While LO-ARM-st-Sae is a more performant generative model in terms of FCD, LO-ARM-st-st--ope is more computationally efficient for training without using separate network. As for inference, their cost would be close as the variational order policy is not needed. Both options could be useful in practice, depending actual needs.
Table 3. LO-ARM's perforamance on dremistry-specitnn metries on ZINC250K. Specifically, lower similarity indicates higher diver- sity for molecules in the dataset. On all these metrics, LO-ARM exceeds or matches the performance the state-of-the-art results of the models. The results of other methods cited from (Mazuz et al., 2023). V. is short for Validity.
Method
QED
SAS
Sim.
V.%*
Interestingly, we find that LO-ARM-st-sepp is more robust to Top-p sampling. Even if we set = 0.3, the downgrade of its FCD is marginal (from 3.229 to 3.859, staying the SOTA). This suggests that the dimensions sampled with the model order policy have highly concentrated probability mass, which therefore implies that the order policy favours dimensions with higher certainty. This aligns with our ob- servation presented in Figure 1.
Ground truth (test data)
0.75
2.76
0.35
100.0
JT-VAE
0.64
4.69
100.00
GCPN
0.65
4.53
99.00
Taiga
0.75
2.89
88.00
LO-ARMM-st-eep (ours)
0.75
3.08
0.34
96.26
address similar problems. Specifically, the Top-p sampler only draws from dimensions with high probability, and p corresponds to the cumulative probability of the top dimen- sions ranked with respect their probabilities. Specifically, p 1.0 regresses to sampling the entire distribution, and the lower the p, the greedier the sampler. We tested this approach by sharpening the distribution of the classifier fe after the policy has sampled a dimension. After incorpo- rating the Top-p sampler, we can see that both the validity of Aoased-AO-ARM-oopp (Top-p- = 0.9) and LO-ARM-st- st-topp (Top-p 0.9) increase by a large margin, while at cost of downgraded FCDs. As expected, the Top-p sampling introduces small bias, shifting the model's 's dis- tribution away from the true data distribution. We provide further ablation studies on this effect in Appendix D.
Thirdly, on the front of practical usefulness, the ultimate goal of the generative model is is to sample from the same chemical space that the training data comes from. On the other hand, the validity and uniqueness metries are fairly saturated by existing methods, and the FCD on its is not strong enough predictor of molecule quality. To ad- dress this issue, we further evaluate LO-ARM-st-sep on the chemistry-specific metrics, including 1) Synthetic Accessi- bility Score (SAS), which measures the case of synthesizing chemical compound (i.e., the lower the score, the easier the compound be synthesized), 2) Quantitative Estimate of Drug-likeness (QED) which evaluates how well pound's physicochemical properties align with those of suc- cessful administered drugs (the higher the score, the more aligned), and 3) pairwise similarity measures the diversity in samples (lower score corresponds to higher diversity). We visualize the distributions on these three metrics Figure 2.
Through incorporating Top-p sampling at test time, LO- ARM-st- st--st-topp achieves similar performance as LO-ARM-
8

---

Learning-Order Antoregressive Models
On both datasets the distributions of these metries calculated on the LO-ARM samples closely match the corresponding ground-truth (GT) data, finding consistent with their fa- vorable FCD scores. We also compare LO-ARM with other methods in the literature on these metrics. As shown in Ta- ble 3, LO-ARM exceeds or matches the previous best results in all three metries.
up LO-ARM to other domains of higher data dimensions, e.g., high-resolution image generation and protein design.
Impact Statement
This paper presents the work whose goal is to advance the field Machine Learning. There are many potential societal consequences of this work, none of which we feel must be specifically highlighted here.
Finally, we visualize the generation ordering learned by LO-ARM on ZINC250K in Figure I and Appendix B.2. Without imposing any inductive bias orderings, LO-
ARM learns to favour an edge-first ordering, i.c, 1) building skeleton of a molecule with chemical bonds, 2) infilling the rest of the adjacency matrix "no-edge" states or the "imaginery" bonds, and finally 3) infilling atoms (See Ap- pendix B.2). This different from the node-first ordering proposed in Kong al. (2023). To analyze the consistency of this learned ordering, we sample 30, 000 molecules with three LO-ARM-st-sep models trained with different random seeds, and find that 99.9% of the molecules are generated with the edge-first ordering (See Appendix C). Note that learning a context- depeedeet ordering (that varies across dif- ferent data/molecules) within the edge or node dimensions is still critical to performance, as we can see by compar- ing LO-ARM with Biased-AO-ARA which just uniformly samples within edge or node dimensions.
Acknowledgements
We would like to extend deep thanks to Yee Whye Teh, Arnaud Doucet and Yujia Li for diligently reviewing this paper and providing valuable comments.
References
Austin, J., Johnson, D. D., Ho, J., Tarlow, D., and Van Den Berg, R. Structured denoising diffusion models in discrete state-spaces. In Advances in Neural Information Processing Systems, 2021.
Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., Agarwal, S., Herbert- Voss, A., Krueger, G., Henighan, T., Child, R., Ramesh, A., Ziegler. D., Wu, J., Winter, C., Hesse, C., Chen, M., Sigler, E., Litwin, M., Gray, S., Chess, B., Clark, J., Berner, C.,, McCandlish, S., Radford, A.,, Sutskever, I., and Amodei, D. Language models are few-shot learners. In Larochelle, H., Ranzato, M., Hadsell, R., Balcan, M., and Lin, H. (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 1877-1901. Curran Associates, Inc., 2020.
6. Conclusions and Discussion
We have introduced LO-ARM, a novel autoregressive method that learn the ordering for generating data with- canonical ordering. We derive simple variational lower bound which can be optimized with unbiased stochastic gra- dients. In addition. we design a scalable training algorithm which makes LO-ARM easily adaptable to real-world tasks. Moreover, we evaluate LO-ARM on two popular molecu- lar generation tasks, and achieve state-of-the-ar results in terms of FCD and QED, and competitive results on molccule validity, uniqueness and SAS. In particular, show that LO-ARM can learn consistent ordering for generating new molecules without requiring inductive bias on gener- ation orderings, which in turn contributes to a significant performance gain for generative modeling.
Bu, J., Mehrab, K. S., and Karpatne, A. Let there be order: Rethinking ordering in autoregressive graph generation. arXiv preprint (.XXxx2305.15500 2023.
Chen, X., Han, X., Hu, J., Ruiz, F., and Liu, L. Order matters: Probabilistic modeling of node sequence for graph generation. In International Conference on Ma- chine Learning, pp. 1630-1639, PMLR, 2021.
Dwivedi, V. P., Joshi, C. K., Luu, A. T., Laurent, T., Bengio, Y., and Bresson, X. Benchmarking graph neural networks. J. Mach. Learn. Res., 24(1), January 2023. ISSN 1532- 4435.
On the task of molecule generation specifically, although LO-ARM yields a more performant generative model in terms of FCD, there is still room to improve the validity of individual molecules compared to diffusion-based meth- ods. We hypothesize that these models may have learned to denoise fragment-like sub-graphs. With LO-ARM, one way to do this is to incorporate more complex molecular constraints into the ordering policies. Another improvement could be to use more memory efficient representation of graphs, such that we do not need store the full adjacency matrix. On a general note, we are also interested in scaling
Eichelsbacher, P. and Reinert, G. Stein's 's method for discrete Gibbs measures. The Annals of Applied Probability, 18 (4):1588-1618, 2008.
Eijkelboom, F., Bartosh, G., Naesseth, C. A., Welling, M., and van de Meent, J.-W. Variational flow matching for graph generation. In The ehirty-eighth Annual Confer- ence on Neural Information Processing Systems, 2024.

---

Learning-Order Antoregressive Models
Figure 2. Evaluating LO-ARM on (a) ZINC250K and (b) QM9) on the chemistry-specific metries i.c, Synthetic Accessibility Score (SAS), Quantitative Estimate of Drug likeness (QED), and pairwise similarity between samples. On both datasets, the distributions of the samples generated LO-ARM on these metrics are well well-matched with the corresponding ground-trutt truth (GT) data, which well-aligned with their performance on FCD scores.
(a] ZINC250K
(b) QM9
QED
Painwise Similarity
Pairwise Similarity
LD-ARM
Deraaty
Denstty
0.0 0.5 1.0
0.0
0.5
1.0
10
0.0
0.5
1.0
0.0
0.5
10
Gu, J. Liu, Q., and Cho, Insertion- based decoding with automatically inferred generation order. Transactions of the Association for Computational Linguistics, 7:661- 676, 11 2019. ISSN 2307-387X. doi: 10.1162/tacl.a 00292. URL https: www.iioorr//11 a_00292.
Sampling Sequences Without Replacement. In Proceed- ings of the 36th International Conference Machine Learning, pp. -3598, PMLR, 2019b.
Langley, Crafting papers on machine learning. In Langley, P. (ed.), Proceedings of the 17th International Conference Machine Learning (ICML 2000), pp. 1207-1216, Stan- ford, CA, 2000. Morgan Kaufmann.
Holtzman, A., Buys, J., Du, L., Forbes, M., and Choi, Y. The curious case neural text degeneration. In International Conference on Learning Representations.
Li, X., Trabucco, B., Park, D. H., Luo, M., Shen, S., Darrell, T., and Gao, Y. Discovering non-monotonic autoregressive orderings with variational inference. In International Conference on Learning Representations, 2021. URL hwt.h:ttt////..... /ttp:://///peeeeeeeeeee net forum?? id= [d=jPlvTH3inC.
Hoogeboom, E., Gritsenko, A. A., Bastings, J., Poole, B., van den Berg, R., and Salimans, Autoregressive dif- fusion models. In International Conference on Learning Representations, 2022.
Jo, J., Lee. S., and Hwang, S. Score-based generative modeling of graphs via the system of stochastic differen- tial equations. In Chaudhuri, K., Jegelka, S., Song, L., Szepesvari, C., Niu, G., and Sabato, S. (eds.), Proceed- ings of the 39th International Conference Machine Learning, volume 162 of Proceedings of Machine Learn ing Research, pp. 13362-10383 PMLR, 17-23 Jul 2022.
Li, Y., Vinyals, O., Dyer, C., Pascanu, R., and Battaglia, P. Learning deep generative models of graphs. arXiv preprint
2018.
Lou, A., Meng, C., and Ermon, S. Discrete diffusion mod- cling by estimating the ratios of the data distribution. International Conference on Machine Learning, 2024.
Kingma, D. P. and Welling, M. Auto-encoding variational Bayes. In International Conference on Learning Repre- sentations, 2014.
Mazuz, E., Shtar, G., Shapira, B., and Rokach, L. Molecule generation using transformers and policy gradient rein- forcement learning. Scientific Reports, 13(1):8799, 2023,
Kong, L.. Cui, J.. Sun, H., Zhuang, Y., Prakash, B. A., and Zhang, C. Autoregressive diffusion model for graph gen- eration. In International Conference on Machine Learn- ing. pp. 17391-17408. PMLR, 2023.
Ou, J., Nie, S., Xue, K., Zhu, F,, Sun, J., Li, Z., and Li, Your absorbing discrete diffusion secretly models the conditional distributions of clean data. arXiv preprint arXiv:2406.037 2024.
Kool, W., Hoof, V., and Welling, M. Buy 4 REIN- FORCE samples, get baseline for free! In DeepRL- StructPred@ICLR 2019a.
Plackett, R. Royal Statistical Society. Series C (Applied Statistics), 24 12):193-202, 1975.
The analysis of permutations. Journal of the
Kool, W., van Hoof, H., and Welling, M. Stochastic Beams and Where To Find Them: The Gumbel-Top-k Trick for
Preuer, K., Renz, P., Unterthiner, T., Hochreiter, S., and Klambauer, G. Fréchet chemnet distance: metric for
10

---

Searning-Ordey Antoregressive Models
generative models for molecules in drug discovery. Jour- nal of Chemical Information and Modeling, 58(9):1736- 1741, 2018.
Yang, Z., Dai, Z., Yang, Y., Carbonell, J., Salakhutdinov, R. R., and Le, Q. V. Xinet: Generalized autoregressive pretraining for language understanding. In Wallach, H., Larochelle, H., Beygelzimer, A., d' d'Alché-Buce Buc, F., Fox, E., and Garnett, R. (eds.), Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
(8():1nng.8(()117333
Ronneberger; O., Fischer, P., and Brox, T. U-net: Convolu- tional networks for biomedical image segmentation. In Navab, N., Hornegger, J., Wells, W. M., and Frangi, A. F. (eds.), Medical Image Computing and Computer-Assisted Intervention -
-ICCAI 2015, pp. 234-241 Cham, 2015. Springer International Publishing.
You, J., Ying, R., Ren, X., Hamilton, W., and Leskovec, J. Graphrnn: Generating realistic graphs with deep auto- regressive models. In International Conference on Ma- chine Learning, pp. 5708-5717. PMLR, 2018.
Sahoo, S., Arriola, M., Schiff, Y., Gokaslan, A., Marro- quin, E., Chiu, J. T., Rush, A., and Kuleshov, V. Simple and effective masked diffusion language models. arXiv preprint arXiv:2406.07524, 2024.
Salimans, T. and Knowles, D. A. On using control vari- ates with stochastic approximation for variational Bayes and its connection to stochastic linear regression. arXiv preprint arXiv:1401.1022 2014.
Shi, J., Zhou, Y.. Hwang. J., Titsias, K. M., and Mackey, L. Gradient estimation with discrete Stein operators. In Advances in Neural Information Processing Systems, 2022.
Shi, J., Han, K., Wang, Z., Doucet, A.. and Titsias, M. K. Simplified and generalized masked diffusion for discrete data. In Advances in Neural Information Processing Systems, 2024.
Uria, B., Murray, I., and Larochelle, H. deep and tractable density estimator. In Xing, E. and Jebara, T. (eds.), Pro- ceedings of the st International Conference on Machine Learning, volume 32 of Proceedings of Machine Learn- ing Research, pp. 467-475, Bejing, China, 22-24 Jun 2014. PMLR.
Vignac, C., Krawczuk, I., Siraudin, A., Wang, B., Cevher, V., and Frossard, P. Digress: Discrete denoising diffusion for graph generation. In International Conference Learning Representations, 2023.
Welleck, S., Brantley, K., lii, H. D., and Cho, K. Non-monotone sequential text generation. In Chaud- huri, K. and Salakhutdinov, R. (eds.), Proceedings of the 36th International Conference on Machine Learn- ing, volume 97 of Proceedings of Machine Learning Research, pp. 6716-6726. PMLR, 09-15 Jun 2019. URL https: pestts::/proceeeee /ttp///roccennssseeee mlr. ptt.spp/e///9///// wel leck1 9a. html.
Xu, K., Hu, W., Leskovec, J., and Jegelka, S. How powerful are graph neural networks? In International Conference on Learning Representations, 2019.

---

Learning-Oree Ordder Antoregressive Models
A. A Toy Example of Generating MNIST with LO-ARM
1bpp1
Icep300
KTRp
sTapy
step:
0002
00003
0000
3330
CA02
00003
0.0093
2310
5.022
GOGS
S
Figure An example data generation MNIST using LO-ARM with shared torso parametrization the order-policy distribution. Each column visualizes step of the sampling process (steps shown are i= 10, 20, 100, 300, 600, 760). The panels in the first row show the order-policy values ps (2x #e/z ( at step depicted as an 28 x 28 image where dark colour indicates small probability value and light colour higher value. The plot in the second current value for 2) depicted as binary masks. Third row shows the progression of generating the actual digit. The figure highlights that the learned order-policy prioritizes unmasking pixels with higher certainty. The moundary-to-cente infilling strategy, shown in both the prioritization via the order-policy (first row) and the unmasking sequence (second row), illustrates this. The rough generation order horders, background, then digit indicates clear preference for of increasing certainty. Initially, the model confidently fills in the black boundaries and background (third row), leading to the gradual, confident clarification of the digit into "5".
shows all sampled pixel locations (immediatcly after having determined the
In addition main quantitative evaluation with molecular graph data, also provide illustrative example of generating new MNIST sample (Figure 3). Firstly, cast image generation generative modelling with discrete tokens, we treat the 256 gray-scale values as discrete categories and augment the token space with one additional mask token. Then generating new MNIST sample autogressively (one token at time) would be similar to generating with masked diffusion models ((Austin et al., 2021; Hoogeboom et al., 2022; Shi et al., 2024; Sahoo et al., 2024)), and the main difference LO-ARM only generates one token at a time, while masked diffusion models could generate multiple tokens one denoising step. a result, the total number sampling steps equals the total number of pixels in one MINST image.
Next, we train the model with Algorithm 1. Specifically, parameterize the categorical classifier with UNet (Ronneberger al., 2015), parametrize the both the variational and order-policy with shared-torso (see Section 3.3 and Appendix F).
Finally, we illustrate the learning order ARM (LO-ARM) algorithm in MNIST and we visualize probabilistic orders when sampling from the model, using the optimized model we apply Algorithm to generate new digits, where we inspect also the autoregressive order the pixels are generated. Figure illustrates different steps of the sampling process; see figure caption for details. We can observe that LO-ARM, which parametrizes the posterior order-policy with shared-torso, prefers sample first the border pixel values and then the "digit" pixels in the centered foreground. Presumably, the reason is that the model learns (through ELBO maximization) prefer order where "easier" border pixel locations tend be specified first. Note that there is inductive bias towards orders that prioritise first more confident data dimensions.
B. Gallery of Generated Molecules
B.I. Individual Molecules
We show the QM9 (Figure 4) and ZINC250K (Figure 5) samples generated with our best LO-ARM models presented in Table and 2.
12

---

Learning-Oree -rrii Antoregressive Models
HO
Figure 4. Generated QM9 samples.
Figure 5, Generated ZINC250K samples.
B.2. How LO-ARM Generates A Molecule
We provide the full generation path for the molecule presented in Figure 1. In particular, only show the full path of building skeleton (Stage in Figure 6) and infilling atoms (Stage 3 in Figure 7), Stage 2 only contains repetitive visible molecules.
C. Consistency Analysis of Learned Ordering for Molecular Graph Generation
As we can see the case study in Figure and Appendix B.1, LO-ARM prefers a simple and specific ordering to generate new samples, which unfolds into three stages: 1) building the skeleton molecules, infilling the "imaginery" bonds in the adjacency matrix, and 3) infilling atoms. With this observation, problem of template matching. Specifically, the template matching unfolds into the following steps:
actually simplify analyzing its consistency to
Firstly, when generating a sample, we label the token generated at each step as one of the three states, i.c., for node/atom, E for edge/bond and for no edge or imaginary bond, and output the token state sequences.
Secondly, for each state sequence, we compress only keep the transitions between different states. For example, state sequence EEEANNNAAA will be compressed EANA, and the expected ordering sequence is compressed to ENA.
Finally, we count the exact matchings between the compressed state sequences and the template of expected ordering sequence, i.e., ENA.
To test the consistency statistically, firstly, we run the experiments to train LO-ARM with three different seeds. Secondly, with each of the three models, we generate 10000 molecules and store their corresponding token state sequences. Note that,
13

---

Learning-Oree Orrii Antoregressive Models
Figure full sample path building the skeleton of molecule presented Figure 1. Masked nodes/atoms are labelled with x, and the bonds in the adjacency unmasked one by one.
we do not filter invalid molecules in this analysis. Finally, for these 30000 state sequences, count the matchings with the expected ordering through the algorithm described above. We find that 99.9% (29973 out of 30000) of the orderings match the ordering pattern that we have observed.
D. Detailed Results
We provide detailed ablation analysis for the architecture options described in Section 3.3 and Appendix F for both QM9 Table 4 and Table 5. In addition, for the experiments with ZINC250k we also provide the results of varying the threshold p in the Top-p (Shi al., 2024) sampler aforementioned in Section 5.
E. Experiment Details for Molecular Graph Generation
We provide the details our experimentation with the QM9 and ZINC250K datasets. The statistics of these two datasets are summarized in Table 6.
14

---

Learning-Onde Antoregressive Models
Figure The full sample path of infilling atoms molecule presented in Figure 1. We also include the last two steps Stage 2, i.c., infilling the no_edge states at the beginning of this figure. During the generation process, atoms unmasked one hy one.
E.1. Data Representation
We follow the standard setup (Vignac et al., 2023; Eijkelboom et al., 2024) transform data represented as SMILES (Simplified Molecular Input Line Entry System) strings dense graphs, G (H,, He, Hm), where H,, is the node vector, H. is the edge adjacency matrix and Hm is the node mask. In particular, the node masks pad molecules with variable sizes to static size, such that they can be processed with mini-batches. Moreover, same as DiGress, to represent the edge connections with a dense adjacency matrix, an "imaginery" bond is introduced, which we also call no eddge. This imaginary bond has no chemical effect when building molecules but helps us distinguish from the entries masked different masks. Now we introduce the masks used in our implementation.
As discussed in Section 2, model autoregressive generation unmasking process. To do this, augment the token vocabulary with one special token to represent the mask, which we call sampling mask. Specifically, for molecular graph generation, we augment vocabularies of both nodes and edges with one zero-valued token respectively. Note that, although both these tokens are zeros in inputs, they may have different latent embeddings as the graph transformer processes nodes and edges differently, as introduced in the next section. Notation-wise in this paper, we call the node mask H,, attention mask to distinguish from the sampling mask.
15

---

Learning-Orde Antoregressive Models
Table Detailed results on on QM9. The negative loglikelihoods (NI.ss) of methods are evaluated against the test set, and other metrics are calculated with the generated samples with the corresponding methods.
Method
Pe
qo
NLI.
Validity %
Uniqueness %
FCD4
Uniform-ARM
24.56
0.12
98.81 0.09
99.08 : 0.03
0.691 0.018
LO-ARM-ent-ss
entropy
shared torso
< 24.05
+ 0.29
99.04 + 0.15
0
0.645 + 0.036
40-ARM-ent-sep
entropy shared torso
separate
< 23.77 +:00+-00 0.11 0.11
98.93 + 0.14
99.03 0.15
0.444 + 0.013
LO-ARM-st
shared torso
< 22.23
+ 0.01
99.03
+ 0.08
98.49
0.465 + 0.027
LO-ARM-st-sep
shared torso
separate
< 21.22
0.47
99.77 + 0.08
98.71 + 0.10
0.264 + 0.081
Table 5. Molecule generation on ZINC250K, In particular, the NLLs of the LO-ARMs with the sampling are omitted as the sampler is only used in inference time and does not change the training- metries.
Method
Top-p
NLL!
Validity %1
Uniqueness %+
FCD4
Uniform-ARM
1.0
S 80.08 + 0.44
31.40 1 1.19
100.00 1 0.00
6.587 + 0.070
Biased-Uniform-ARM
1.0
78.16 + 0.27
33.68 + 0.35
100.00 + 0.00
5.075 1.075:10000 0.051
Biased-Unioom--AMM--n ARM-ttpp
0.9
79.11 + 1.03
100.00 + 0.00
22.492 0.227
LO-ARM-ent-s
1.0
77.92 0.11
4
100.00 0.00
4
LO-ARM-ent sepp
1.0
S 77.04 + 0,22
41.50 0.99
100.00 10.00:000 0.00
4.770
LO-ARM-st-SS
1.0
S 69.92 + 0.67
90.58 0.35
100.00
4.122 = 0.039
L.O-ARM-st--ttoopp
0.9
96.36 +
100.00
9.388 0.100
LO-ARMM-tt-tttopp
0.75
0.44
9
16.416
1.0-ARM-st-xep
1.0
< 68.91
+ 0.45
95.97
99.99
3.33 0.130
LO-ARMMstt-eep-on
0.9
96.16 0.48
100.00 0.00
3.586 0.118
LO-ARM -oo-s-----pp
0.3
96.56 1 0.13
100.00 0.00
3.852 1 0.035
E.2. Parameterizing Classifier, Posterior and Variational Order Policies
To ensure fair comparison with the baselines, we employ the same graph transformer network which is used in (Vignac et al., 2023; Eijkelboom et al., 2024). Just as done in DiGress and CatFlow, our graph transformer takes as masked graph G = (H,, He, Hm) and predicts a distribution over the clean graphs, represented as G = (H,, H.). Again, the attention mask H., stays constant during whole computation. For our classifier, we implement the shared-torso parameterization for the posterior order policy po (-((x2e)) through augmenting the output layers of node and edge logits with one extra dimension respectively, the output of which are used the logits of the posterior order policy, and each logit corresponds the position of node edge. Then during sampling, we obtain the logits of the posterior order policy through concatenating the extra node logits and the flattened extra edge logits.
For the variational order-policy, its shared-torso parameterization is the same as the posterior order-policy. Moreover, when implementing it with a separate neural network, we employ a smaller graph transformer with the same architecture. This is because the variational order-policy also takes graphs as inputs, i.e., go a = G), and we want to leverage the nice cquivariance properties (Vignac al., 2023) brought by the graph transformer to improve sample efficiency. Moreover, intuitively, in addition to homogeneous inputs, the architectural homogeneity would also help optimize the KL term between the variational and posterior order policies. Finally, the output dimension the ge network is only one for both node and edge logits. Same as the Pe network, we obtain the output logits after flattening and concatenating them. Moreover, to enforce symmetry of the adjacency matrix,
Note that, there may just be our design considerations specific to the task molecular graph generation. In principle, we pose no constraints on designing the three components. There could be more efficient parameterization choices, and we will explore as part of the follow-up work.
E.3. Enforcing Symmetry in Masking and Unmasking
Recall that during training 1, after sampling an latent ordering z (zęi, Z2z), where zςi a partial ordering, and agi the unordered set, and we need to mask the tokens dimensions 2z4. To enforce symmetry of the adjacency matrix A. for edges, we only process the dimensions in the order-policy corresponding to the upper half of Hm. with the lower half being masked out according to the attention mask. After zeroing out ø21, we flip the upper half the lower half. We
16

---

Learning-Order Antoregressive Models
Table 6. Statisties of QM9 and ZINK250K datasets used in molccule generation tasks.
Dataset
Number of molecules
Number of nodes
Number of node types
Number of edge types
QM9
133, 885
V
4
4
ZINC250k
249, 455
S V
38
4
Figure 8. Illustration of the process of masking molecular graphs. A graph is represented as a node vector and an dense adjacency matrix which is symmetric with respect the diagonal. The upper row indicates the active dimensions (green blocks) in the order-policy and the lower row shows the corresponding masked graph. When an edge dimension is sampled (yellow block). to ensure symmetry. we mask both the corresponding data dimensions (blue blocks) the graph and its symmetric dimension. Note that sampling an ordering from the order- order-policy is sampling- sampling-without wethout-replacone problem, and once dimension is sampled in the order-policy. we mask it out with the sampling mask (red blocks).
Policy over token positions
node sampled
An edge sampled
Fully masked policy
Node vector
Adjacency matrix
Input graph
Masking the node
Masking the edges
Fully masked
Active policy dimension
Allention mask Sampled dimeension
Active data dimension
Sampling mask
always perform this operation regardless whether an edge is sampled not. Note that, except for the transformer outputting nodes and edges separately, we provide no inductive bias about edge and node positions to the order-policy. We provide an illustration of the masking process in Figure 8.
We apply similar process during sampling Algorithm as shown Figure 9. particular, instead of starting with a fully masked graph with full adjacency matrix, we only consider the upper half of the adjacency matrix, and flipping the upper half after one sampling step. In this way, we will always obtain a symmetric adjacency matrix regardless whether an edge unmasked.
E.4. Experimental Setup
We report the hyperparameters in Table 7. Moreover, the hidden dimensions for the classifier network are kept the same (Vignac et al., 2023) and all data is kept the same in (Jo et al., 2022). All experiments until convergence.
F. Options of Model Parameterization and further details about REINFORCE LOO
We provide the details of different choices of parameterizing the distributions needed in Algorithm I and Algorithm As mentioned in the main paper, we assume that data are discrete and each dimension (token) x; takes m categorical values.
17

---

Learning-Orden Antoregressive Models
Figure 9. Illustration of the generation process unmasking for molccular graphs. Generating new molecular graph reserves the masking process. The only difference is that for the molecule graph, we start with the upper half of the adjacency matrix, and each time an edge dimension is sampled with the polerr we unmask the corresponding data dimension the adjacency matrix and symmetric dimension.
Policy over token positions
An edge is sampled
node is sampled
Fully masked policy
vector
Adjacency matrix
Molecule to be generated
Unmasking the edges
Unmasking the node
Fully unmaskedigenneatted
Aclive policy dimenssoo
Allention mask
Sampled dimension
Active data dimension
Sampling mask
Table 7. Hyperparameter setup.
Hyperparameter
QM9
ZINC250K
Optimizer
AdamW
AdamW
Scheduler
Cosine Annealing
Cosine Annealing
Learning Rate
1 10-5
1.5 10
Weight Decay
1 10
1 10
EMA
0.9999
0.9999
We represent each 2; as 1-dimensional (rather than m-dimensional) one-hot vector where the final m + 1-th value is special one since it indicates that the category of is currently masked or unspecified.
We assume a Neural Network (NN) that takes L. inputs, i.e., L m + 1-dimensional one-hot vectors of the discrete tokens, and gives in the output I vectors of classifier logits, each corresponding to the m categorical values (excluding the m + 1-th masked value) for each token v,. More precisely, we denote the NN output logit vectors by
(0 E R.m k 1, Lt,
so that each fe. consists of the logits of the zı-th output classifier for = k. The NN takes as input a L X (m + 1) matrix &zc. which in all indices zςi contains the corresponding observed (or generated) data values E that precede token 2,, and the remaining indices 2zj is filled in with the mask token value, i.e., with the one-hot vector that has the value the final m 1-th dimension. Another way to view this is that & is modification of a that has all z tokens of æ unmasked and the remaining ø2; tokens of æ masked. Then, based on these NN outputs the probability distribution that specifies the value for xz, is categorical distribution or classifier of the form
Po (z (x22.) softmax(fo.s. (ff (Zzz.)),
where the logits are passed through softmax to provide a probability vector. The other quantities needed to be specified are the model order-policy factors po(2; k k//zi, Fzey) and the corresponding variational posterior-order distribution factor
18

---

Learning- Autoregressive Models
(2 k|zzú, R). As mentioned in the main text, the difference between these two distributions is that the first exploit information only from the previous specified dimensions Zzzc, while the second condition the full We specify these distributions as described below.
E.1. Form of Order-Policy Distribution
We explore two choices for parametrizing each order-policy, factor po(2: classifiers pe 22 (xx2))) ), defined above, and form pa(2; k|zci, Txcc) using entropy-base uncertainties. In the second option we model pe (zg Lzzi) by adding an extra final linear layer to the NN feature representation. This latter approach requires no modification or extension of the NN architecture. We detail both cases below.
( Ric Fzgz). In the first option re-use the NN
Entropy- parametrization. the first option we incorporate some inductive bias towards more confident less confident tokens under the model. More precisely, we entropy-baae -baed uncertainty logits so that each factor is written as
pa(z; = k| Fzci) Σ Eeezz e -HH(po(zz (
33M(soltmma fo (2
Evesge
HM( fo (ec))y)))
where is the entropy distribution and B E IR is scalar parameter. For 3 this becomes the uniform distribution the L 7- i values in Zzg- For B 0 the distribution favors the selection of z values for which the model data conditionals (3rz, (azes) are more certain about the actual categorical value Fz, should assign, while similarly for B < O the preference is reversed towards more uncertain tokens. 3 is treated as additional model parameter that is optimized jointly with 9 using the variational inference method described in Section 3.2.
Shared-torso parametrization. For second option of the order-policy, referred to as shared-torso, we add an extra final linear layer to the NN with L scalars h-) (-) IR, k I....,,. L. Then, each order-policy factor obtained by
(o
pa(z; Alzciitzee)
EMErs, (aar(Bse))
where the extra model parameters needed to define these L scalar outputs optimized jointly with the remaining parameters of the NN architecture.
F.2. Variational Distribution
As explained in the main text, model each variational factor go(z: k|zcc Dx.:))
(00,)(0))
qo (2, ke 21 x) Σ Evers, (20
(15)
so that we only need to parametrize the vector function
(x)))= E R.L.
The first option we consider is adding an extra final layer with L outputs another head the main neural architecture (this is the shared-torso option). In the second option, and somehow more flexible, we construct separate NN where in the final layer it outputs L real values to model the vector ga (x:)
Computing and sampling from the variational distribution is very fast since it requires single forward pass with input æ from the NN (either the shared architecture or the separate network) to obtain and store the vector of logits go(x). Then sampling a full path z can be done at once in parallel using the Gumbel-top-k trick (Kool et al., 2019b). Note that sampling two paths 2 z2 from the variational distribution is needed in order to compute the objective in Equation (16); see next.
19

---

Learning-Order Antoregressive Models
F.3. Objective to with Automatic Differentiation
As discussed in the main text, training done using an unbiased REINFORCE leave-one-out (RLOO) gradient computed separately for each data point in the minibatch (then all these gradients averaged over the minibatch). For single data point this unbiased gradient is given by Equation (11) in Section 3.2.
For implementation convenience there is way to obtain the unbiased gradient by programming a certain objective function and then apply automatic differentiation to it. Specifically, what is required is to implement the following objective (to be maximized):
{ (log (=<2 log go (=22 z)) stopgrad [F] + Fo x) + Fo(22..)
(16)
2
where stopgrad [F] stops the gradient computation the AF. However, note that this objective is just a trick to easily obtain the gradient, while the actual objective that we maximize is the ELBO. This means that monitoring convergence done by computing the following stochastic ELBO (per data point):
L
(()) (Fo x) + (0)
(17)
which is just the two-sample version of the one-sample stochastic ELBO from Equation (9). Given a training minibatch these stochastic ELBOs are further averaged over the minibatch.
20


Reviews:
[{"review_id": "NrE70jhfeH", "reviewer": "Reviewer_H7W8", "summary": "This paper propose a new generative modeling framwork named  Learning-Order Autoregressive Models. The core of this framework is to extend the traditional ARs to learning a dynamic order of sampling. Speicifically train a order polocy to determine the order. To train such model, the authers use a  variational lower bound on the exact log-likelihood , and optimize via stochastic gradient estimation. It achieved quite good results in  QM9 and Zink250K.\n\n## Update after rebuttal\nI appreciate the authors' thorough response to the concerns raised. I have no further questions and will maintain my original score as positive unchanged. Thank you for your efforts.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the time you’ve taken to review our work and for the positive and constructive feedback! We are glad that you found the problem we dealt with important and our approach novel and well-grounded. In response to the weaknesses and questions:\n\n### 1. Generalization to other high-dimensional data (images, graphs)\nWe would like to clarify that the focus of this work is on designing the order learning mechanism in ARMs and its application to molecular graph generation as the main testbed. We chose this task because we believe it is more suitable – As the reviewer has pointed out, an important issue in graph autoregressive generation is the difficulty to determine an order for sampling. Indeed, our learning-order ARM was able to discover an autoregressive order that outperforms fixed or random order used in prior graph generation work. Our focus is not on images – we included MNIST only as a sanity check to confirm the model can learn a meaningful order that distinguishes between the digits and the background. While there are no practical constraints for applying LO-ARMs to higher dimensional natural images – it is unclear whether there is a similarly meaningful and interpretable ordering to learn there and therefore we decided not to pursue it in this work.\n\nTo demonstrate the scalability of our algorithm to higher dimensional graph datasets, we conduct an additional experiment on the GuacaMol dataset, which is a larger molecule dataset with 3906 input dimensions (which is larger than 1482 of ZINC250k and 90 of QM9). We have provided the comparison of the three molecule datasets in the discussion with Reviewer Pa4s. We believe that these three datasets can provide a comprehensive evaluation matrix to support our claims.\n\nPerformance-wise, our preliminary results show that the order-policy still yields an advantage, exceeding or matching the SOTA performance, as you can see from the details we have provided in the discussion thread with Reviewer Pa4s.\n\n### 2. Application to 3d point cloud, pixels, voxels\nThank you for the suggestion. We are interested in exploring these modalities in future work.\n\n### 3. Pipeline illustration figure\nWe will follow the reviewer’s suggestion to add an illustration figure in the final version.\n\nWe hope that our responses have addressed your concerns and questions. If so, we would kindly ask the reviewer to reconsider the decision in light of our responses and update their score accordingly. We are also eager to address any additional concerns the reviewer may have.", "claims_and_evidence": "I agree that the authors are dealing with a important issue in  graph autoregressive generation as it is hard to determine a order for sampling. I think use a order-policy to sample the order seems reasonable. The claim of LO-ARMs Learn a Meaningful and Consistent Order for Generation is acceptable.\nWhile LO-ARMs Can Generalize to Other High-Dimensional Data (Images, Graphs) shoule be given more evidence. The MNIST experiment is more of a toy example rather than rigorous proof that the method generalizes to all high-dimensional data, and quantitative metrics (e.g., likelihood scores, FID, etc.) should be  provided for image generation, making it easy to measure actual performance compared to baselines.", "methods_and_evaluation_criteria": "Overall the the method proposed is suitable for the problem . The HyperEdge enhanced EGNN and the docking scorefeatures are reasonable. The evaluations are on the commonly used unconditional molecule graph generation.", "theoretical_claims": "There are no theoretical claims in this paper", "experimental_designs_or_analyses": "The experiment settings are from the existing related work, and the authors use the same settings, which are reasonable", "relation_to_broader_scientific_literature": "Appling the order policy to other generation modalitys like 3d point cloud, pixels, voxels maybe interesting.", "essential_references_not_discussed": "Related works are discussed properly.", "other_strengths_and_weaknesses": "Strengths\nThe paper presents a novel approach to learning orderings in autoregressive models, which is a non-trivial extension of traditional ARMs. The paper provides a well-structured training algorithm and a clear sampling procedure, making the method well-grounded\nWeaknesses\nWhile LO-ARMs are tested on MNIST as a toy example,  i think this is not sufficient to claim general applicability to high-dimensional data like images, more examples of benchmark in image generation would be better.\nWould be better if a pipeline illstration figure is provided. Currently only a case study figure shown in the paper main part.", "other_comments_or_suggestions": "No", "questions_for_authors": "See the Strengths And Weaknesses part", "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": "", "extra_scores": {}}, {"review_id": "pRu2lDK7iS", "reviewer": "Reviewer_Stjb", "summary": "This paper proposes a method to learn an optimal generation order for autoregressive models in data domains that do not possess a natural canonical ordering (e.g., graphs or images). By framing the ordering itself as a latent variable with a dynamic, learnable distribution (the “order-policy”), the authors unify the ideas of AO-ARMs with variational inference. They demonstrate that a variational bound on the log-likelihood can be optimized by sampling permutations from a “posterior” distribution and matching them to the model’s own “order-policy”. Empirically, the method achieves competitive results on two molecular generation datasets, QM9 and ZINC250k. The paper’s thorough ablations highlight design trade-offs for parameterizing the order-policy, and the final approach shows consistent improvements over both uniform and biased AO-ARMs.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the positive feedback! We are glad that you find our theoretical results sound, our improvement consistent and our ablations thorough. We address your concerns as below, especially your concern on training stability with REINFORCE.\n\n### 1. Performance gains on QM9 and ZINC250k\nWe believe our improvements on QM9 and ZINC250k are substantial. Although baselines score well on validity/uniqueness, our results indicated that there is still large room for improvement in distributional similarity metrics such as FCD — highlighted by Reviewer Stjb as a recognized measure for drug-like molecules. Our LO-ARMs model achieved new state-of-the-art FCD scores, reducing them to 0.240 (QM9) and 3.229 (ZINC250k) from prior bests of 0.441 and 13.21, respectively.\n\n### 2. Smaller and larger scale benchmarks (e.g., FreeSolv, Lipophilicity, HOPV)\nThank you for your suggestion, and we will certainly include the results on smaller benchmarks in the final version of our work. Given the time constraints of this rebuttal period, we made a strategic decision to prioritize the GuacaMol experiment, which is a much larger dataset than ZINC250k, to maximize the efficiency of our response. This is because we anticipate that training stability might pose a greater challenge for larger datasets compared to smaller ones, as their training processes tend to exhibit more stochasticity. Furthermore, the scalability of LO-ARM is a concern shared by several reviewers. The detailed results on GuacaMol dataset are presented in the response to Reviewer Pa4s. In short, our preliminary results show that LO-ARM learns an order-first generation ordering [1] and still exceeds or matches performance of the current state-of-the-art-models.\n\n### 3. Analysis of variance and convergence of REINFORCE-based training\nTo provide a clearer understanding of the variance and convergence of REINFORCE and the effectiveness of our variance reduction method, we have visualized the following quantities during training course of the GuacaMol experiments [2]:\n\n* Negative Evidence Lower Bound (ELBO)\n* Maximum and minimum q-logits from the variational order policy network.\n\nSpecifically, we conducted an ablation study on the learning rate in two experiments, keeping all other settings constant. Given that the gradient variance primarily arises from the stochasticity of the variational order policy, the maximum and minimum q-logits can reflect this variance throughout the training.\n\nAs seen in the plot, most of the time the RLOO variance reduction keeps the optimization smooth -- still occasionally training instability is observed with a large learning rate (2e-5), indicated by loss spikes and discontinuities in the maximum q-logits. Fortunately, reducing the learning rate to 1.5e-5 effectively stabilizes training, as demonstrated by the green curves. \n\nWe hope that our responses have addressed your concerns and questions. If so, we respectfully ask that you reconsider your decision based on our responses and update your recommendation score accordingly. We are also eager to address any further concerns you may have.\n\n[1] Generation trajectory of GuacaMol sample: https://drive.google.com/file/d/1a5HU4FZ98bqS_9JFK4Eb60QfhXiksaue/view?usp=drive_link\n\n[2] Variance visualization: https://drive.google.com/file/d/19b2wQocvLJc82bAm6eUG5qbAcmVAra0V/view?usp=drive_link", "claims_and_evidence": "The main claims of the paper are:\n1. By introducing a trainable policy over generation order, the autoregressive mode; avoids the limitations of either a single fixed order or a uniformly random permutation of dimensions.\n\n2. The authors propose a variational approach, using REINFORCE with a leave-one-out baseline to reduce gradient variance. They argue that this is practical enough to handle the large combinatorial space of permutations.\n\nThese claims are generally supported by the quantitative comparisons. \n\nMy concerns are 1. the convergence of REINFORCE-based training is not thoroughly studied in the analysis session; 2. the performance gains on the two datasets (QM9 and ZINC250k) are very modest as the baselines have already achieved very high results.", "methods_and_evaluation_criteria": "The authors follow common practice in generative modeling by: evaluating sample quality on the standard Validity/Uniqueness metrics for molecules, and measuring distributional similarity with FCD, a recognized measure for drug-like molecules.\n\nAlthough the results are strong, the paper focuses primarily on mid-scale tasks (QM9, ZINC). It would be beneficial to demonstrate that the approach still offers consistent improvements on small-scale benchmarks [1,2,3], as the baseline results on these benchmarks are less optimal compared to QM9 and ZINC.\n\n[1] FreeSolv: a database of experimental and calculated hydration free energies.\n[2] In silico evaluation of logD7.4 and comparison with other prediction methods.\n[3] The Harvard organic photovoltaic dataset.", "theoretical_claims": "The derivation of the method as a variational ARM with a learned ordering distribution is a clean extension to the standard AO-ARM framework. The theoretical foundation seems sound: the ordering z serves a hidden latent variable for the distribution of data x, so the standard variantional inference can play a role here for the optimization of joint probability distribution.", "experimental_designs_or_analyses": "The experiments on molecular generation are thorough, with ablations revealing the roles of shared vs. separate neural networks for the order-policy and the choice of top-p sampling strategies. The authors also visualize generation paths for example molecules, giving a nuanced look at how the dynamic ordering policy emerges in practice.\n\nHowever, concerns remain about REINFORCE’s notorious variance. While the authors do implement a leave-one-out method to reduce variance, readers may question the stability of training at larger scales or on smaller datasets with fewer training points (like FreeSolv or Lipophilicity). An in-depth discussion or demonstration of how many epochs/hyperparameter adjustments were needed would be welcome.", "relation_to_broader_scientific_literature": "The paper situates its method well with respect to AO-ARMs, masked discrete diffusion, and existing graph-generation methods. It clarifies that past approaches either fixed or uniformly randomized the ordering.", "essential_references_not_discussed": "None.", "other_strengths_and_weaknesses": "None.", "other_comments_or_suggestions": "None.", "questions_for_authors": "1. Have you measured or visualized the gradient variance over the course of training? Are there scenarios (like larger, more complex data) where you suspect REINFORCE might fail to converge easily?\n\n2. Do you have preliminary results on smaller tasks (e.g., FreeSolv, Lipophilicity, HOPV) that confirm the order-policy still yields an advantage?", "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": "", "extra_scores": {}}, {"review_id": "1wYDerqoTI", "reviewer": "Reviewer_Pa4s", "summary": "This paper addresses a fundamental limitation in Autoregressive Models (ARMs)—the assumption of a fixed generation order, which may not be optimal for complex data types like graphs. The authors introduce Learning-Order Autoregressive Models (LO-ARMs), a novel approach where the model learns a context-dependent order for sequential data generation rather than relying on a predefined or random order. This paper presents a major advancement in autoregressive modeling, showing that learning an optimal generation order improves sample efficiency and generation quality.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate your positive feedback on the quality of our work. We now address each of your concerns as below.\n\n## 1. Preliminary results on larger bioactive molecule dataset (ChEMBL)\n\nTo address your concern regarding the scalability of LO-ARM to larger datasets, we include a preliminary result on the ChEMBL dataset (also known as the GuacaMol benchmark [1]) without hyperparameter tuning. Specifically, we preprocessed this dataset using the utility provided in [2], a method consistently employed in our benchmarks [2, 3]. We have summarized the statistics of the three datasets used in our evaluation—namely QM9, ZINC250k, and ChEMBL/GuacaMol—as presented below. We believe this inclusion provides a more comprehensive assessment of LO-ARM's capabilities.\n\n\n\n|          | #samples | #nodes           | #node types | Input dims |\n|----------|----------|------------------|-------------|------------|\n| QM9      | 130k     | 1 <= \\|V\\| <= 9  | 4           | 90        |\n| ZINC250k | 250k     | 6 <= \\|V\\| <= 38 | 9           | 1482        |\n| GuacaMol | 1.2M     | 2 <= \\|V\\| <= 62 | 12          | 3906       |\n\n\nOur preliminary results show that LO-ARM still exceeds or matches performance of the current state-of-the-art-models.\n\n|               | Input data | Validity | Uniqueness | Novelty   | Raw FCD  |\n|----------|--------|--------|---------|---------|--------|\n| LSTM          | SMILES     | **95.9** | **100.0**  | 91.2      | **0.46** |\n| MCTS          | Graph      | 92.9     | 95.5       | 100.0     | 21.00    |\n| DiGress       | Graph      | 85.2     | 100.0      | 99.9      | 1.92     |\n| LO-ARM (ours) | Graph      | *94.2*   | **100.0**  | **100.0** | 3.73     |\n\nTo ensure consistency in our evaluation across the QM9 and ZINC250k datasets (as is common in the literature), we have converted the reported FCD scores from prior work, which are often presented on an exponential scale, back to their raw FCD values.\n\nSpecifically, the LO-ARM model trained on the GuacaMol dataset exhibits a preference for an atom-first generation order. This means it tends to generate atoms first, followed by the real bonds connecting them, and lastly, it fills in the non-existent or \"imaginary\" bonds. We illustrate an example of this generation process in the accompanying figure [4]. This learned ordering contrasts with the edge-first ordering observed in models trained on the ZINC250k and QM9 datasets. This difference likely arises because generating node-related dimensions first is potentially simpler for the GuacaMol samples, given that the number of edge dimensions increases quadratically with the number of nodes.\n\nOur preliminary findings indicate that LO-ARM surpasses other graph-based methods in terms of validity and uniqueness, achieving performance levels close to the state-of-the-art SMILES-based approaches. However, we have observed a small gap in FCD compared to DiGress. We hypothesize that this difference arises because for datasets with higher dimensionality, a generation strategy operating at a coarser granularity (e.g., dimension blocks) than individual dimensions might be more effective. For example, tokenizing molecules into fragments, as discussed in our Discussion section, could potentially improve performance. We defer the task of integrating domain-specific tokenization into LO-ARM for future investigation.\n\n## 2. Some aspects of the methodology and evaluation could be further discussed.\nIn addition to these new results of GuacaMol experiments, we have provided an in-depth analysis on training stability with REINFORCE in the thread of Reviewer Stjb.\n\n## 3. Improve the presentation of the main tables and enrich related work\nThank you for your thorough review and your advice, and we will refine the main tables and reference the surveys of graph generation in the final version.\n\nWe hope that our responses have addressed your primary oncerns regarding the scalability of our algorithm. If so, we respectfully request that the reviewer reconsider their decision in light of our responses and update their recommendation score accordingly. We remain available and eager to address any further concerns the reviewer may have.\n\n[1] https://arxiv.org/abs/1811.09621\n\n[2] https://arxiv.org/abs/2209.14734\n\n[3] https://arxiv.org/abs/2406.04843\n\n[4] Generation trajectory of GuacaMol sample: https://drive.google.com/file/d/1a5HU4FZ98bqS_9JFK4Eb60QfhXiksaue/view?usp=drive_link", "claims_and_evidence": "The majority of the claims in the paper are supported by well-structured evaluations.", "methods_and_evaluation_criteria": "The proposed methods and evaluation criteria mostly align well with the problem of learning dynamic generation orderings in autoregressive models, especially for molecular graph generation. However, some aspects of the methodology and evaluation could be further discussed. \n\nThe paper states, \"LO-ARM generalizes across different molecular graphs\" (p. 7), but only evaluates on QM9 and ZINC250K, both of which contain relatively simple organic molecules. The datasets such as larger bioactive molecules (ChEMBL) could be helpful to check the ability of the model on the complex molecular graphs.", "theoretical_claims": "No issues.", "experimental_designs_or_analyses": "All good", "relation_to_broader_scientific_literature": "The key contributions of this paper build upon and extend several established areas in machine learning, particularly autoregressive models (ARMs), variational inference, graph generation, and molecular generative modeling.", "essential_references_not_discussed": "No", "other_strengths_and_weaknesses": "See above", "other_comments_or_suggestions": "The bold result in table is not clear, it looks like authors are bold the best result on the other method and the on their proposed method in table 1. But in table 2, the best result on proposed method is not bolded. It would be great to keep them consistent.\n\nAlso please considering to cite the survey papers for molecule/graph generation.", "questions_for_authors": "No.", "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": "", "extra_scores": {}}, {"review_id": "65FbmUBAKs", "reviewer": "Reviewer_u8KG", "summary": "This paper introduces a method for learning the order in which discrete elements from a masked set should be generated. The authors assume the generation process begins with a set of masked elements and, in addition to predicted the unmasked element, the model learns a policy which dictates the order in which elements should be unmasked from the set. The methods uses a variational approach and show how the previously proposed REINFORCE leave-one-out estimator can be used to backprop gradients to the variational distribution network. The authors focus the application of their method on molecule generation using two commonly used datasets and also show how the method can be applied to generate MNIST images.\n\n# Update after rebuttal\n\nI thank the authors for answering my questions and providing additional experimental results. I remain unconvinced that molecule generation is the right application for this method, and I think the additional sampling steps required vs. SMILES models remains a problem. However, I think the method is interesting, well-supported and could be useful for other applications, so I will increase my score.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We extend our sincere gratitude to the reviewer for their expert insights and valuable suggestions regarding the incorporation of chemistry-specific metrics to enhance our evaluation. We address their points below.\n\n## 1. Report results against chemistry-specific metrics\nWe evaluated our best-performing models, trained on the ZINC250k dataset, using the Synthetic Accessibility Score (SAS), Quantitative Estimation of Drug-likeness (QED), diversity and novelty metrics. Diversity was measured by calculating the pairwise Tanimoto similarity within a set of generated molecules. The table below summarizes the results, including baselines and measurements on ground truth data, and more detailed visual comparisons of the metric distributions between our models and the ground truth data are available at [1]. As the results indicate, the distribution of these metrics in the samples generated by LO-ARM closely resembles that of the ground truth data.\n\n| mean | QED           | SAS           | Diversity     | Novelty     |\n|----------|----------|---------------|---------|--------|\n| Ground truth  | 0.75  | 2.76  | 0.35  | -           |\n| JT-VAE | **0.76** | 3.37 | - | 100.0 |\n| GCPN | 0.61 | 4.62 | - | 100.0\n|MolecularRNN| 0.68         | 3.59        | - | 100.0\n| LO-ARM        | **0.75**  | **3.08**  | 0.34   | 100.0 |\n\n## 2. Add baselines of SMILES-based transformers\nConsistent with our baselines (e.g., DiGress and CatFlow), we focus on graph representation of molecules in this work. For the reviewer’s interest, we included a SMILES-based baseline in our GuacaMol results (please see our reply to Reviewer Pa4s). Morever, LO-ARM is agnostic to data representations and could also be applied to SMILES strings. As suggested by the reviewer, SMILES would be more practically useful, adapting LO-ARM to SMILES generation is a promising direction for our future work.\n\n## 3. Time complexity of inference\nRegarding the cost of sampling, we first note that we can skip the edge sampling stage if the order policy determines it as a no-edge dimension. Second, the learned ordering within LO-ARM effectively separates the generation of non-existent bonds (no-edges) from atoms and real bonds. This separation allows us to generate all no-edges in a single inference step without compromising the chemical validity of the generated molecules, substantially reducing inference overhead. We present a comparison of generic and ordering-informed sampling approaches evaluated on the ZINC250k dataset below.\n\n|                   | Validity | Uniquness | FCD   | Avg. sampling steps |\n|-------------------|----------|-----------|-------|---------------------|\n| Generic sampling  | 96.26    | 100.      | 3.229 | 330.7               |\n| Ordering-informed | 96.13    | 100.      | 3.319 | 48.8                |\n\nWe plan to further explore this learned ordering to scale up inference in future work. A possible yet more principled direction could be leverating the equivalence between random order ARMs and masked diffusion models to enable parallel generation. Our preliminary experiment shows that we can obtain 90.0% validity in half of the number of steps using the masked diffusion formulation of our model.\n\n## 4. Suitability of learning ordering for molecular graph generation\nWhile generation order constraints reduce the permutation space, our results show a clear benefit to learning a data-dependent order. For example, LO-ARM significantly outperforms GraphARM [2], which unmasks a node and its edges in one step, across all metrics.\n\nTo ensure a fair comparison with our baseline models, we have only evaluated unconditional generation. To further tailor LO-ARM for molecular graph generation, one could integrate inductive biases into the backward sampling process during training, which we also recognize as a promising avenue for future development of LO-ARM.\n\n## 5. Should the equation for qtheta at the start of page 6 have a $x$ instead of $x_{z<i}$?\nYes, this is a typo, thanks for examining our work thoroughly.\n\n## 6. Is each component of z sampled autoregressively, or are all sampled at the same time?\nTo clarify: z is sampled autoregressively via the order-policy for generation, but for the variational distribution during training, all its components are generated in one pass by the q-network using the Gumbel top-k trick [3].\n\n## 7. Enrich essential references\nWe will update the Related Work section with the references you’ve suggested in the final version of the paper.\n\nWe hope that our responses have addressed your concerns and questions. If so, we would be grateful if you would reconsider your decision in light of our clarifications and update your recommendation score accordingly. We remain available and eager to address any further concerns you may have.\n\n[1] Chemistry-specific metrics: https://drive.google.com/file/d/1miSlh2vYYRfqHJ-e5QmYFmht5xMQcqEO/view?usp=sharing\n\n[2] https://arxiv.org/pdf/2307.08849 \n\n[3] https://arxiv.org/pdf/1903.06059", "claims_and_evidence": "The authors are well supported by their evaluation results, however, please see below for shortcomings with the evaluation.", "methods_and_evaluation_criteria": "- For molecule generation tasks the authors should also report the diversity (eg. often defined using tanimoto similarity of the fingerprints) and novelty of the generated molecules. This is to help ensure that the model is not sampling only from a restricted part of chemical space.\n\n- Since the validity and uniqueness metrics are fairly saturated by existing methods, and the FCD on its own is not a strong enough predictor of molecule quality, I think it would be important to report metrics which attempt to measure distances to the data distribution. For example comparing the QED, synthetic accessibility, number of aromatic rings, etc. from a set of generated molecules to the training distribution. Ultimately the goal of the generative model is to sample from the same chemical space that the training data comes from.\n\n- In order to make the evaluation more practically meaningful the authors should report baseline results for a SMILES-based transformer model (eg. Molformer, Chemformer or Molecular Transformer would all be simple starting points), using both canonical and SMILES augmented orderings. This type of model is very commonly used in practice and is a very important baseline for autoregressive molecule generation.", "theoretical_claims": "The theoretical claims of the paper are well supported and the method derivation is very clearly presented.", "experimental_designs_or_analyses": "The experiments are well designed and follow previous work in the area.", "relation_to_broader_scientific_literature": "The paper introduces a novel method for generating graph-structured data (or any other data without a specific ordering). The method is related to masked discrete diffusion and flow-matching methods but introduces a variational framework and a different training strategy. The method produces strong results on graph generation and the authors show example generation traces and provide some intuition for the decisions made by the model. However, I think some important baselines have been missed - see above for evaluation and below for additional diffusion and FM graph generation models.", "essential_references_not_discussed": "- DeFoG [1] and GruM [2] have both recently been proposed as graph generative models with strong performance on molecule generation using flow-matching and diffusion, respectively.\n\n[1] https://arxiv.org/abs/2410.04263\n\n[2] https://arxiv.org/abs/2302.03596", "other_strengths_and_weaknesses": "- Is each component of graph (ie. nodes and edges) sampled one-at-a-time? Meaning for a graph with $n$ nodes the method requires $n^2 + n$ sampling steps? This seems like a major limitation in comparison to existing methods for autoregressive molecule generation where the length of the sequence is usually only slightly larger than $n$. It would also become a significant bottleneck for scaling the LO-ARM method to larger graphs.\n\n- I think the overall method for learning the order in which items are generated is an interesting idea and a potentially useful framework for other applications, however, I am not convinced that molecule generation is the right application for this. A number of very strong baselines exist for molecule generation and would likely have significantly faster sampling times than the proposed method. Additionally, once a starting point has been defined, the generation order of molecules could be constrained relatively easily by only unmasking atoms which are connected to already unmasked atoms, significantly reducing the space of permutations and making learning the ordering less important.", "other_comments_or_suggestions": "- Should the equation for $q_{theta}$ at the start of page 6 have a $\\mathbf{x}$ instead of $\\mathbf{x_{z_{\\lt i}}}$?", "questions_for_authors": "- I am not really clear on how you achieve single step generation for $\\mathbf{z}$. Is each component of $\\mathbf{z}$ sampled autoregressively, or are all sampled at the same time?", "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": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: The paper introduces an auto-regressive model (ARM) that generates images and graphs using a probabilistic ordering inferred from data. This is achieved by using a trainable probability distribution that dynamically decides the sampling order of the data dimensions, which is optimized using stochastic gradient estimation the variation lower bound of the exact log-likelihood. Experimentally, the paper shows good results on molecular generation (QM9 and ZINC250k benchmarks).

This paper was well-received by the reviewers, in particular the clarity of exposition and the motivations behind the paper, the novelty and elegance of the approach, and the satisfactorily experimental part. The rebuttal phase seems to have addressed most concerns, and while some still remain (see the response by reviewer u8KG which points up a higher number of sampling steps with respect to e.g. SMILES generators) they seem minor.

I recommend acceptance.
Decision: Accept (poster)