Title: Learning Directed Graphical Models with Optimal Transport
Venue: NeurIPS 2023
Accepted: no
Keywords: Bayesian networks

probabilistic graphical models
Abstract: Estimating the parameters of a probabilistic directed graphical model from incomplete data remains a long-standing challenge. This is because, in the presence of latent variables, both the likelihood function and posterior distribution are intractable without further assumptions about structural dependencies or model classes. While existing learning methods are fundamentally based on likelihood maximization, here we offer a new view of the parameter learning problem through the lens of optimal transport. This perspective licenses a framework that operates on many directed graphs without making unrealistic assumptions on the posterior over the latent variables or resorting to black-box variational approximations. We develop a theoretical framework and support it with extensive empirical evidence demonstrating the flexibility and versatility of our approach. Across experiments, we show that not only can our method recover the ground-truth parameters but it also performs competitively on downstream applications, notably the non-trivial task of discrete representation learning.  

1

Abstract

Estimating the parameters of a probabilistic directed graphical model from incomplete data remains a long-standing challenge. This is because, in the presence of latent variables, both the likelihood function and posterior distribution are intractable without further assumptions about structural dependencies or model classes. While existing learning methods are fundamentally based on likelihood maximization, here we offer a new view of the parameter learning problem through the lens of optimal transport. This perspective licenses a framework that operates on many directed graphs without making unrealistic assumptions on the posterior over the latent variables or resorting to black-box variational approximations. We develop a theoretical framework and support it with extensive empirical evidence demonstrating the flexibility and versatility of our approach. Across experiments, we show that not only can our method recover the ground-truth parameters but it also performs competitively on downstream applications, notably the non-trivial task of discrete representation learning.

1 Introduction

Learning probabilistic directed graphical models (DGMs, also known as Bayesian networks) with 16 latent variables is an important ongoing challenge in machine learning and statistics. This paper 17 focuses on parameter learning, i.e., estimating the parameters of a DGM given its known structure. 18 Learning DGMs has a long history, dating back to classical indirect likelihood-maximization ap19 proaches such as expectation maximization [EM, 15]. However, despite all its success stories, EM 20 is well-known to suffer from local optima issues. More importantly, EM becomes inapplicable when 21 the posterior distribution is intractable, which arises fairly often in practice. 22

A large family of related methods based on variational inference [VI, 30, 27] have demonstrated 23 tremendous potential in this case, where the evidence lower bound (ELBO) is not only used for 24 posterior approximation but also for point estimation of the model parameters. Such an approach 25 has proved surprisingly effective and robust to overfitting, especially when having a small number of 26 parameters. From a high-level perspective, both EM and VI are based on likelihood maximization 27 in the presence of latent variables, which ultimately requires carrying out expectations over the 28 commonly intractable posterior. In order to address this challenge, a large spectrum of methods 29 have been proposed in the literature and we refer the reader to [5] for an excellent discussion of 30 these approaches. Here we characterize them between two extremes. At one extreme, restrictive 31 assumptions about the structure (e.g., as in mean-field approximations) or the model class (e.g., 32 using conjugate exponential families) must be made to simplify the task. At the other extreme, when 33 no assumptions are made, most existing black-box methods exploit very little information about the 34 structure of the known probabilistic model (for example, in black-box and stochastic variational 35 inference [44, 27], hierarchical approaches [45] and normalizing flows [42]). 36

Addressing the problem at its core, we hereby propose an alternative strategy to likelihood maxi37 mization that does not require the estimation of expectations over the posterior distribution. Con38 cretely, parameter learning is now viewed through the lens of optimal transport [54], where the data 39 distribution is the source and the true model distribution is the target. Instead of minimizing a Kull40 back-Leibler (KL) divergence (which likelihood maximization methods are essentially doing), here 41 we aim to find a point estimate θ ∗ that minimizes the Wasserstein distance [WD, 31] between these 42 two distributions. 43

44

45

46

47

48

49

50

This perspective allows us to leverage desirable properties of the WD in comparison with other metrics. These properties have motivated the recent surge in generative models, e.g., Wasserstein GANs [1, 9] and Wasserstein Auto-encoders [50]. Indeed, the WD is shown to be well-behaved in situations where standard metrics such as the KL or JS (Jensen-Shannon) divergences are either infinite or undefined [43, 4]. The WD thus characterizes a more meaningful distance, especially when the two distributions reside in low-dimensional manifolds [9]. Ultimately, this novel view enables us to pursue an ambitious goal towards a model-agnostic and scalable learning framework.

51

52

53

54

55

56

57

58

59

60

61

62

63

64

65

66

67

68

69

70

71

72

73

74

75

76

Contributions. Wepresent an entirely different view that casts parameter estimation as an optimal transport problem [54], where the goal is to find the optimal plan transporting 'mass' from the data distribution to the model distribution. To achieve this, our method minimizes the WD between these two distributions. This permits a flexible framework applicable to any type of variable and graphical structure. In summary, we make the following contributions:

We introduce OTP-DAG - an O ptimal T ransport framework for P arameter Learning in D irected A cyclic G raphical models. OTP-DAG is an alternative line of thinking about parameter learning. Diverging from the existing frameworks, the underlying idea is to find the parameter set associated with the distribution that yields the lowest transportation cost from the data distribution.

We present theoretical developments showing that minimizing the transport cost is equivalent to minimizing the reconstruction error between the observed data and the model generation. This renders a tractable training objective to be solved efficiently with stochastic optimization.

We provide empirical evidence demonstrating the versatility of our method on various graphical structures. OTP-DAG is shown to successfully recover the ground-truth parameters and achieve competitive performance across a range of downstream applications.

2 Background and Related Work

We first introduce the notations and basic concepts used throughout the paper. We reserve bold capital letters (i.e., G ) for notations related to graphs. We use calligraphic letters (i.e. X ) for spaces, italic capital letters (i.e. X ) for random variables, and lower case letters (i.e. x ) for their values.

A directed graph G = ( V , E ) consists of a set of nodes V and an edge set E ⊆ V 2 of ordered pairs of nodes with ( v, v ) / ∈ E for any v ∈ V (one without self-loops). For a pair of nodes i, j with ( i, j ) ∈ E , there is an arrow pointing from i to j and we write i → j . Two nodes i and j are adjacent if either ( i, j ) ∈ E or ( j, i ) ∈ E . If there is an arrow from i to j then i is a parent of j and j is a child of i . A Bayesian network structure G = ( V , E ) is a directed acyclic graph (DAG), in which the nodes represent random variables X = [ X i ] n i =1 with index set V := { 1 , ..., n } . Let PA X i denote the set of variables associated with parents of node i in G .

In this work, we tackle the classic yet important problem of learning the parameters of a directed 77 graph from partially observed data . Let O ⊆ V and X O = [ X i ] i ∈ O be the set of observed nodes 78 and H := V \ O be the set of hidden nodes. Let P θ and P d respectively denote the distribution 79 induced by the graphical model and the empirical one induced by the complete (yet unknown) data. 80 Given a fixed graphical structure G and some set of i.i.d data points, we aim to find the point es81 timate θ ∗ that best fits the observed data X O . The conventional approach is to minimize the KL 82 divergence between the model distribution and the empirical data distribution over observed data 83 i.e., D KL ( P d ( X O ) , P θ ( X O )) , which is equivalent to maximizing the likelihood P θ ( X O ) w.r.t θ . 84 In the presence of latent variables, the marginal likelihood, given as P θ ( X O ) = ∫ X H P θ ( X ) dX H , 85 is generally intractable. Standard approaches then resort to maximizing a bound on the marginal 86 log-likelihood, known as the evidence lower bound (ELBO), which is essentially the objective of 87 EM [38] and VI [30]. Optimization of the ELBO for parameter learning in practice requires many 88

89

90

91

92

93

94

95

96

97

98

99

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

considerations. For vanilla EM, the algorithm only works if the true posterior density can be computed exactly. Furthermore, EM is originally a batch algorithm, thereby converging slowly on large datasets [36]. Subsequently, researchers have tried exploring other methods for scalability, including attempts to combine EM with approximate inference [56, 40, 14, 10, 13, 36, 41].

When exact inference is infeasible, a variational approximation is the go-to solution. Along this line, research efforts have concentrated on ensuring tractability of the ELBO via the mean-field assumption [11] and its relaxation known as structured mean field [47]. Scalability has been one of the main challenges facing the early VI formulations since it is a batch algorithm. This has triggered the development of stochastic variational inference (SVI) [27, 26, 16, 29, 8, 7] which applies stochastic optimization to solve VI objectives. Another line of work is collapsed VI that explicitly integrates out certain model parameters or latent variables in an analytic manner [23, 32, 48, 34]. Without a closed form, one could resort to Markov chain Monte Carlo [18, 19, 21], which however tends to be slow. More accurate variational posteriors also exist, namely, through hierarchical variational models [45], implicit posteriors [49, 58, 37, 49], normalizing flows [33], or copula distribution [51]. To avoid computing the ELBO analytically, one can obtain an unbiased gradient estimator using Monte Carlo and re-parameterization tricks [44, 57]. As mentioned in the introduction, an excellent summary of these approaches is discussed in [5, §6]. Extensions of VI to other divergence measures than KL divergence e.g., α -divergence or f -divergence, also exist [35, 24, 55]. In the causal inference literature, a related direction is to learn both the graphical structure and parameters of the corresponding structural equation model [60, 17]. These frameworks are often limited to additive noise models while assuming no latent confounders.

3 Optimal Transport for Learning Directed Graphical Models

We begin by explaining how parameter learning can be reformulated into an optimal transport problem [53] and thereafter introduce our novel theoretical contribution.

Weconsider a DAG G ( V , E ) over random variables X = [ X i ] n i =1 that represents the data generative process of an underlying system. The system consists of X as the set of endogenous variables and U = { U i } n i =1 as the set of exogenous variables representing external factors affecting the system. Associated with every X i is an exogenous variable U i whose values are sampled from a prior distribution P ( U ) independently from other exogenous variables. For the purpose of this work, our framework operates on an extended graph consisting of both endogenous and exogenous nodes (See Figure 1b). In the graph G , U i is represented by a node with no ancestors that has an outgoing arrow towards node i . Consequently, for every endogenous variable, its parent set PA X i is extended to include an exogenous variable and possibly some other endogenous variables. Henceforth, every distribution P θ i ( X i | PA X i ) can be reparameterized into a deterministic assignment

The ultimate goal is to estimate θ = { θ i } n i =1 as the parameters of the set of deterministic functions ψ = { ψ i } n i =1 . We will use the notation ψ θ to emphasize this connection from now on.

𝑿𝟏

𝑿𝟐

𝑿𝟑

𝑿𝟒

(a) DAG

𝑼𝟒

𝑼𝟐

𝑼𝟑

𝑼𝟏

𝑿𝟏

𝑿𝟐

𝑿𝟑

𝑿𝟒

(b) Extended DAG

𝑼𝟒

𝑼𝟐

𝑼𝟑

𝑼𝟏

𝑿𝟏

𝑿𝟐

𝑿𝟑

𝑿𝟒

(c) Algorithmic DAG

Given the data distribution P d ( X O ) and the model distribution P θ ( X O ) over the observed set O , 125 the optimal transport (OT) goal is to find the parameter set θ that minimizes the cost of transport 126 between these two distributions. The Kantorovich's formulation of the problem is given by 127

where P ( X ∼ P d , Y ∼ P θ ) is a set of all joint distributions of ( P d ; P θ ) and c : X O ×X O ↦→R + is any measurable cost function over X O (i.e., the product space of the spaces of observed variables) that is defined as c ( X O , Y O ) := ∑ i ∈ O c i ( X i , Y i ) where c i is a measurable cost function over a space of a certain observed variable.

Let P θ (PA X i , U i ) denote the joint distribution of PA X i and U i factorized according to the graphical model. Let U i denote the space over random variable U i . The key ingredient of our theoretical development is local backward mapping. For every observed node i ∈ O , we define a stochastic 'backward' map ϕ i : X i ↦→ Π k ∈ PA Xi X k ×U i such that ϕ i ∈ C ( X i ) where C ( X i ) is the constraint set given as

Essentially, ϕ i pushes the data marginal of X i forward to the model marginal of its parent variables. If PA X i are latent variables, ϕ i can be viewed as a stochastic decoder mapping X i to the conditional density ϕ i (PA X i | X i ) .

Theorem 1 presents the main theoretical contribution of our paper. Our OT problem is concerned with finding the optimal set of deterministic 'forward' maps ψ θ and stochastic "backward" maps { ϕ i ∈ C ( X i ) } i ∈ O that minimizes the cost of transporting the mass from P d to P θ over O . While the formulation in Eq. (1) is not trainable, we show that the problem is reduced to minimizing the reconstruction error between the data generated from P θ and the observed data. To understand how reconstruction works, let us examine Figure 1c. Given X 1 and X 3 as observed nodes, we sample X 1 ∼ P d ( X 1 ) , X 3 ∼ P d ( X 3 ) and evaluate the local densities ϕ 1 ( PA X 1 | X 1 ) , ϕ 3 ( PA X 3 | X 3 ) where PA X 1 = { X 2 , X 4 , U 1 } and PA X 3 = { X 4 , U 3 } . The next step is to sample PA X 1 ∼ ϕ 1 ( PA X 1 | X 1 ) and PA X 3 ∼ ϕ 3 ( PA X 3 | X 3 ) , which are plugged back to the model ψ θ to obtain the reconstructions ˜ X 1 = ψ θ 1 ( PA X 1 ) and ˜ X 3 = ψ θ 3 ( PA X 3 ) . We wish to learn θ such that X 1 and X 3 are reconstructed correctly. For a general graphical model, this optimization objective is formalized as

Theorem 1 For every ϕ i as defined above and fixed ψ θ ,

The proof is provided in Appendix A. It is seen that Theorem 1 set ups a trainable form for our optimization solution. Notice that the quality of the reconstruction hinges on how well the backward maps approximate the true local densities. To ensure approximation fidelity, every backward function ϕ i must satisfy its push-forward constraint defined by C . In the above example, the backward maps ϕ i and ϕ 3 must be constructed such that ϕ 1 #( X 1 ) = P θ ( X 2 , X 4 , U 1 ) and ϕ 3 #( X 3 ) = P θ ( X 4 , U 3 ) . This gives us a constraint optimization problem, and we relax the constraints by adding a penalty to the above objective.

The final optimization objective is therefore given as

where D is any arbitrary divergence measure and η > 0 is a trade-off hyper-parameter. D ( ϕ, P θ ) is a short-hand for divergence between all pairs of backward and forward distributions.

This theoretical result provides us with several interesting properties: (1) to minimize the global OT cost between the model distribution and the data distribution, one only needs to characterize the local densities by specifying the backward maps from every observed node to its parents and optimizing them with appropriate cost metrics; (2) all model parameters are optimized simultaneously within a single framework whether the variables are continuous or discrete ; (3) the computational process can be automated without deriving an analytic lower bound or restricting to certain graph-

128

129

130

131

132

133

134

135

136

137

138

139

140

141

142

143

144

145

146

147

148

149

150

151

152

153

154

155

156

157

158

159

160

161

162

163

164

165

166

167

168

ical structures. In connection with VI, OTP-DAG is also optimization-based. We in fact leverage 169 modern VI techniques of reparameterization and amortized inference [6] for solving it efficiently 170 via stochastic gradient descent. However, unlike such advances as hierarchical VI, our method does 171 not place any prior over the variational distribution on the latent variables underlying the variational 172 posterior [45]. For providing a guarantee, OTP-DAG relies on the condition that the backward maps 173 are sufficiently expressive to cover the push-forward constraints. We prove further in Appendix A 174 that given a suitably rich family of backward functions, our algorithm OTP-DAG can converge to the 175 ground-truth parameters. Details on our algorithm can be found in Appendix B. In the next section, 176 we illustrate how OTP-DAG algorithm is realized in practical applications. 177

178

179

180

181

182

183

184

185

186

187

188

189

190

191

192

193

194

195

196

197

198

199

200

201

202

203

204

205

206

207

208

209

210

4 Applications

We apply OTP-DAG on 3 widely-used graphical models for a total of 5 different sub-tasks. Here we aim to demonstrate the versatility of OTP-DAG: OTP-DAG can be exploited for various purposes through a single learning procedure. In terms of estimation accuracy, OTP-DAG is capable of recovering the ground-truth parameters while achieving the comparable or better performance level of existing frameworks across downstream tasks. 1

We consider various directed probabilistic models with either continuous or discrete variables. We begin with (1) Latent Dirichlet Allocation [12] for topic modeling and (2) Hidden Markov Model (HMM) for sequential modeling tasks. We conclude with a more challenging setting: (3) Discrete Representation Learning (Discrete RepL) that cannot simply be solved by EM or MAP (maximum a posteriori). It in fact invokes deep generative modeling via a pioneering development called Vector Quantization Variational Auto-Encoder (VQ-VAE) [52]. We investigate an application of OTP-DAG algorithm to learning discrete representations by grounding it into a parameter learning problem.

Note that our goal is not to achieve the state-of-the-art performance, rather to prove OTP-DAG as a versatile approach for learning parameters of directed graphical models. Figure 2 illustrates the empirical DAG structures of the 3 applications. Unlike the standard visualization where the parameters are considered hidden nodes, our graph separates model parameters from latent variables and only illustrates random variables and their dependencies (except the special setting of Discrete RepL). We also omit the exogenous variables associated with the hidden nodes for visibility, since only those acting on the observed nodes are relevant for computation. There is also a noticeable difference between Figure 2 and Figure 1c: the empirical version does not involve learning the backward maps for the exogenous variables. This stems from an experimental observation that sampling the noise from an appropriate prior distribution at random suffices to yield accurate estimation. We find it to be beneficial in that training complexity can be greatly reduced. In the following, we report the main experimental results, leaving the discussion of the formulation and technicalities in Appendix C. In all tables, we report the average results over 5 random initializations and the best ones are highlighted in bold. In addition, ↑ , ↓ indicate higher/lower performance is better, respectively.

𝑼

𝑾

𝒁

𝜽

𝑵

𝑴

(a) LDA

𝑿𝟐

𝒁 𝟏

𝑿𝟏

𝒁 𝟐

𝒁 𝟑

𝑿𝟑

𝑿𝟒

𝒁 𝟒

𝑼𝟏

𝑼𝟐

𝑼𝟑

𝑼𝟒

(b) HMM

𝒁

𝑿

𝝁 𝑪

𝑼

𝑪

(c) Discrete RepL

4.1 Latent Dirichlet Allocation

Let us consider a corpus D of M independent documents where each document is a sequence of N words denoted by W = ( W 1 , W 2 , · · · , W N ) . Documents are represented as random mixtures over K latent topics, each of which is characterized by a distribution over words. Let V be the size of a vocabulary indexed by { 1 , · · · , V } . Latent Dirichlet Allocation (LDA) [12] dictates the following generative process for every document in the corpus:

216

217

218

219

where Dir ( . ) is a Dirichlet distribution. θ is a K -dimensional vector that lies in the ( K -1) -simplex and γ k is a V -dimensional vector represents the word distribution corresponding to topic k . In the standard model, α, β, K are hyper-parameters and θ, γ are learnable parameters. Throughout the experiments, the number of topics K is assumed known and fixed.

220

221

222

223

224

225

226

227

228

229

230

231

232

233

234

235

236

237

238

239

240

241

242

243

244

245

246

247

248

Parameter Estimation. To test whether OTP-DAG can recover the true parameters, we generate synthetic data in a simplified setting: the word probabilities are parameterized by a K × V matrix γ where γ kn := P ( W n = 1 | Z n = 1) ; γ is now a fixed quantity to be estimated. We set α = 1 /K uniformly and generate small datasets for different number of topics K and sample size N . Inspired by the setup of [20], for every topic k , the word distribution γ k can be represented as a square grid where each cell, corresponding to a word, is assigned an integer value of either 0 and 1 , indicating whether a certain word is allocated to the k th topic or not. As a result, each topic is associated with a specific pattern. For simplicity, we represent topics using horizontal or vertical patterns (See Figure 3). Following the above generative model, we sample 3 sets of data w.r.t 3 sets of configuration triplets { K,M,N } : { 10 , 1000 , 100 } , { 20 , 5000 , 200 } and { 30 , 10000 , 300 } .

We compare OTP-DAG with Batch EM [38] and SVI [25, 27]. For the baselines, only γ is learnable whereas α is set fixed to be uniform, whereas for our method OTP-DAG, we take on a more challenging task of learning both parameters . We report the fidelity of the estimation of γ in Table 1 wherein OTP-DAG is shown to yield estimates closest to the ground-truth values. At the same time, our estimates for α (averaged over K ) are nearly 100% faithful at 0 . 10 , 0 . 049 , 0 . 033 (recall that the ground-truth α is uniform over K where K = 10 , 20 , 30 respectively).

Figure 3 illustrates the model topic distribution at the end of training. OTP-DAG recovers all of the ground-truth patterns, and as further shown Figure 4, most of the patterns in fact converge well before training ends.

Ground

truth

OTP

EM

SVI

Topic Evaluation. In this application, we use OTP-DAG to infer the topics of 3 real-world datasets: 2 20 News Group, BBC News and DBLP. We here revert to the original generative process where the topic-word distribution follows a Dirichlet distribution parameterized by the concentration parameters β , instead of having γ as a fixed quantity. β is now initialized as a matrix of real values ( β ∈ R K × V ) representing the log concentration values. Table 2 reports the quality of the inferred topics from OTP-DAG, in comparison with Batch EM and SVI. For every topic k , we select top 10 most related words according to γ k to represent it. Topic quality is evaluated via the diversity and coherence of the selected words. Diversity refers to the proportion of unique words, whereas Coherence is measured with normalized pointwise mutual information [2], reflecting the extent to which the words in a topic are associated with a common theme.

249

250

251

252

253

254

255

256

257

258

259

260

261

262

263

264

Ground

truth

Epoch 100

Epoch 200

Epoch 300

4.2 Hidden Markov Models

Poisson Time-series Data Segmentation. This application deals with time-series data following a Poisson hidden Markov model (See Figure 2b). Given a time series of T steps, the task is to segment the data stream into K different states, each of which is associated with a Poisson observation model with rate λ k . The observation at each step t is given as

Following [39], we use a uniform prior over the initial state. The Markov chain stays in the current state with probability p and otherwise transitions to one of the other K -1 states uniformly at random. The transition distribution is given as

Let P ( Z 1 ) and P ( Z t | Z t -1 ) respectively denote these prior transition distributions. We generate a synthetic dataset D of 200 observations at rates λ = { 12 , 87 , 60 , 33 } with change points occurring at times (40 , 60 , 55) . We would like to learn the concentration parameters λ 1: K = [ λ k ] K k =1 through which segmentation can be realized, assuming that the number of states K = 4 is known.

Table 3 demonstrates the quality of our estimates, in comparison with MAP estimates. Our estimation approaches the ground-truth values comparably to MAP. We note that the MAP solution requires the analytical marginal likelihood of the model, which is not necessary for our method. Fig-

ure 5a reports the most probable state for each observation, inferred from our backward distribution 265 ϕ ( X 1: T ) . It can be seen that the partition overall aligns with the true generative process the data. 266

0

25

50

75

100

125

150

175

200

time

20

40

60

80

100

latent rate

Inferred latent rate over time

inferred rate

observed counts

(a) Poisson time-series segmentation

4

8

12

16

20

24

No. hidden states

K

0

2

4

6

8

10

Negative Log-Likelihood

OTP

MAP

5

10

15

20

25

30

35

40

Time (minutes)

(b) Polyphonic music modeling

Polyphonic Music Modeling. We consider another application of HMM to model sequences of polyphonic music. The data under analysis is the corpus of 382 harmonized chorales by J. S. Bach [3]. The training set consists of N = 229 sequences, each of which has a maximum length of T = 129 and D = 51 notes. The data matrix is a Boolean tensor of size N × T × D . We follow the standard preprocessing where 37 redundant notes are dropped. 3

The observation at each time step is modeled using a factored observation distribution of the form

where B d ( k ) = P ( X td = 1 | Z t = k ) and k = 1 , · · · , K . Similarly, we use a uniform prior over the initial state. Following [39], the transition probabilities are sampled from a Dirichlet distribution with concentration parameters α 1: K , where α k = 1 if the state remains and 0 . 1 otherwise,

The parameter set θ is a matrix size D × K where each element θ ij ∈ [0 , 1] parameterizes B d k ( . ) . The goal is to learn these probabilities with underlying HMM sharing the same structure as Figure 2b. The main difference is that the previous application only deals with one sequence, while here we consider a batch of sequences. For larger datasets, estimating MAP of an HMM can be expensive. Figure 5b reports negative log-likelihood of the learned models on the test set, along with training time (in minutes) at different values of K . Our fitted HMM closely approaches the level of performance of MAP. Both models are optimized using mini-batch gradient descent, yet OTP-DAG runs in constant time (approx. 3 minutes), significantly faster than solving MAP with SGD.

4.3 Learning Discrete Representations

Many types of data exist in the form of discrete symbols e.g., words in texts, or pixels in images. This motivates the need to explore the latent discrete representations of the data, which can be useful for planning and symbolic reasoning tasks. Viewing discrete representation learning as a parameter learning problem, we endow it with a probabilistic generative process as illustrated in Figure 2c. The problem deals with a latent space C ∈ R K × D composed of K discrete latent sub-spaces of D dimensionality. The probability a data point belongs to a discrete sub-space c ∈ { 1 , · · · , K } follows a K -way categorical distribution π = [ π 1 , · · · , π K ] . In the language of VQ-VAE, each c is referred to as a codeword and the set of codewords is called a codebook . Let Z ∈ R D denote the latent variable in a sub-space. On each sub-space, we impose a Gaussian distribution parameterized by µ c , Σ c where Σ c is diagonal. The data generative process is described as follows:

267

268

269

270

271

272

273

274

275

276

277

278

279

280

281

282

283

284

285

286

287

288

289

290

291

292

293

294

295

296

297

298

299

300

301

302

303

304

305

306

307

308

309

310

311

312

313

314

315

316

317

318

319

320

321

322

323

324

325

326

327

328

329

330

331

332

333

334

335

336

337

where ψ is a highly non-convex function with unknown parameters θ and often parameterized with a deep neural network. Q refers to the quantization of Z to µ c defined as µ c = Q ( Z ) where c = argmin c d z ( Z ; µ c ) and d z = √ ( Z -µ c ) T Σ -1 c ( Z -µ c ) is the Mahalanobis distance.

The goal is to learn the parameter set { π, µ, Σ , θ } with µ = [ µ k ] K k =1 , Σ = [Σ k ] K k =1 such that the model captures the key properties of the data. Fitting OTP-DAG to the observed data requires constructing a backward map ϕ : X ↦→ R D from the input space back to the latent space. In connection with vector quantization, the backward map is defined via Q and an encoder f e as

Following VQ-VAE [52], our practical implementation considers Z as an M -component latent embedding. We experiment with images in this application and compare OTP-DAG with VQ-VAE on 3 popular datasets: CIFAR10, MNIST and SVHN. Since the true parameters are unknown, we assess how well the latent space characterizes the input data through the quality of the reconstruction of the original images. Our analysis considers various metrics measuring the difference/similarity between the two images on patch (SSIM), pixel (PSNR), feature (LPIPS) and dataset (FID) levels. We also compute Perplexity to evaluate the degree to which the latent representations Z spread uniformly over K sub-spaces. Table 4 reports our superior performance in preserving high-quality information of the input images. VQ-VAE suffers from poorer performance mainly due to an issue called codebook collapse [59] where most of latent vectors are quantized to few discrete codewords, while the others are left vacant. Meanwhile, our framework allows for control over the number of latent representations assigned to each codeword through learning π , ensuring all codewords are utilized. See Appendix C.3 for detailed formulation and qualitative examples.

5 Limitations

Our framework employs amortized optimization that requires continuous relaxation or reparameterization of the underlying model distribution to ensure the gradients can be back-propagated effectively. For discrete distributions and for some continuous ones (e.g., Gamma distribution), this is not easy to attain. To this end, a recent proposal on Generalized Reparameterization Gradient [46] is a viable solution. OTP-DAG also relies on the expressivity of the backward maps. Since our backward mapping only considers local dependencies, it is however simpler to find a good approximation compared to VI where the variational approximator should ideally characterize the entire global dependencies in the graph. We use neural networks to model the backward conditionals. With enough data, network complexity, and training time, the difference between the modeled distribution and the true conditional can be assumed to be smaller than an arbitrary constant ϵ based on the universal approximation theorem [28].

6 Conclusion and Future Work

This paper contributes a novel approach based on optimal transport to learning parameters of directed graphical models. The proposed algorithm OTP-DAG is general and applicable to any directed graph with latent variables regardless of variable types and structural dependencies. As for future research, this new perspective opens up promising avenues, for instance applying OTP-DAG to structural learning problems where edge existence and directionality can be parameterized for continuous optimization, or extending it to learning undirected graphical models.

References 338

Jonas Adler and Sebastian Lunz. Banach wasserstein gan. Advances in neural information 339 processing systems , 31, 2018. 2 340

Nikolaos Aletras and Mark Stevenson. Evaluating topic coherence using distributional seman341 tics. In Proceedings of the 10th international conference on computational semantics (IWCS 342 2013)-Long Papers , pages 13-22, 2013. 6 343

Moray Allan and Christopher Williams. Harmonising chorales by probabilistic inference. Ad344 vances in neural information processing systems , 17, 2004. 8 345

Luca Ambrogioni, Umut Güçlü, Yagmur Güçlütürk, Max Hinne, Marcel A.J. Van Gerven, and 346 Eric Maris. Wasserstein variational inference. Advances in Neural Information Processing 347 Systems , 2018-December(NeurIPS):2473-2482, 2018. 2 348

Luca Ambrogioni, Kate Lin, Emily Fertig, Sharad Vikram, Max Hinne, Dave Moore, and 349 Marcel van Gerven. Automatic structured variational inference. In Arindam Banerjee and 350 Kenji Fukumizu, editors, Proceedings of The 24th International Conference on Artificial In351 telligence and Statistics , volume 130 of Proceedings of Machine Learning Research , pages 352 676-684. PMLR, 13-15 Apr 2021. 1, 3 353

Brandon Amos. Tutorial on amortized optimization for learning to optimize over continuous 354 domains. arXiv preprint arXiv:2202.00665 , 2022. 5 355

Animashree Anandkumar, Rong Ge, Daniel Hsu, Sham M Kakade, and Matus Telgarsky. Ten356 sor decompositions for learning latent variable models. Journal of machine learning research , 357 15:2773-2832, 2014. 3 358

Animashree Anandkumar, Daniel Hsu, and Sham M Kakade. A method of moments for mix359 ture models and hidden markov models. In Conference on Learning Theory , pages 33-1. 360 JMLR Workshop and Conference Proceedings, 2012. 3 361

Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein generative adversarial net362 works. In International conference on machine learning , pages 214-223. PMLR, 2017. 2 363

Matthew J Beal and Zoubin Ghahramani. Variational bayesian learning of directed graphical 364 models with hidden variables. 2006. 3 365

Christopher M Bishop and Nasser M Nasrabadi. Pattern recognition and machine learning , 366 volume 4. Springer, 2006. 3 367

David M Blei, Andrew Y Ng, and Michael I Jordan. Latent dirichlet allocation. Journal of 368 machine Learning research , 3(Jan):993-1022, 2003. 5 369

Olivier Cappé and Eric Moulines. On-line expectation-maximization algorithm for latent 370 data models. Journal of the Royal Statistical Society: Series B (Statistical Methodology) , 371 71(3):593-613, 2009. 3 372

Bernard Delyon, Marc Lavielle, and Eric Moulines. Convergence of a stochastic approximation 373 version of the em algorithm. Annals of statistics , pages 94-128, 1999. 3 374

A. P. Dempster, N. M. Laird, and D. B. Rubin. Maximum Likelihood from Incomplete Data 375 Via the EM Algorithm . Journal of the Royal Statistical Society: Series B (Methodological) , 376 39(1):1-22, 1977. 1 377

Nick Foti, Jason Xu, Dillon Laird, and Emily Fox. Stochastic variational inference for hidden 378 markov models. Advances in neural information processing systems , 27, 2014. 3 379

Tomas Geffner, Javier Antoran, Adam Foster, Wenbo Gong, Chao Ma, Emre Kiciman, Amit 380 Sharma, Angus Lamb, Martin Kukla, Nick Pawlowski, et al. Deep end-to-end causal inference. 381 arXiv preprint arXiv:2202.02195 , 2022. 3 382

Alan E Gelfand and Adrian FM Smith. Sampling-based approaches to calculating marginal 383 densities. Journal of the American statistical association , 85(410):398-409, 1990. 3 384

Walter R Gilks, Sylvia Richardson, and David Spiegelhalter. Markov chain Monte Carlo in 385 practice . CRC press, 1995. 3 386

Thomas L Griffiths and Mark Steyvers. Finding scientific topics. Proceedings of the National 387 academy of Sciences , 101(suppl_1):5228-5235, 2004. 6 388

John Hammersley. Monte carlo methods . Springer Science & Business Media, 2013. 3 389

Ernst Hellinger. Neue begründung der theorie quadratischer formen von unendlichvielen 390 veränderlichen. Journal für die reine und angewandte Mathematik , 1909(136):210-271, 1909. 391 7 392

393

394

James Hensman, Magnus Rattray, and Neil Lawrence. Fast variational inference in the conjugate exponential family. Advances in neural information processing systems , 25, 2012. 3

Jose Hernandez-Lobato, Yingzhen Li, Mark Rowland, Thang Bui, Daniel Hernández-Lobato, 395 and Richard Turner. Black-box alpha divergence minimization. In International conference on 396 machine learning , pages 1511-1520. PMLR, 2016. 3 397

Matthew Hoffman, Francis Bach, and David Blei. Online learning for latent dirichlet alloca398 tion. advances in neural information processing systems , 23, 2010. 6 399

Matthew D Hoffman and David M Blei. Structured stochastic variational inference. In Artifi400 cial Intelligence and Statistics , pages 361-369, 2015. 3 401

Matthew D Hoffman, David M Blei, Chong Wang, and John Paisley. Stochastic variational 402 inference. Journal of Machine Learning Research , 2013. 1, 3, 6 403

Kurt Hornik, Maxwell Stinchcombe, and Halbert White. Multilayer feedforward networks are 404 universal approximators. Neural networks , 2(5):359-366, 1989. 9 405

406

407

408

Matthew Johnson and Alan Willsky. Stochastic variational inference for bayesian time series models. In International Conference on Machine Learning , pages 1854-1862. PMLR, 2014. 3

409

410

411

412

413

414

415

416

417

418

419

420

421

422

423

424

425

426

427

428

Michael I Jordan, Zoubin Ghahramani, Tommi S Jaakkola, and Lawrence K Saul. An introduction to variational methods for graphical models. Machine learning , 37:183-233, 1999. 1, 2

Leonid V Kantorovich. Mathematical methods of organizing and planning production. Management science , 6(4):366-422, 1960. 2

Nathaniel J King and Neil D Lawrence. Fast variational inference for gaussian process models through kl-correction. In Machine Learning: ECML 2006: 17th European Conference on Machine Learning Berlin, Germany, September 18-22, 2006 Proceedings 17 , pages 270-281. Springer, 2006. 3

Durk P Kingma, Tim Salimans, Rafal Jozefowicz, Xi Chen, Ilya Sutskever, and Max Welling. Improved variational inference with inverse autoregressive flow. Advances in neural information processing systems , 29, 2016. 3

Miguel Lázaro-Gredilla, Steven Van Vaerenbergh, and Neil D Lawrence. Overlapping mixtures of gaussian processes for the data association problem. Pattern recognition , 45(4):1386-1395, 2012. 3

Yingzhen Li and Richard E Turner. Rényi divergence variational inference. Advances in neural information processing systems , 29, 2016. 3

Percy Liang and Dan Klein. Online em for unsupervised models. In Proceedings of human language technologies: The 2009 annual conference of the North American chapter of the association for computational linguistics , pages 611-619, 2009. 3

Dmitry Molchanov, Valery Kharitonov, Artem Sobolev, and Dmitry Vetrov. Doubly semi429 implicit variational inference. In The 22nd International Conference on Artificial Intelligence 430 and Statistics , pages 2593-2602. PMLR, 2019. 3 431

Todd K Moon. The expectation-maximization algorithm. IEEE Signal processing magazine , 432 13(6):47-60, 1996. 2, 6 433

Kevin P Murphy. Probabilistic machine learning: Advanced topics . MIT Press, 2023. 7, 8 434

Radford M Neal and Geoffrey E Hinton. A view of the em algorithm that justifies incremental, 435 sparse, and other variants. Learning in graphical models , pages 355-368, 1998. 3 436

Ronald C Neath et al. On convergence properties of the monte carlo em algorithm. Advances 437 in modern statistical theory and applications: a Festschrift in Honor of Morris L. Eaton , pages 438 43-62, 2013. 3 439

George Papamakarios, Eric Nalisnick, Danilo Jimenez Rezende, Shakir Mohamed, and Balaji 440 Lakshminarayanan. Normalizing flows for probabilistic modeling and inference. Journal of 441 Machine Learning Research , 22:1-64, 2021. 1 442

Gabriel Peyré, Marco Cuturi, et al. Computational optimal transport. Center for Research in 443 Economics and Statistics Working Papers , (2017-86), 2017. 2 444

Rajesh Ranganath, Sean Gerrish, and David Blei. Black box variational inference. In Artificial 445 intelligence and statistics , pages 814-822. PMLR, 2014. 1, 3 446

Rajesh Ranganath, Dustin Tran, and David Blei. Hierarchical variational models. In Interna447 tional conference on machine learning , pages 324-333. PMLR, 2016. 1, 3, 5 448

Francisco R Ruiz, Titsias RC AUEB, David Blei, et al. The generalized reparameterization 449 gradient. Advances in neural information processing systems , 29, 2016. 9 450

Lawrence Saul and Michael Jordan. Exploiting tractable substructures in intractable networks. 451 Advances in neural information processing systems , 8, 1995. 3 452

Yee Teh, David Newman, and Max Welling. A collapsed variational bayesian inference algo453 rithm for latent dirichlet allocation. Advances in neural information processing systems , 19, 454 2006. 3 455

Michalis K Titsias and Francisco Ruiz. Unbiased implicit variational inference. In The 22nd In456 ternational Conference on Artificial Intelligence and Statistics , pages 167-176. PMLR, 2019. 457 3 458

Ilya Tolstikhin, Olivier Bousquet, Sylvain Gelly, and Bernhard Schoelkopf. Wasserstein auto459 encoders. arXiv preprint arXiv:1711.01558 , 2017. 2 460

Dustin Tran, David Blei, and Edo M Airoldi. Copula variational inference. Advances in neural 461 information processing systems , 28, 2015. 3 462

Aaron Van Den Oord, Oriol Vinyals, et al. Neural discrete representation learning. Advances 463 in neural information processing systems , 30, 2017. 5, 9 464

Cédric Villani. Topics in optimal transportation , volume 58. AMS Graduate Studies in Math465 ematics, 2003. 3 466

Cédric Villani et al. Optimal transport: old and new , volume 338. Springer, 2009. 2 467

Neng Wan, Dapeng Li, and Naira Hovakimyan. F-divergence variational inference. Advances 468 in neural information processing systems , 33:17370-17379, 2020. 3 469

Greg CG Wei and Martin A Tanner. A monte carlo implementation of the em algorithm and 470 the poor man's data augmentation algorithms. Journal of the American statistical Association , 471 85(411):699-704, 1990. 3 472

Ming Xu, Matias Quiroz, Robert Kohn, and Scott A Sisson. Variance reduction properties of 473 the reparameterization trick. In The 22nd International Conference on Artificial Intelligence 474 and Statistics , pages 2711-2720. PMLR, 2019. 3 475

Mingzhang Yin and Mingyuan Zhou. Semi-implicit variational inference. In International 476 Conference on Machine Learning , pages 5660-5669. PMLR, 2018. 3 477

Jiahui Yu, Xin Li, Jing Yu Koh, Han Zhang, Ruoming Pang, James Qin, Alexander Ku, 478 Yuanzhong Xu, Jason Baldridge, and Yonghui Wu. Vector-quantized image modeling with 479 improved vqgan. In International Conference on Learning Representations . 9 480

Yue Yu, Jie Chen, Tian Gao, and Mo Yu. Dag-gnn: Dag structure learning with graph neural 481 networks. In International Conference on Machine Learning , pages 7154-7163. PMLR, 2019. 482 3 483

Peer reviews:
comments: SUMMARY:
This paper uses optimal transport for learning the parameters of a DAG structure that represents a Bayesian network given samples drawn from it (data). In the proposed algorithm the data can be incomplete (i.e. some random variables may be latent). 
The algorithm can be seen as a generalization of the existing "Optimal transport-based divergence minimization" on target distributions that are factorized as a product of conditional univariate densities.

STRENGTHS:
1. The generalization of the Optimal Transport to Bayesian networks is interesting and as far as I can say novel (though, more or less straightforward).
2. The paper is well-written (In particular the first two sections).
3. Several experimental models are presented.

WEAKNESSES:
1. This method is only suitable for learning the parameters of Bayesian network DAGs. It can not be used for learning the DAG structure. 
2. The description of the proposed algorithm (section 3 of the draft) is quite compact and many details are missing.
More details are presented in the supplementary material (Section B). Still, even there, the cost function and push-forward divergence measure that are used in the paper's experiments are not spelled out.

QUESTIONS:
I would suggest moving section B of the supplementary material to the main text and moving the details of the experimental models to the supplementary material.

LIMITATIONS:
yes
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-04

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn this submission, the authors propose an optimal transport-based method to infer the parameters of probabilistic directed graphical models from partial observations. \nIn particular, given a DAG associated with the target model, the proposed method reparameterizes the probability of a node conditioned on its parents by an encoder with external perturbation. \nAccordingly, a stochastic decoder is applied to map each node to the conditional density of its parents. \nThe above two modules lead to a model with an auto-encoding architecture, which can be learned like a Wasserstein autoencoder (WAE), as shown in Eqs. (2, 3).\nExperiments on the inference of LDA, HMM, and discrete representation models demonstrate the potentials of the proposed method.\n\n--- After rebuttal ---\n\nThanks for the authors' efforts in the rebuttal phase. After reading other reviewers' comments, my main concern about this work is still its similarity to WAE, especially the theoretical part. Although the authors claimed that WAE can be viewed as a special case of this work, in my opinion, it is more likely that this work is a special case of WAE. I am satisfied with the other part of this work, so my final score is kept as \"borderline accept\".\n\nSTRENGTHS:\n1. The paper is well-written and easy to follow.\n\n2. The proposed method is reasonable --- the objective function is based on the theoretical part of WAE, whose rationality has been guaranteed. In addition, the implementation of the proposed method is simple.\n\n3. The authors consider various applications, demonstrating the universality of the proposed method.\n\n4. The limitations of the proposed method are discussed, and the potential solutions are provided at the same time.\n\nWEAKNESSES:\n1. If my understanding is correct, Theorem 1 in this submission is a special case of Theorem 1 in the WAE work [a]. The final objective (Eq.(3)) approximates the Wasserstein distance by relaxing the constraint of phi_i to a regularizer, which is also similar to the strategy of WAE. The authors should discuss the connections and the differences between the proposed method (and theory) and that in [a]. \n\n[a] Tolstikhin, Ilya, Olivier Bousquet, Sylvain Gelly, and Bernhard Schoelkopf. \"Wasserstein Auto-Encoders.\" In International Conference on Learning Representations. 2018.\n\n\n2. As the authors mentioned, the proposed method leverages the amortization strategy to reparameterize the conditional distributions. Therefore, in the experimental part, the authors should consider some amortization methods as baselines, e.g., those in [b, c, d].\n\n[b] Kim, Yoon, Sam Wiseman, Andrew Miller, David Sontag, and Alexander Rush. \"Semi-amortized variational autoencoders.\" In International Conference on Machine Learning, pp. 2678-2687. PMLR, 2018.\n\n[c] Agrawal, Abhinav, and Justin Domke. \"Amortized variational inference for simple hierarchical models.\" Advances in Neural Information Processing Systems 34 (2021): 21388-21399.\n\n[d] Huang, Chin-Wei, Shawn Tan, Alexandre Lacoste, and Aaron C. Courville. \"Improving explorability in variational inference with annealed variational objectives.\" Advances in neural information processing systems 31 (2018).\n\n\n3. The datasets used in the experimental part are relatively simple and small. Especially in the experiments of discrete representation learning, I wonder 1) whether the proposed method can deal with images with larger sizes, e.g., face images in CelebA, and 2) besides reconstruction, whether the proposed method can generate images with tolerable qualities.\n\nQUESTIONS:\nSee the comments in the section on weaknesses. \n\nTo demonstrate the novelty of the proposed method, the authors should clarify the connections and differences between the proposed method and WAE.\n\nLIMITATIONS:\nThe authors discussed the limitations of using the amortization strategy at the end of the submission. \nSome potential solutions are proposed at the same time.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors propose a method for learning parameters $\\theta$ of a DAG within an optimal transport (OT) framework, minimizing the Wasserstein distance between the data distribution $P_d$ and the model distribution $P_\\theta$ in $\\theta$. The Kantorovich formulation of this problem is a minimization of an expected cost $c(X,Y)$ over all joint distributions on $X,Y$ such that marginally $X \\sim P_d, Y\\sim P_\\theta$; this is implemented in practice by empirically drawing $X_i$ from the data, and then $PA_{X_i}$ conditionally on $X_i$ to satisfy $PA_{X_i} \\mid X_i \\sim P_\\theta(\\textrm{PA}_{X_i})$ through the use of stochastic “backward” mappings. This makes optimization tractable over the space of backward mappings $\\phi$ and model parameters $\\theta$, so long as the constraint above is relaxed to a regularization term.\n\nSTRENGTHS:\n* DAGs represent an extremely rich family of models, and the work also generalizes to settings with unobserved variables, making this method easy to apply in a variety of settings.\n* Unlike variational methods, evidence bounds need not be computed; the proposed method is more “direct” in this sense. \n* The final optimization objective is easy to compute and computationally cheap.\n* The framing of the problem from the lens of OT is novel to my knowledge, and provides an interesting formulation for optimization.\n* The most significant strength of the paper is the evaluation of the method on a rich test suite of interesting problems such as LDA and Poisson time series segmentation. Comparisons show that the proposed method outperforms existing methods such as Batch EM and SVI in a variety of scenarios and metrics.\n\nWEAKNESSES:\n* The method reuqires that the random variables be reparameterizable in the sense of the equation at line 122 (this equation should be numbered); this may limit the family of joint distributions that can be considered. Though the authors say on line 167 that discrete variables can be used, reparameterization is tricky in these cases (as the authors acknowledge in the limitations section). \n* The backward maps $\\phi_i$ are a confusing quantity (and potentially hard to fit); see the Questions section. This could be due to a lack of understanding of my part. As I understand it, however, it raises questions about the efficacy of the proposed method.\n* The formality of the OT framing is appealing, yet this formality is ultimately dropped for a regularized analogue that does not solve exactly the same problem that is posed.\n* Some notation is confusing; for example, does $PA_{X_i}$ include exogenous variable $U_i$? The discussion around line 120 and line 147 suggest so, but then notationally the equations at line 122 and 136 suggest otherwise.\n\nQUESTIONS:\n* Is the data distribution $P_d(X_O)$ (line 125) a mixture of point masses, or a continuous (but unknown) distribution? This might be worth mentioning explicitly. In the case of the former, it could motivate use of the Kantorovich formulation instead of the Monge formulation of the OT problem.\n* Line 128 can be expressed more precisely, and means to define $\\mathcal{P}(X \\sim P_d, Y \\sim P_\\theta)$ as the set of joint distributions on $X, Y$ with $P_d$, $P_\\theta$ as marginals. \n* In line 137, is this meant to say “backward” as it’s a backward map? The # notation should be defined somewhere. \n* The constraint set on $\\phi_i$ requires the following: that for $X_i \\sim P_d$, the random variable $\\phi_i(X_i) \\sim P_\\theta(\\textrm{PA}_{X_i},  U_i)$. This ensures that $(X_i, \\phi_i(X_i))$ are a joint distribution that marginally follow the requirements of the Kantorovich formulation. Is this all correct so far?\n* If so, the constraint on the backward maps seems potentially ill-posed: by construction the distribution $P_\\theta(PA_{X_i},  U_i)$ does not depend on the *value* of $X_i$ (it does depend on $X_i$ structurally, as we must know what nodes are its parents). For any given $\\theta$, we are able to directly sample from $P_\\theta(PA_{X_i},  U_i)$ via ancestral sampling. It seems to be then that training a neural network decoder such that $\\phi_i(X_i) \\sim P_\\theta(PA_{X_i},  U_i)$ could result in a collapse to the “prior” defined the model, in other words the decoder could learn to ignore the input $X_i$ and simply mimic ancestral samples from $ P_\\theta(PA_{X_i},  U_i)$, which presumably is what the stochastic backward maps are trained with. In such a case, the optimization over the space $\\mathcal{P}(X \\sim P_d, Y \\sim P_\\theta)$ would only occur over a very restricted subset of joint distributions, namely those where $X, Y$ are independent or nearly so. This seems like a possibility to me that is not addressed in the paper.\n* In this extreme case, the final optimization problem becomes very simplistic: the regularization term disappears; and one simply draws ancestral samples from the model $P_\\theta$, computes the cost function with a draw from the data, and minimizes the result in $\\theta$. For simple models (or low dimensional $\\theta$) this still might perform well. To ensure the situation above is not occurring maybe the quantity $KL(P_d P_\\theta \\mid \\mid  P_d \\phi)$ could be evaluated as a metric, or included via an additional regularization term to ensure that the backward maps $\\phi$ are indeed resulting in distributions with dependence. This quantity should be significantly different from zero in that case.\n\nLIMITATIONS:\nThe settings where the work can be applied are discussed by the authors. A limitations section is included that proactively assesses some limitations of the proposed method.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors propose a framework for learning the parameters of directed graphical models based on the idea of fitting by selecting the parameter values that minimize the Wasserstein distance (WD) between the data and model distributions. They prove (Thm. 1) that these distances can be characterized as the result minimizing a cost functional over a family of constrained stochastic 'backwards mappings', which yields a solvable objective when the constraint is relaxed and a regularization term is added to the objective.  \n\nThe approach is illustrated via a series of experiments, including 3 real-world datasets, in which the proposed OPT-DAG method outperforms various baselines across disparate tasks.\n\nSTRENGTHS:\nI'm am not particularly familiar with the literature on DAG learning, nor on learning with 'optimal transport' objectives, so the extent of the originality of this paper is impossible to gauge. Operating under the assumption that this is the first work to introduce an OPT objective in this context, I would feel confident in stating that this is a fairly significant innovation. \n\nThe presentation of the paper is, in my opinion, of a high standard (modulo one or two concerns that I will voice in the next section). The appendices provide a considerable depth of exposition of their procedure, and I am satisfied that the most pressing immediate concerns are addressed, either there or in the main body.\n\nWEAKNESSES:\nI believe that the experimental section as well as the discussion could be improved. Again, I am unaware of what constitutes the benchmarking standard for DAG learning methods, but it seems to me as though the following are not adequately addressed:\n\na) The authors mention that their goal is not to achieve state-of-the-art performance, but rather to demonstrate the inherent versatility of their method. Situations in which their method might not be feasible are alluded to in Section 5 (Limitations), but the discussion here is extremely terse. Do these situations pose problems for competing alternatives as well? The chosen examples strike me as being somewhat simple. Do the authors claim that these examples are roughly representative of DAG learning problems generally? \n\nb) For the topic evaluation example (239-250) OPT-DAG is outperformed by the baselines in the 'diversity' metric in all three datasets, and outperformed on 'coherence' on the DBLP data. These results are reported with precisely no discussion or explanation. \n\nc) Section 4.2, first part (251- 266) could use some clarification. For example, the authors say \"we generate a synthetic dataset D with 200 observations at rates {12, 87, 60, 33} with change points occurring at times (40, 60, 55).\" If the model is as they say, (with a change of state happening with probability $1-p$, of which the lowest value is when $p = 0.95$, then a) why are there so few change points? At $p = 0.95$ would we not expect 10? Perhaps I am confused about what the authors are doing. Have they fixed the dataset, and computed the estimates of the rate parameters assuming that $p$ is fixed at the values indicated in the table (i.e. they fit the model 6 times, each time with a different assumed p)? Also, averaged over just the values $p = 0.75, 0.95$, OPT-DAG is actually inferior to MAP. Again, no explanation or discussion is forthcoming. \n\nd) Generally, it is not clear to me if the baselines being compared against are the most appropriate. For instance, in the final example, VQ-VAE is used as a baseline, and its poor performance is linked to a phenomenon called 'codebook collapse'. The paper provided as a reference in fact proposes an extension to the vanilla version (of VQ-VAE) that they seem to be comparing against, which seems to indicate that not only is their baseline not state-of-the-art, it is in fact very well known to not be so...\n\nQUESTIONS:\nSee above.\n\nLIMITATIONS:\nLimitations are discussed, but as per my comment above, it does not seem to me as though this discussion is adequate. The authors seem to have a variety of situations in mind in which their method will either not work or not be competitive, and a more explicit discussion here would be welcome.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-27"}