Complete Likelihood Objective for Latent Variable Models
Authors: Mikhail Arkhipov
Maria Vikhreva
Venue: Submitted to ICLR 2023 (2023)
Keywords: Posterior Collapse
Latent Variable Models
Complete Likelihood
Empirical Distribution
Assignment Problem

Abstract: In this work, we propose an alternative to the Marginal Likelihood (MaL) objective for training latent variable models, Complete Latent Likelihood (CoLLike). We analyze the objectives from the perspective of matching joint distributions. We show that MaL corresponds to a particular $KL$ divergence between some target \emph{joint} distribution and the model joint. Furthermore, the properties of the target joint explain such major malfunctions of MaL as uninformative latents (posterior collapse) and high deviation of the aggregated posterior from the prior. In CoLLike approach, we use a sample from the prior to construct a family of target joint distributions, which properties prevent these drawbacks. We utilize the complete likelihood both to choose the target from this family and to learn the model. We confirm our analysis by experiments with expressive low-dimensional latent variable models, which also indicate that it is possible to achieve high accuracy unsupervised classification using CoLLike objective.

Under review as a conference paper at ICLR 2023
COMPLETE LIKELIHOOD OBJECTIVE FOR LATENT VARIABLE MODELS
Anonymous authors Paper under double-blind review
ABSTRACT
In this work, we propose an alternative to the Marginal Likelihood (MaL) ob- jective for learning representations with latent variable models, Complete Latent Likelihood (Col.Like). We analyze the objectives from the perspective of match- ing joint distributions. We show that Mal. corresponds gence between some jurgetjoint distribution and the model joint. Furthermore, the properties of the target joint explain such major malfunctions (from the represen- tation learning perspective) of MaL uninformative latents (posterior collapse) and high deviation of the aggregated posterior from the prior. In the CoLLike approach, we use sample from the prior to construct a family of target joint distributions, which properties prevent these drawbacks. We utilize the complete likelihood both to choose the target from this family and to learn the model. We confirm our analysis by experiments with oow-dimensional latents, which also in- dicate that it is possible achieve high-accuracy unsupervised classification using CoLLike objective.
particular T.. diver-
INTRODUCTION
In the latent variable setting, the model defines a joint distribution over both observed variables w and latent variables z, while the training data contains only observed variables. The problem can be treated as an unknown 2|æ target conditional distribution. There are at least two possible solutions to this problem: try to come up with meaningful target 2|æ distribution and train the model similarly to supervised setting, or give up and focus on matching only marginals in the x domain. The latter the choice of the Mal. objective. In this work, we follow the former approach. However, instead of picking up single target conditional we construct an entire family of possible distributions and use the model likelihood to decide which conditional to use as a target.
To construct family of possible conditionals, we use sample from prior of the same size as the dataset the observed domain. All possible assignments observed samples latent ones span a family of empirical joint distributions. This can be represented as permutations of the latent samples. Despite the size of the permutations set being tremendous and growing as a factorial of the dataset size, the search of the permutation with the best likelihood can be done efficiently using combina- torial optimization. The resulting optimization procedure resembles the expectation maximization algorithm (Dempster al. 1977), where expectation replaced with the combinatorial assignment problem. Furthermore, since the proposed algorithm uses gradient-free optimization for obtaining the target distribution, the objective can be seamlessly applied to both continuous and discrete latent variables, while the discrete latents case is challenging for approaches based on the MaL(Mnih & Gregor 2014 Mnih & Rezende 2016 Tucker et al. 2017).
We analyze the objectives from the perspective of matching joint distributions. We show that MaL corresponds to specific choice of the target 2|.r conditional, while our approach takes into con- sideration family of possible conditionals. The choice of target conditional is responsible for two major failures that arise during training with the MaL objective: inability to learn informative latents, also known "posterior collapse" (Bowman et al. 2016 Razavi et al.) 2019 He et al. 2019). and divergence between the prior and the aggregated posterior (Hoffman & Johnson 2016) Makhzani et al. 2015: Zhao et al. 2019 Kim & Mnih| 2018). These characteristics are vital for latent vari- able models because posterior collapse prevents learning meaningful representation and sampling from the regions of high deviation of the latent marginals are subjected to severe quality degradation

---

Under review as a conference paper at ICLR 2023
(Rosca al. 2018). The form of the target joint also motivates the success of the complete likeli- hood in these challenges. Namely, the target distribution for CoLLike has high mutual information and matches prior.
We verify our analysis with experiments. In this work, we focus on low-dimensional latent variables to perform direct comparison with the exact MaL.. Models trained with CoLLike stably maintain high mutual information and low divergence from the prior. In turn, MaL. inevitably leads either to posterior collapse to a highly divergent aggregated posterior. Previously, for simple linear mod- els. it has been shown that there is posterior collapse during the optimization of the exact likelihood Lucas et al. (2019) Our experiments demonstrate that can as well happen with expressive models trained with exact likelihood. Along with informativeness and latent distribution matching, CoL- Like indicates no degradation of likelihood compared Mal.. Furthermore, we show that CoLLiiee objective alone can achieve high accuracy in unsupervised classification.
We show that CoLLike unifies range of existing approaches that lack probabilistic justification. Constrained K-means (Bennett et al., 2000), Permutation Invariant Training (Yu et al. 2017 Luo & Mesgarani| 2019), and Noise as Target Booannoski & Joulin (2017) are among these approaches. This allows us to extend them to different factorizations of the joint and perform analysis from the probabilistic perspective. Furthermore, CoLLike bridges likelihood and optimal transport (OT) frameworks. From this perspective, the negative likelihood plays the role of both mapping from latent to visible domain and distance function.
2 COMPLETE LIKELIHOOD OBJECTIVE
In the regular latent variable setting, we are given a dataset (a1,..., and the model po 2) po(#|=))p<>) The missing 2 can be treated as the missing P=(2(1:) part of the target joint. If we cannot come up with a reasonable 2/2 target, we can at least match the marginals in the observed domain with K Lpp(<x)) po (z))) in hope that the model will learn an informative relation between and z. This is equivalent to the maximization of MaL:
N
Maa (0) = (o = Σ log po(x,, z)dz
(I)
Justification of the MaL. comes from the equivalence of maximization of (T) and minimization of the Kullback-Leibler divergence K Lpp (x)) (1)))) which measures the discrepancy between the target empirical data distribution ps I and the model distribution p=(x) (Murphy) 2022 4.2.2).
Note that this justifies MaL only for learning distributions of observed variables, not learning repre-
sentations. The fact that MaL does not promote informativeness (Alemi et al.) 2018) clearly shows the lack of justification of Mal. for learning representation because informativeness is undoubtedly a fundamental requirement for any useful representation.
Despite the family of all possible target p(2]z) distributions being tremendous, we do not need to consider it entirely. Firstly, the target distribution must be informative. Secondly, the marginal of the target joint distribution the latent domain should match the prior p(z). The fixed prior implies that the desired marginal distribution of is known. These requirements can be interpreted (Huszár) 2017) as Infomax principle (Linsker 1988). It is not hard to get a rich family of distributions with such properties. We can obtain collection (21,..., =N) by sampling from the prior and pair this collection with the dataset (I1, xxv ). The pairing produce an empirical distribution. Empirical distribution attains highest possible Mutual Information (MT) under the assumption that there no
repeated values of x in the dataset (see Appendix B for derivation). This ensures the first require- ment. Sampling from the prior addresses the second requirement, since the collection of samples converges to p(z) (Cover & Thomas 2006) Theorem 11.2.1). However, the sampling effects can be significant problem for figh-dimensional latents. We express each pairing as some permuta- tion ग, which produces complete collection ([z), ²T(1), ..., (IN. V=(N))) and an empirical joint P8x (x. 2) = ( (2(x). Given a family of distributions, we need to decide which member of the family our target. We propose to pick the one with the highest complete likelihood relying on the model inductive biases. For this target we then once again optimize the complete likelihood of
Wee find the Greek letter 8 especially suitable for data distribution hecause it is consonant with "data" and reflects the Delta-function-liae form of the empirical distribution.
2

---

Under review as a conference paper at ICLR 2023
the (Ii, (x)( pairs with the optimal permutation x². These considerations lead to the CoLLike objective:
(0, Σ log Pe (1is (x)))
which we maximize both with respect 0 and π. An alternative view on the objective can be the following: we sample values from prior and assume that they are ground truth targets for the training dataset with unknown pairing. Figure[] the main difference between the objectives: CoLLike maximizes specific points of the joint distribution, while MaL is aimed at maximization of whole lines along the joint.
Z
Z
X
X
Figure I: Illustration of the CoLLike (left) and MaL (right) objectives. Triangles depicts sample val- ues. Filled circles represent pa (2;, z) for all possible (ar, z) pairs. Double circles indicate optimal m. Bold lines and double circles are areas of the joint to maximized.
3 OBJECTIVE ANALYSIS
We start our analysis by proving that MaL corresponds to the matching a specific joint distribution and the model joint:
K (ippssa)pn(lllllloooes = F2 (epppo((l)) log ps(5)po(2la)
log po (x)
Pa ps(x))
= Exppp(
po(x)po(222)
= -E) [log (x)ps(a) [log (o))
1
Σ log log N
1
(0)
where C' is a constant. The joint K form of the MaL brings new perspectives on the objective. It might be tempting think about MaL workaround for unknown latents that allows you not to specify the target 2|æ conditional. However, the joint form reveals that the target conditional is actually specified and equals Po (2(2)2) if we ask what distribution we want mimic. This implies that we are aiming to keep the model posterior unchanged. In addition, the form also highlights the intimate connection between Mal. and posterior.
CoLLike and common variational (Jordan et al. 1999) approximation of Mal., Evidence Lower Bound (ELBO), can also be expressed as KL divergences between joint distributions (see Table D.. We refer to Appendix A foror derivation of the equivalence. Note the elegant similarity between objectives which becomes obvious in the joint L form. All divergences share the model pa(a, 2) the second argument, which implies that the first argument is the target joint distribution. For all objectives the target joint contains the data distribution ps (æ) as a marginal a domain, thus the only difference is in the target z|æ conditional. Therefore, all the considered objectives belong the family of the following form:
[())= KT(ps(())o(ellllloocc...
KT(pp(e)) KT.(s(r))l + (pp()) [KT[v(=]<))lpo((<c)
Since the second term in is non-negative, all objectives in the family lower bounds the likelihood up to an additive constant. Note that the 2| target conditional is used to minimize the overall divergence. This affects the second term of (3) to make the lower bound tighter.
(2(xr) approximate posterior distribution parametrized by

---

Under review as a conference paper at ICLR 2023
Table 1: Considered objectives and their joint KL forms.
Original Objective
Joint KL form
CoLLike
Σ i= log po (ris Ra(i))
KIApps(a)oe<<llallooeee K(lppo((aroo(lllllloo0dd (L(ps()) (po
Mal.
D po
ELBO
Em= E. (e log po(eus)
Despite the common traits, the objectives are different. We will highlight few differences and go deeper in the following sections. Firstly, the target conditional for CoLLike pa (=|ar) is empirical, while its counterparts (=(2))) and (a(2(z) are not. Secondly, in MaL. approach, we construct particular joint distribution ps (x)pp (z(rr) and use it as target while, in CoLLike, we construct an entire family of joint distributions with desired properties. Thirdly, the target posterior readily available in CoLLike and ELBO cases, while for Mal. it could be intractable. Furthermore, CoLLike can be seamlessly applied discrete variables, while optimization of ELBO for discrete latents is challenging (Mnih & Gregor, 2014, Mnih & Rezende, 20166) Tucker et al. 2017). Lastly, the CoLLike objective allows learning models with a reverse factorization pa (z pp) (= (r), while MaL and ELBO do not. Reverse factorization is another inductive bias that can be useful or not. Furthermore, it can be significantly faster compared to regular factorization if (z(x) is assumed to be uniform and p(=|z) is factorized.
3.1 MUTUAL INFORMATION OF THE TARGET DISTRIBUTION
Mutual Information is the key property of the joint distribution in latent variable setting. It charac- terizes how dependent the observed and latent variables are. We would like know what MI value our model is targeted at for each objective. Since our objective can be expressed as K T divergence between model and target joint distributions (Table we can investigate MI values for each target joint. We define MI between æ and under p(2, distribution as:
p(r, 2) p(z)p(=)
MI 2)) E Ex,zop(a, log
(4)
For MaL, the MI of the target pa (2))pe((2)2) is determined by the model's 's current posterior po lr). Most models have no class preferences at initialization, which results in low MI of p=(2))p(22 Moreover,
aimed at keeping it unchanged, since target posterior. So, low MI at initialization might induce learning non-meaningful factorized joint throughout the training procedure. Since for ELBO the approximate posterior aligns to the true model posterior this argument is applicable to ELBO too. Furthermore, uninformative posterior is common problem when learning latent variable model (Bowman et al., 2016 Alemi et al. 2018
using the current posterior
Lucas et al. 2019 Razavi et al. 2019 He et al. 2019) known as "posterior collapse".
CoLLike target is an empirical joint distribution. represents a deterministic mapping and has constantly high MI by construction, as shown in Appendix[B] Therefore, we are aimed at mimicking a high MI distribution with our model distribution. Furthermore, CoLLike can interpreted as some realization of InfoMax principle Huszár (2017), where prior limits the entropy and deterministic mapping maximize MI.
3.2 MATCHING IN THE LATENT DOMAIN
The joint form of the objectives from Ta convenient for obtaining perspective distribution matching in the latent space. After treating ps (=())xo(2(z) as joint 2)) and rewriting the original MaL objective as:
Ea (x)a))ao(s)) log (2)
K (=)) = KT(pse(L, 2)) 2))
pe po(x)2)po(2)
(x) [K (xs0(2)=) (pa)) + K (=)llpoo)))
(5)
4

---

Under review as a conference paper at ICLR 2023
see that matching in space requires matching in space. Namely, K where pso (z) is called an aggregated posterior. It signifies that even though MaL is constructed such that z given z conditional part of the KLL between joints is zero, end up in situation where none of the model marginals match target marginals. Moreover, the learning signal from the first term of (5) might be significantly larger compared to the second term signal if the dimensions of x and differ a lot. This might lead to a sacrifice of the second divergence in favor of the first one.
o,
Matching in latent domain is considered as known challenge of latent variable modelling (Hoff- man & Johnson 2016). Mismatch with prior results in unnatural samples from areas with high deviation aggregated posterior from the prior (Rosca al., 2018). number of works focused on this problem. They either utilize additional losses that penalize discrepancy between marginals (Makhzani et al. 2015 Zhao et al., 2019 Kim Mnih 2018) or introduce learnable prior (Bauer
& Mnih 2019: Tomezak & Welling 2018).
In turn, CoLLike addresses this problem by constructing a conditional, which marginal matches prior in the latent domain. Obviously, the target marginal in z domain for CoLLike is always ps(r). In turn, the target aggregate posterior is always a sample from the prior since por (z) Ja (s(())) (=|ar)dx
P. (=) for all values, where (z) is the distribution of the sample produced by sampling from the prior. While it intuitively obvious that the empirical distribution converges to the underlying distribution, one can show that KL between the empirical sample and the prior Theorem 11.2.1).
converges in probability to (Cover & Thomas 2006
4 ALGORITHM
The objective includes maximization with respect to two parameters: т and 0. We approach it by alternating between maximization with respect to π and B. We apply stochastic minibatch technique similar to Bojanowski & Joulin (017)) which performs maximization of both æ and A for minibatch instead of the entire dataset and returns the latents back to the dataset in the optimal order. Furthermore, we interpret maximization with respect as a linear sum assignment problem (L.AP) to utilize efficient combinatorial optimization techniques (see Appendix foor the derivation). Algorithm D describes the resulting stochastic optimization procedure.
Algorithm 1 Stochastic optimization of CoLLike
Require: X = .......,NN) pa, =) po (=()))pp(2) batch size B, learning rate v) Sample Z = (21, 2N) from prior (z) ~ p(2))
while not converged do Sample random indices (i1,...,is))
Compute matrix C E BBxB for which Cq.kk = log po (o 2
Compute πt that maximizes (2) for (Ii, ..., Fin): ). (Zi, Zin) by applying LAP solver to C (0, Put (Σiy .... 2 back into Z in the optimal order (Zx*(11)) end while
0 + 0
(in))
The core of the algorithm is in computation of the matrix C and optimal permutation πτ. Both parts potentially computationally intense and challenging.
Computation of the matrix C requires B² forward passes. Note that backward passes are not re- quired for this step, hence, memory requirements low-dimensional discrete latents. Assuming the number of categories of the latent variable is lower than B, the number of all possible values of z in log Po( 2) equals K instead of B. Thus, can calculate all needed values of C using only K B forward passes instead of B², In the super- vised case, we only need the 0 update part of the entire algorithm, which requires B forward passes and backward passes. As rough estimate, we can assume that forward and backward passes take the same time, CoLLike will then require K/2 more compute time compared to the supervised setting.
mild. Furthermore, in this work focus
'This is similar to EM algorithm Dempster et al. 1977 in the sense that EM alternates between maxi- mization of the lower bound tightness (expectation step) and maximization of the the resulting tight bound (maximization step).
5

---

Under review as a conference paper at ICLR 2023
To find x² we use Hungarian algorithm Kuhn 1955 as a LAP solver. The algorithm requires the cost matrix C E Raxh as input to produce the optimal permutation It" in the form of optimal ordering of (Zi, ( 2ig) indices. The complexity of the algorithm is O(B)). The complexity of the LAP solver potentially limits the applicability of CoI-iie to large batch sizes. However, for batch sizes regularly used in practice, solving LAP results in only a minor increase in the overall computation time. For instance, in our experiments, we used batch size 64, Solving the LAP took orders of magnitude less time compared even the supervised setting. See Appendix [E for detailed timings. We also highlight that optimization with large batches is not only challenging but also could significantly reduce generalization (Xing et al. 2018 You et al., [019] 2020p. However, as gracefully shown ([luszár[(011)) this kind of minibatch combinatorial optimization provides only locally optimal solutions. Nevertheless, the size of the gap between local global optimum is still to be determined.
The result of Algorithm T a trained model. However, we are interested in the posterior pe(z(zz). For low-dimensional categorical we can exactly compute the posterior using Bayes rule po (2|xx) (a 2) (x) since po(x) = E, Po(a, = is tractable. For other cases, we can fit approximate posterior using regular variational techniques. We can also use CoLLike objec- tive to obtain estimates of & values. After training model with CoLLike objective we have X and Z arrays that are matched. We add new x samples X and sample extra values from the prior to extend Z. The inference can then be performed by optimizing 2 with respect to т.
5 CONNECTIONS
Connections with existing techniques not only give alternative perspectives on CoLLike objective, but also provide probabilistic grounding to some existing algorithms. Many well-known objectives actually use CoLLike while being motivated as an ad-hoc empirical risk minimization. We show that these objectives not only seem reasonable but are also probabilistically motivated.
While traditional K-means algorithm (MacQueen 1967 Lloyd 1982) has a probabilistic grounds (Murphy) 2022, 21.,.1.1), its constrained counterpart (Bennett et al., 2000) lacks probabilistic jus- tification. Constrained K-means is equivalent to CoLLike under factorized Gaussian po(x])) and uniform categorical p(z), which has a number of states equal to the number of clusters. This con- nection allows extending the constrained K-means approach to different generative distributions and priors. Nevertheless, a probabilistic interpretation is present in Jitta & Klami (2018), however, the choice of the complete likelihood as an objective is not explained.
Permutation Invariant Training (PIT) (Yu et al. 2017 Luo & Mesgarani| 2019) used in source sepa- ration solutions can also be expressed as CoLLike objective. For instance, in cocktail part problem, want separate a mixture of K sources. During training, we have K isolated mixture compo- nents and a network that produces K estimates of the components based on a single mixture. We don't know which network output corresponds which source and we pick permutation that pro- duces minimal total mismatch between outputs and sources. This procedure corresponds to training latent variable model with CoLLike objective, where categorical latent variable of dimension K determines the source identity. In this setting, we treat mixture components as samples in the dataset.
The closest predecessor of the CoLLike is Noise As Target (NAT) (Bojanowski & Joulin 2017).
This is an unsupervised approach to learn an image encoder. In this approach, the representations produced by network are assigned to fixed collection of vectors sampled from the uniform dis- tribution on a sphere. After this, the network parameters are adjusted to make encodings closer to the assigned vectors. This approach is equivalent to CoLLike with reverse model factorization po(a, 2) = p=(=(=)p(x) and factorized Gaussian P& (=]ar). Another approaches that obtain clear probabilistic interpretation using CoLLike include: Sinkhorn Autoencoders (Patrini et al. 2019) simultancous clustering and representation learning (Asano et al., 2020). and (Jeong & Song 2019
Bojanowski & Joulin (2017) noticed that NAT objective has Optimal Transport (OT) roots. OT framework can be used to measure diserepancy between distributions. Particularly, for a given non- negative cost function the optimal transport distance between distributions ps and p. is defined
as
OT(ps = min Ex [c(x,
pc)(p.p.))

---

Under review as a conference paper at ICLR 2023
Table Results for tractable categorical latents. MNIST, CIFAR BPD; AG News NLL.
Dataset
Objective
Accuracy
NLL/BPD 4
Agg. KL. +
MIT
CIFAR
CoLLike
14.5
3,45
0.01
2.20
CIFAR
MaL
14.0
3,46
0.74
1.50
MNIST
CoLLike
14.1
1.27
0.01
1.95
MNIST
MaL.
12.5
1.29
1.61
0.58
AG News
CoLLike
82.1
250.79
0.00
1.32
AG News
MaL
31.6
249.73
0.00
0.00
where r(ps.p) is the set of all joint distributions z and with marginals (z( and p. (2) re- spectively. Furthermore, if use a parametric model Po place P. we can fit by minimizing the distance. Note that in this case we minimize the function that already has min function inside.
When both ps and Pe are empirical, the search space I becomes countable and finite. Now it contains only pairings between points in and points in p.. Given an arbitrary initial pairing, we can express all other pairings through permutation applied to either a or z. In this case, the cost becomes
(T Pe) min Σ e (1t, xell
where II is the set of all permutation functions. This expression is almost the CoLLike objective 2). Choosing the cost function e to be log (1, 2) and switching to maximization make them equivalen(" Thus, CoLLike bridges maximum likelihood methods with OT. This connection allows bringing latest developments in OT to improve likelihood- methods. Furthermore, Appendix C provide example of the equivalence between CoLLike and Wasserstein distance. In the case, the model's complete likelihood plays the roles both a mapping from z to æ domain and a distance metric.
6 EXPERIMENTS
In this work, we focus on low-dimensional discrete latents. This type of latent variables allows to perform direct comparison with the exact likelihood. Furthermore, emphasize focus learning useful ||r instead of simplifying the model with factorized x| conditional.
Models with tractable likelihood perfect for comparing likelihood-basee algorithms because they remove the problem of the likelihood estimation precision. For this type of models, all quantities of interest can be computed exactly. Moreover, tractable likelihood allows comparing CoLLike directly with MaL instead of its approximations like ELBO.
po(x|2)
(a(xl=) x)
Transformer Block
0
Affine Coupling
0
Transformer Block
Gow
p(2)
p(z)
Inv. 1x1 Conv.
i
Emb.
Emb.z
Ø ...
Emb.z
Emb..r
Cal t. +1)
ActNorm
X,
Figure 2: Architectures used for image (left) and text (right) domains. For image domain, the three flow blocks are repeated 21 (CIFAR) and 14 (MNIST) times. Every coupling block is conditioned on
use MNIST (LeCun et al. 1998) and CIFAR (Krizhevsky 2009) datasets for image modality and AG News (Zhang et al. 2015) for text domain. All these datasets are equipped with class labels.
A cautious reader might note that for continuous variables non negativity of log po can be violated, how- ever, all model densities used in practice are finite and the corresponding cost can be made positive just by an additive constant which does change the optimization prohlem.

---

Under review as a conference paper at ICLR 2023
For images, train Glow-like normalizing flow conditioned on discrete latent variable with 10 categories through all coupling layers. For text we use a Transformer Language Model conditioned on discrete latent variable with 4 categories using additive embedding for all tokens. The size of the discrete variables is equal to the number of classes in the underlying dataset. Small number of categories allows to compute exact marginal likelihood value and speed up computation of the The schematic representations of the architectures are provided in
cost matrix C in Algorithm |
Figure 2
Table 2 presents the results of training the latent variable models for CoLLike and MaL objectives averaged across 4 runs. Both objectives exhibit similar performance in terms of likelihood across datasets. However, other characteristics vary.
MI is high for CoLLike objective on every dataset. Furthermore, attains approximately maximal value for AG News and CIFAR. MI for MaL. objective ranges from zero to values significantly lower than those of CoLLike. Zero MI indicates posterior collapse cases, which are mainly observed in ELBO optimization and recently discovered by Lucas et al. (2019) for MaL applied to simple linear models. This experiment indicates important observation: posterior collapse can as well happen in deep latent variable models during optimization of exact MaL despite usually being corresponded to the structure of ELBO. Importantly, for the MNIST dataset, half of the experiments exhibits posterior collapse.
CoLLike exhibits near-zero aggregated KL for all experiments. It implics that the model joint marginal in the latent domain perfectly matches the prior. For MaL., aggregated KL is zero only for AG News dataset which also has uninformative factorized joint. For other datasets, aggregated posterior significantly deviates from the prior. We also note that for MNIST dataset, MaL puts all probability mass to a single category in half of the runs.
To estimate the quality of unsupervised classification, we perform the optimal assignment of latent categories to classes. For all cases except CoLLike objective on AG News dataset, the quality of the unsupervised classification is similar and is low. On AG News the unsupervised accuracy is exceptionally good. However, the variance of the proposed solution is relatively high. The standard deviation of the accuracy across runs 5.4 with the highest value of 87.1 and the lowest of 73.3. In the following section, we show that it is possible to achieve significantly higher unsupervised accuracy and lower variance by latent variable ensembling.
Overall, CoLLike clearly outperforms Mal. in the tractable likelihood setting. Moreover, it shows high unsupervised classification accuracy for text modality. For Mal., experiments depict a variety of possible failures from posterior collapse to degenerate aggregated posterior, which extends find- ings of (Lucas et al., 2019) to expressive models and exact likelihood. However, despite CoLLike producing informative latents in terms of MI, unsupervised classification might be challenging in these cases. We believe that the key to high-performance unsupervised classification should be in the right inductive biases in conditioning and probabilistic model type.
6.1 LATENT ENSEMBLING
To reduce the high variance of CoI.Like unsupervised classification accuracy and increase its ac- curacy we propose perform ensembling of multiple models trained on the same data but using different seeds at initialization. Although there is no correspondence between labels for latent vari- able models, we can try find the labels assignment based on the agreement between them. This approach is motivated by direct cluster ensembling (Boongoen & lam-on 2018). The agreement between two labels of different ensemble members the number of intersecting samples with those labels. To align the latents we iteratively find the assignment with the highest intersection between labels. Finally, we find the assignment between aligned latents and ground truth labels.
In our experiments, we use 8 models per ensemble and train 4 independent ensembles. The simplest ensembling method is averaging of the predictions. It increases the mean unsupervised accuracy from 82.1 to 84.5 and reduces the standard deviation from 5.4 to 1.7. We further significantly im- prove these results by utilizing the agreement score, which also used for alignment of the labels. We pick top-k models with highest maximum coherence across other models in the ensemble. Aver- aging predictions of those top-k models further increases accuracy to 86.6 and lowers the standard deviation to 0.2.
8

---

Under review as a conference paper at ICLR 2023
0.3
1PET
prompts
87.5
8.8
Coli
86.0
0.7
Samoles
86.2
AOOUEE
0.6
LM-BFF
sangles
95.4
DocSCAN
CoLL latent: enserble
0.4
DeBERTa v3 ensemble
prompts
8.3
DeBERTa
v3
Accuracy
size
(a) CoLLike ensemble unsupervised and semi- (b) Supervised DeBERTa v3 base CoLLike ensem- supervised approaches.
Figure 3: Comparison of ensembled CoLLike with supervised (a) and unsuperviised/few-sho meth- ods(b).
We compare CoLlike results with the following unsupervised and supervised approaches: PET and iPET (Schick & Schütze 202I). EFL (Wang al., 2021). LM-BFF (Gao et al. (02II, DocSCAN
Stammbach & Ash [201] DocSCAN is purely unsupervised, while other approaches rely on engineering multiple textual descriptions of classes (prompts) or labeled data. All methods
heavy pre-trained Transformers (Vaswani et al.) (0117) as an initialization, while in CoLLike we small 2-layer Transformer with random initialization. Figure 3a presents the comparison of the methods. CoLLike clearly outperforms both unsupervised and most of the supervised methods. To determine how much training data we need without, possibly laborious, prompt engineering we use DeBERTa v3 (He et al. 2021). We vary training set sizes from 32 to 2048 and apply additional ensembling of models with different initializations and train-validation splits. Figure [3b reveals that Like can be better alternative to labeling more than hundred samples, which, in turn, requires an extensive data analysis. Besides, note the high difference between the ensemble and the single model for small dataset sizes in a supervised setting, which is interesting result by itself.
DISCUSSION AND FUTURE WORK
In this work, we propose to switch from the MaL paradigm of matching only marginals in the ob- served domain to CoLLike paradigm of finding an exact target joint by selection from family of joints with desirable properties. Furthermore, we show that matching of marginals utilized by Mal. corresponds to a specific choice of target joint, which motivates such failures as posterior collapse and divergence between target and model marginals in the latent domain. We experimentally show the ability of CoLLike to learn useful representations. Connection of CoLLike with OT allows to borrow techniques from the latter. For instance, Sinkhorn Relaxation (Cuturi 2013) can be used to speed up the assignment problem. Investigation of alternatives complete likelihood for target lection is of special interest. The right inductive biases for inducing useful properties using CoLLike still to be discovered, least until want to get the desired without specifying what want. We believe that the further extension of CoLLike to high-dimensional latents would be exciting and challenging. Other lines of research can be devoted the application of other divergences to the constructed family of joint target candidates and extension of CoLLike to learnable priors.
REPRODUCIBILITY
To promote reproducibility open-source code the link is hidden for double-bling review. check the supplementary materials) Furthermore, we describe details flow architectures in Appendix
F1 and Transformers in Appendix F22 Along with models, describe details of the training

---

Under review as a conference paper at ICLR 2023
procedures and data pre-processing. We also devote special attention to setting all necessary seeds, including CUDA, and to removing stochasticity from the BPE tokenizer.
REFERENCES
Alexander A. Alemi, Ben Poole, Ian Fischer, Joshua V. Dillon, Rif A. Saurous, and Kevin Mur- phy. Fixing a broken ELBO. In Jennifer G. Dy and Andreas Krause (eds.). Proceedings of the 35th International Conference on Machine Learning, ICML 2018, Stockholmsmüssan, Stockholm, Sweden, July 10-15, 2018, volume 80 of Proceedings of Machine Learning Research, pp. 159- 168. PMLR, 2018. wtt.tttt:/ httph:ttppprceeeiinngge
em 18a htmm
Yuki Markus Asano, Christian Rupprecht, and Andrea Vedaldi. Self-labelling via simultaneous clustering and representation learning. In 8th International Conference on Learning Represen- tations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net, 2020. URL
ew. net htt///forummiifi/ee -
Matthias Bauer and Andriy Mnih. Resampled priors for variational autoencoders. In Kama- lika Chaudhuri and Masashi Sugiyama (eds.), The 22nd International Conference Artifi- cial Intelligence and Statistics, AISTATS 2019, 16-18 April 2019, Naha, Okinawa, Japan, vol- ume 89 of Proceedings of Machine Learning Research, pp. 66-75. PMLR, 2019. |RLhttp: //proceedings mlr. prwss http://////baar19.......
K.P. Bennett, P.S. Bradley, and A. Demiriz. C clustering Technical Report MSR- TR-2000-65, May 2000. URL httts::/www.niiooooo//ot.oooooot ctm/ en- uan-ss/ze wss.usearch/
Piotr Bojanowski and Armand Joulin. Unsupervised learning by predicting noise. In Doina Pre- cup and Yee Whye Teh (eds.), Proceedings of the 34th International Conference on Machine Learning, ICMI. 2017, Sydney, NSW, Australia, 6-11 August 2017, volume 70 of Proceedings of Machine Learning Research, pp- 517-526. PMLR, 2017. URL htt.hhttp///gt http::://prooceediggg// ngs. mlr. wil..pres///e v 70/ /bo Janowsk 1 7a .html
Tossapon Boongoen and Natthakan lam-on. Cluster ensembles: A survey of approaches with recent extensions and applications. Comput. Sci. Rev., 28:1-25, 2018. doi: co016//.osree.2218... 003. Uttptttss:
//oi///o////o htwp//........... 1016/111i.... cosrev. 2018 8. 01 003
Samuel R. Bowman, Luke Vilnis, Oriol Vinyals, Andrew M. Dai, Rafal Józefowicz, and Samy Bengio. Generating sentences from a continuous space. In Yoav Goldberg and Stefan Rie- zler (eds.), Proceedings of the 20th SIGNLL Conference on Computational Natural Language Learning, CoNLL 2016, Berlin, Germany, August 11-12, 2016, pp. 10-21. ACL, 2016. doi: 10.18653/v1/kke 1002.
Thomas M. Cover and Joy A. Thomas. Elements of information theory (2. ed.). Wiley, 2006. ISBN 978-0-471-24195-9, URL http: wwwww/wwww...oo.w.wwo.... wow.eeeetooooooooooooooottaaaaaa www.:oorrmmttttttoo....o.m.mm.
Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. In Christopher J. C. Burges, Léon Bottou, Zoubin Ghahramani, and Kilian Q. Weinberger (eds.), Advances Neural Information Processing Systems 26: 27th Anmual Conference on Neural Information Processing Systems 2013. Proceedings of a meeting held December 5-8, 2013, Lake Tahoe, Nevada, United States, pp. 2292-2300, 2013. URL
neu.eerdinnnnnnnnnnne....e.....e wwwriiuuppp.... CC/ cc/paper/ htmm
a SM12222227|222222220022222233333
A. P. Dempster, N. M. Laird, and D. B. Rubin. Maximum likelihood from incomplete data via the em algorithm. Journal of the Royal Statistical Society. Series B (Methodological), 39(1):1-38, 1977. ISSN 00359246. URL htRp:t/p: / //ww/ /wwwwwww. www.ostor......oooo/o/oo
jstor. orgg http///abbeeeee22ee//e/e
Tianyu Gao, Adam Fisch, and Dangi Chen. Making pre-trained language models better few-shot learners. In Chengqing Zong, Fei Xia, Wenjie Li, and Roberto Navigli (eds.), Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th Interna- tional Joint Conference on Natural Language Processing, ACL/IJCNLP 2021. (Volume 1: Long Papers), Virtual Event, August 1-6, 2021, pp. 3816-3830. Association for Computational Lin- guistics, 2021. doi: 10.18653/v 00553//1/202111 1/55:l/1/22l. lon. URL https: //doi. orwpoo..rg/00.......
v1/2021 acl ong. 295
10

---

Under review as a conference paper at ICLR 2023
Junxian He, Daniel Spokoyny, Graham Neubig, and Taylor Berg- Kirkpatrick. Lagging inference networks and posterior collapse in variational autoencoders. In 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA. May 6-9, 2019. OpenReview.net, 2019.
ew net/ fww. wet. d=ry DFCCFF
Pengcheng He, Xiaodong Liu. Jianfeng Gao, and Weizhu Chen. Deberta: decoding-enhand enhanced bert with disentangled attention. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021. URL https: //openreview net/ffruum ( d
Matthew D Hoffman and Matthew Johnson. Elbo surgery: yet another way to carve up the varia- tional evidence lower bound. In Workshop in Advances in Approximate Bayesian Inference, NIPS, volume 1, 2016.
Ferenc Huszár.
Unsupervised learning by predicting noise: an information max-
imization view. sttps://www. http:///s/ww..iiieeeeeeooooaaaaa Leanini
orge:pacce...aaee.a....o.ooooooa 2022-11-17.
2017. Accessed:
Yeonwoo Jeong and Hyun Oh Song. Learning discrete and continuous factors of data via alter- nating disentanglement. In Kamalika Chaudhuri and Ruslan Salakhutdinov (eds.), Proceedings of the 36th International Conference on Machine Learning. ICML 2019, 9-15 June 2019, Long Beach, California, USA, volume 97 of Proceedings of Machine Learning Research, pp. 3091- 3099. PMLR, 2019. URL UtL/http://proceeee http://roceeddnn.........
ngs. ml wwwps/sss99 jeongl 9d.html
Aditya Jitta and Arto Klami. On controlling the size of clusters in probabilistic clustering. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 32, 2018.
Michael Jordan, Zoubin Ghahramani, Tommi S. Jaakkola, and Lawrence K. Saul. An introduction to variational methods for graphical models. Mach. Learn. 37(2):183-233, 1999. doi: 10.1023/A: 1007665907178. URL httpp:::: /ttp//ii.g/////e//////oeooeoeett 11.x///10.0
10. 10l. //2 00000//2100007577777898
Hyunjik Kim and Andriy Mnih. Disentangling by factorising. In Jennifer G. Dy and Andreas Krause (eds.), Proceedings of the 35th International Conference on Machine Learning, ICML 2018, Stockholmsmässan. Stockholm, Sweden, July 10-15, 2018, volume 80 of Proceedings of Machine Learning Research, pp. 2654-2663, PMLR, 2018. URL
mlr. prwss/eeesp..... v80 im1 8b html
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Yoshua Bengio and Yann LeCun (eds.), ) 3rd International Conference on Learning Representations, ICLR 2015, San Diego CA, USA, May 7-9, 2015, Conference Track Proceedings, 2015. URLhttp:
// //arxiv..co www.vorrgaaaa abs 1412. 6980
Diederik P. Kingma and Max Welling. An introduction to variational autoencoders. Found. Trends Mach. Learn., 3774333773322 2019. doi: 10. 10.1561/22000000066 URL
10. Nb.1561/22000 05561/220000005/
Alex Krizhevsky. Learning multiple layers of features from tiny images. Technical report, 2009.
Harold W. Kuhn. The hungarian method for the assignment problem. Naval Research Logistics Quarterly, 2(1-2):3397, 1955. doi:
URL https:
wi http:cco//ooooaa/
Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient-based learning applied to document recognition. Proc. IEEE, 86(11):278-2224, 1998. doi: 10.1109/5.72679. URL
http:://60..orr/100..1000
726791
Ralph Linsker. Self-organization in a perceptual network. Computer, 21(3):105-117, 1988. doi:
10. 10.1/09/..36. httpt://s::: /otp::::///oi..r/1..... 1t../....19//..........
12..109922.. /6
Stuart P. Lloyd. Least squares quantization in PCM. JEEE Trans. Inf. Theory, 28(2):129-136,
1982. doi: 00.1109/TT.11922.0555533 URL https: //ooi. wwg/10..100... 11000....011110 TTE.O///TTTTTTT 1982.
1056489

---

Under review as a conference paper at ICLR 2023
James Lucas, George Tucker, Roger B. Grosse, and Mohammad Norouzi. Don't blame the elbo! A linear VAE perspective on posterior collapse. Hanna M. Wallach, Hugo Larochelle, Alina Beygelzimer, Florence d' Aluch-BBuc, Emily B. Fox. and Roman Garnett (eds.), Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Pro- cessing Systems 2019, NeurlPS 2019, December 8-14, 2019, Vancouver, BC, Canada, pp-
9403-9413, 2019, URL httpss/pprroeeeeingg.... neurips............p............ wtt.:/eeerr.....ppprrrer.e..e Abstract. html
Yi Luo and Nima Mesgarani. Conv-tasnet: Surpassing ideal time finefreeyyncy magnitude masking for speech separation. IEEE ACM Trans. Audio Speech Lang. Process., (786) 126612266,,6, 2019. doi: 10. 10.10997AASLP22199291111
URL htttps://cci.oooo 1109TTAAAA 2019. 09/TAST/
2915167
J MacQueen. Classification and analysis of multivariate observations. In In 5-th Berkeley Symposium on Mathematical Statistics and Probability. pp. 281 281-297, 1967.
Alireza Makhzani, Jonathon Shlens, Navdeep Jaitly, and Ian J. Goodfellow. Adversarial autoen- 05644
coders. CoRR, abs/151 1.05644, 2015. URL
Andriy Mnih and Karol Gregor. Neural variational inference and learning in belief networks. In Proceedings of the 31th International Conference on Machine Learming, ICML 2014, Beijing, China, 21-26 June 2014, volume 32 of JMLR Workshop and Conference Proceedings, pp. 1791- 1799. JMLR.org. 2014. URL http: http://roceeeengge/eoooo http:://prcceee.ggsss.
mlr www.pppwwss....e. mmthll html
Andriy Mnih and Danilo Jimenez Rezende. Variational inference for monte carlo objectives. In Maria-Florina Balcan and Kilian Q. Weinberger (eds.), Proceedings of the 33nd International Conference on Machine Learning, ICMI. 2016, New York City, NY, USA, June 19-24, 2016, vol- ume 48 of JMLR Workshop and Conference Proceedings, pp. 2188-2196. JMLR. org, 2016. URL wwa.iig..m......e.. mmlr.presse / / mn www.mmnhhhlh.....mmmmmmmmmmmmmmm html
Kevin P. Murphy. probml.
Probabilistic Machine Learning: An introduction.
MIT Press, 2022. URL
Christos H. Papadimitriou and Kenneth Steiglitz.
Combinatorial Optimization: Algorithmss and Complexity. Prentice-Hall, 1982. ISBN 0-13-152462-3.
Giorgio Patrini, Rianne van den Berg, Patrick Forré, Marcello Carioni, Samarth Bhargav, Max Welling, Tim Genewein, and Frank Nielsen. Sinkhorn autoencoders. In Amir Globerson and Ricardo Silva (eds.), Proceedings of the Thirty-Fifth Conference on Uncertainty in Artificial In- telligence, UAI 2019, Tel Aviv, Isruel, July 22-25, 2019, volume 115 of Proceedings of Machine Learning Research, pp. 733-743, AUAI Press, 2019. URL greas/v115/patrinne
inwroeeedingg...... . mlr.
html
Ali Razavi, Aäron van den Oord, Ben Poole, and Oriol Vinyals. Preventing posterior collapse with delta-vaes. In 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019. OpenReview.net, 2019. URL htt.ps:://ppereviieeee
met/fooumm: 1d=3JeeenDDD...
Mihaela Rosca, Balaji Lakshminarayanan, and Shakir Mohamed. Distribution matching in vari- ational inference. CoRR, 802.0807200 2018. URL http:///ax.i..org/ab///// 06847
Timo Schick and Hinrich Schütze. Exploiting cloze-questions for few-shot text classification and natural language inference. In Paola Merlo, Jörg Tiedemann, and Reut Tsarfaty (eds.), Pro- ceedings of the 16th Conference of the European Chapter of the Association for Computa- tional Linguistics: Main Volume, EACL 2021, Online, April 19 23, 2021, pp. 255-269, As- sociation for Computational Linguistics, 2021. doi: 10. 1/ htttps://doi.
URL
Dominik Stammbach and Elliott Ash. Docscan: Unsupervised text classification via learning from neighbors. CoRR, 011221155.4422 2021. URL wtt.ttttts///aaxiv..tm
05.00.0.... 04024
12

---

Under review as a conference paper at ICLR 2023
Jakub M. Tomezak and Max Welling. VAE with vampprior. In Amos J. Storkey and Fer- nando Pérez-Crus (eds.). International Conference on Artificial Intelligence and Statistics, AIS- TATS 2018, 9-11 April 2018. Playa Blanca, Lanzarote, Canary Islands, Spain, volume 84 of Proceedings of Machine Learning Research, pp. 1214-1223. PMLR, 2018. URL http: //proceeding ngs ....m.m.. ml pt..as.......s.////............a
/ v8 ////4///mmzaaaa 8a html
George Tucker, Andriy Mnih, Chris J. Maddison. Dieterich Lawson, and Jascha Sohl-Dickstein. REBAR: low-variance, unbiased gradient estimates for discrete latent variable models. In Is- abelle Guyon, Ulrike von Luxburg, Samy Bengio, Hanna M. Wallach, Rob Fergus, S. V. N. Vish- wanathan, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA, pp. 2627-2636, 2017. URL htactps://procceedingg hwwp://cc.eeennn....r...........
neurips. .html
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Isabelle Guyon, Ulrike von Luxburg, Samy Bengio, Hanna M. Wallach, Rob Fergus, S. V. N. Vishwanathan, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA, pp. 5998 6008, 2017. URL http:ttttpp.tcced..nnnte iew..eeeeedding................. htcplipsscccpppppe pc/pppprr/01///aa.. 3FFa2233777c6111003ccccaaaaaa--3
html
Sinong Wang. Han Fang, Madian Khabsa, Hanzi Mao, and Hao Ma. Entailment as few-shot learner. CoRR, abs/2104. 14492104.14690, 2021. URL
www.aawwwwww.oooaaaaaaaaaaooeeae abs/ 2104 14690
Chen Xing, Devansh Arpit, Christos Tsirigotis, and Yoshua Bengio. A walk with SGD. CoRR, ab.//100. 2018.
http::////aaall.oooo//////////// http://b/118000 1ttp//11122000700)
Yang You, Jonathan Hseu, Chris Ying, James Demmel, Kurt Keutzer, and Cho-Jui Hsieh. Large- batch training for Istm and beyond. In Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis, SC '1,, New York, NY, USA, 2019. Association for Computing Machinery. ISBN 9781450362290. doi: 10.1145/295500.336611.. 3356137
URL https: / htt/oooo.oooooooo 110 10.
Yang You, Jing Li, Sashank J. Reddi, Jonathan Hseu, Sanjiv Kumar, Srinadh Bhojanapalli, Xiao- dan Song, James Demmel, Kurt Keutzer, and Cho-Jui Hsieh. Large batch optimization for deep learning: Training BERT in 76 minutes. In 8th International Conference on Learning Repre- sentations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net, 2020. URL Syx
net/
Dong Yu, Morten Kolbæk, Zheng-Hus Tan, and Jesper Jensen. Permutation invariant training of deep models for speaker-independent multi-talker speech separation. In 2017 IEEE International Conference on Acoustics, Speech and Signal Processing, ICASSP 2017, New Orleans, LA, USA, March 5-9, 2017, pp. 241-245. IEEE, 2017. doi: 10. 1A.10///ASSP.221.7992111... URL https:
7/doi aro.i.org/10 10. 1109 ICASSICCAPPPPPPP 2017. 7952154
Xiang Zhang. Junbo Jake Zhao, and Yann LeCun. Character-level convolutional networks for text classification. In Corinna Cortes, Neil D. Lawrence, Daniel D. Lee, Masashi Sugiyama, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 28: Annual Confer- ence on Neural Information Processing Systems 2015, December 7-12, 2015, Montreal, Que- bec, Canada, pp. 649-657, 2015, URL http:: ittp: http:///ccccceeei..ee/cc new.igg..e.ippssp cc/paper/ Abt.oat--aaaatt.t..a..oac htmi
Shengjia Zhao, Jiaming Song, and Stefano Ermon. Infovae: Balancing learning and inference in variational autoencoders. In The Thirty-Third AAAI Conference on Artificial Intelligence, AAAI
2019, The Thirty-First Innovative Applications of Artificial Intelligence Conference, IAAI 2019,
The Ninth AAAI Symposium on Educational Advances in Artificial Intelligence, EAAI 2019,
Honolulu, Hawaii, USA, January 27 - February 1, 2019, pp. 5885 5892. AAAI Press, 2019.
doi: 10. a0.1609/aaii.v3338 133101.3301118888 URLbttps: /oi: org 10. 1009/aai.. aaai. v33101.
3301 5885
13

---

Under review as a conference paper at ICLR 2023
A DERIVATION OF KL FORMS OF THE CONSIDERED OBJECTIVES
Equivalence between CoLLike objective 2 and its K divergence form from Table can be derived follows:
KLipstehpe(ell)llood..... -E (e log (s
Po Po(a; 2)
pa (ala) (log (x))))x((222)) log po (2,
(6)
C - (ep Po x)]
log (1
=C -C-LL00
where the first term in 66 is treated as constant with the assumption that all samples from ps (x) take distinct values, which reasonable for such high-dimensional objects images, texts, and sounds. Thus, the K L form of the objective is equivalent to the CoLLike objective up to a multiplicative factor and additive term. For proof of the constancy see Appendix B
The derivation of equivalence between and its KT. from Table is as follows
log ps (=))
K [Llpp(e))oo(lll]loooo..... = Ba.anps(e)no(all)
po (0( po
log px)
po pg(x).
Bewpp(u) (1) [log pa (I)] Ezxps(a) (logpo(t)) (x))
Σ log (z:)
N
1
C MMaa
The KL form of ELBO objective from Table D can be found in many works Zhao al. (2019);
Kingma & Welling (2019), however, we provide a derivation here to make the paper self-contained.
(L(ps(e)go((ll))loooo... E, mo log (=(2)ss<=12)
po (x
log 4o (2(zr) po (ze 2)
Exrpp((>) [logps(2))] E (np
log Po(x, 2)
=C Exppp(z) E Ezmnu(a(o))
(0(2)x)
N Σ CELBO(EN,
= =C
B ENTROPY AND MUTUAL INFORMATION OF EMPIRICAL JOINT
In this appendix derive useful properties of the empirical joint distributions produced by sampling from the prior. The joint distribution ps (=) (2(22) depends on æ. We focus on how π influences such distribution characteristics entropy and mutual information.
Consider a joint distribution over discrete & and z. This kind of distribution can be visualized as table, such depicted in Figure 4 If there multiple samples taking the same value both in
14

---

Under review as a conference paper at ICLR 2023
100 O Z 10000 O
z
Z
X
X
Figure 4: Example of three joint distributions with discrete and Lines the left and in the bottom depict the number of samples from empirical marginals with the corresponding value of the random variable. Squares reflect the joint probability value. Each line crossing the square corresponds to I/N probability added to the corresponding (z, z) random variable pair.
and z domain, the permutation can change the entropy of the joint. For instance, for the distribution
on the left, the entropy H (ps(())) (2))) distribution in the center, the entropy H (xs(2) (=)))=))( (2[2)) - (4- _ log 1.39 nats. So, depending on ı we might end up with more and less entropic distributions.
( log + log + log ()) ~ 1.04 nats. For the
However, when restrict any empirical marginal to take only distinct values, like in the right part of Figure 4 the situation changes. Namely, each distinet pair (r. 2) can be chosen at most once, because to choose it twice need duplicate sample in both domains. This be verified using the right part of Figure 4 Just try to construct a joint with some square having grater than one line assuming that each value æ marginal has only one line. Moreover, for æ, this is a reasonable assumption since usually the domain of is nigh-diimssional. For the case, the joint will contain N non-zero points each with probability 1/N. Thus, the entropy of the empirical distribution equal
- -N N log N = log N.
The observation above allows to easily derive mutual information of the empirical joint. Mutual information defined as
p(x, E Bx,xo((n.s) log p(z)p(2)
p(z[z)
MI(p(a,2)) 2)) E Ex,zop(n,s) log
p(2)
Under the assumption that pa()) contains only distinct elements the conditional par (2(())))) = for all values from the ps(x). So, choosing uniquely determines the value of 2, be seen from the right part of Figure 4 Then the mutual information is given by
Par (2|x) psn (2)
1 psx psx(2) (z)
(1(p) Be. log
(E) log
(1 (8)
So, the mutual information is always equal to the entropy of the empirical prior. It is possible to show that the value is the maximum possible one. This becomes obvious from the entropic factorization of the mutual information
MI(p(z,2) H(((z)) Expp(>) [H((2=(=))
Since the entropy non-negative, the mutual information can be decreased only through the second term of (9). which equals because value is completely determined by
When we try to extend the observations above to continuous cases we face the following challenge: empirical distribution has infinite values at the sample points. This drives the differential entropy as well as mutual information to infinity. However, adding noise to the empirical distribution solves this problem. Adding uniform noise with the interval smaller than the precision of the floating point makes the entropy finite and constant with respect to π. One can show that the resulting mutual information of the empirical joint also equals log N.
15

---

Under review as a conference paper at ICLR 2023
C WASSERSTEIN DISTANCE AND COLLIKE
Optimal Transport cost becomes Wasserstein distance when e is a metric. A very illustrative example from this family is equality of Wasserstein-2 (c is the Euclidean distance) and CoLLike for some setups. Specifically, the following objective can be produced both by Wasserstein distance and CoLLike
Lw (0) = min Σ (5. (o:=_o(=m(o)) fo(=)())) TEll
To get this objective from OT perspective define the model distribution be produced by passing fixed sample from prior through deterministie decoder fz(z). The result is empiric distribution x domain. Wasserstein distance between two empiric distributions is determined by optimal pairing between points from data distribution pa(x) and model distribution spanned by empiric latents. The same objective is produced by factorized Gaussian (=(22) and uniform prior p(z).
This connection demonstrates that the model (=]2)) defines both mapping from to domain and "topology" of the z space (how we measure distance between objects). However, approaches based on the Wasserstein distance are limited to continuous variables, while CoLLike is applicable both for discrete and continuous domains. Moreover, CoLLike provides probabilistic basis for the choice of the cost function.
D OPTIMAL PAIRING BY COMBINATORIAL OPTIMIZATION
Having at hand log-likelihood values for all possible 2; Σy pairs, we are ready find the optimal permutation. naive way to do so is to evaluate the sum 2 for every possible permutation æ. Despite need only sum different pre-computed values, the search space for is tremendous N!. However, we can cast this problem to combinatorial optimization one. Following Papadimitriou & Steiglitz (1982), the assignment problem is stated follows
minimize
Σ G3g0i.j
subjected to
Σ Oij Dai= =1 j=0,..,
Σ a.j =1 0,..,, N
aij E {0,
where Crj is the cost of picking the element i,j and a., is the indicator variable. The constrains of this problem define a set of permutation matrices. By choosing the cost to be negative log- likelihood and replacing indicator variable with permutation we end up with CoLLike objective. This combinatorial optimization problem can be solved efficiently with Hungarian algorithm Kuhn (1955) with complexity of ((N)).
E
THE COMPLEXITY OF THE LAP SOLVER
BigureBdepiict dependency between LAP problem size and time consumed by Hungarian algorithm to solve the problem. The input to the algorithm is matrix E BxB where B the size of the problem.
F MODELS DESCRIPTION
F.I NORMALIZING FLOWS
We use Glow-like normalizing flow for all image experiments. We choose the learning rate by start- ing from le 2 learning rate search was done. Below we provide details on the model parameters.
1 and gradually decrease until there is no instabilities during training. No extensive
CIFAR model:
16

---

Under review as a conference paper at ICLR 2023
Hungar lan LAP execution tine,
36-
10-
16-1
Size of the problem
Figure 5: Time to solve LAP with Hungarian algorithm for different sizes of the problem.
architecture: Glow
number of flows per scale: 7
flow coupling: affine
coupling ResNet, 3 blocks, hidden size 96
permutation flow: invertible convolution with lower-upper factorization
normalization flow: ActNorm
number of scales:
scale factor: 2 height, width
MNIST model (same as CIFAR except the number of scales):
architecture: Glow
number of flows per scale: 7
flow coupling: affine
coupling net: ResNet, 3 blocks, hidden size 96
permutation flow: invertible 1 convolution with lower-upper factorization
normalization flow: ActNorm
number of scales:
scale factor: height, width
During training we use marginal likelihood to validate both MaL and CoLLike. We did not search over possible optimizers, and use Adam (Kingma Ba] 2015) with default parameters. The vali- dation split is chosen to be 0.05 because no significant variations of the likelihood were observed during training. To summarize, we use the following parameters both for CoLLike and Mal.:
epochs: 256
learning rate: 56 MNIST; -CFARR
batch size: 64
validation part of the training set: 0.05
validation criterion: marginal likelihood
optimizer: Adam, 8 = (0.9, 0.999): 1e -8
As data pre-processing step, we used only dequantization with uniform noise, with the range equal to the quantization step.
17

---

Under review as a conference paper at ICLR 2023
F.2 TRANSFORMER
We used simple two-layer transformer across our experiments. The model description:
number of layers:
hidden size: 128
feedforward dimension: 128
embedding dimension: 128
number of attention heads: 4
number of embeddings: 4000
The training details are similar to CIFAR configuration:
epochs: 256
learning rate: 2e -4 CIFAR
batch size: 64
validation part of the training set: 0.05
validation criterion: marginal likelihood
optimizer: Adam, B (0.9, 0.999); e 1
In pre-processing step, we truncate the sequences longer than 192 tokens. Truncation affects less than 0.3% of the samples. Nevertheless, the tokenizer is trained the full-length sequences. Data pre-processing can be summarized as follows:
maximum length truncation: 192
BPE tokenization
vocabulary size: 4000
18


Reviews:
[{"review_id": "jqOmYSZJhG4", "reviewer": "Reviewer_vRN3", "summary": "", "questions": "", "limitations": "", "rating": 1, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The submission proposes a Complete Latent Likelihood (CoLLike) objective, which is given as (3), as an alternative to commonly used log-likelihood objectives (which the authors call Marginalized Likelihood, MaL) for learning latent variable models. In essence, when the prior over latent variables are known or given, the proposed method generates i.i.d. latent variables from the prior, and find a permutation between the generated latent variables and observed samples. \n\nThe motivation comes from scenarios where computing a marginalized probability by integrating over latent variables can be intractable, typically required in EM style algorithms. Instead, by finding the optimal pairing of observed samples and (generated and fixed) latent variables, which can basically be reduced to a bipartite matching (or assignment) problem, authors could avoid the intractable integration step. \n\nExperiments are performed on datasets with relatively small number of latent classes, and compare the performance of CoLLike and MaL. ", "summary_of_the_review": "I do not think the paper contains enough contribution or originality to be published. ", "strength_and_weaknesses": "\nThere are a number of controversial heuristics unjustified in my view. \n\n1. One basic principle of the MLE is its (asymptotic) consistency (which is why in the first place people want to do ML estimation). What can be said about the proposed algorithm? Is it consistent at least? \n\n2. The proposed algorithm is not a strict generalization of K-means, since empirical distribution of true labels does not necessarily match the true prior. Impact of such sampling errors should be considered (e.g., when clusters are very well-separated). \n\n3. For AG News data, the huge gap in Accuracy is suspicious. More thorough and detailed discussion on this phenomenon (on why it happens) should follow. \n\n\n\nOther weaknesses\n\n- Discussions in Section 5 are reasonable, but they are sort of trivial and do not add much value to the paper in my opinion. \n\n- What can be done if the prior of latent variables are not known in advance? \n\n- In the experiments, it is not clear what NLL / Agg. KL stand for (and also Accuracy is not precisely defined).\n\n- It is unusual to use $-\\log$ as a distance measure for OT. So why this should be better than using a typical Euclidean distance measure?\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "1: The contributions are neither significant nor novel.", "empirical_novelty_and_significance": "1: The contributions are neither significant nor novel."}, "extra_text": {}}, {"review_id": "gccyOWsfQ9", "reviewer": "Reviewer_G4Xy", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposes a complete latent likelihood (CoLLike), an alternative to the marginal likelihood (MaL) objective that is standard in ML algorithms (MaL) today. A systematic comparison between COLLike and Mal in the KL framework is compared. By discussing posterior collapse and high deviation of the aggregated posterior from the prior, authors show CoLLike can prevent these drawbacks. Empirical studies on image and NLL tasks show consistent improvement of CoLLIke over MaL. ", "summary_of_the_review": "a fundamental and potentially impactful work to address the training objective in ML algorithms today. ", "strength_and_weaknesses": "Strength:\nThis paper address a fundamental problem in all machine learning training and it potentially has a big impact. \nMany connection and differences with existing objective are discussed - it was enlightening.\n\nWeakness:\nSome concern on the permutation estimation exist due to inefficiency. \n\n\nComments:\n- \" target conditional is actually specified and equals $p_{\\theta} (z|x)$ if we ask what distribution we want to mimic\":  the exact form of $p_{\\theta} (z|x)$ is still not specified, and it could be anything since it canceled out. If $p_{\\theta} (z|x)$ is changed to any other distribution,  MaL  remains the same. This seems to contradict the statement. \n- \" the CoLLike objective allows learning models with reverse factorization\": ELBO has been used in VAE like structure, which resembles the function form. In addition, it would be good to provide examples on advantage or necessity of reverse factorization.\n- some references are not cited correctly. e.g., \" InfoMax principle Huszar (2017))\".\n- Eq 3: does the permutation mean all possible label assignments for all samples? How is it a permutation if the number of each class within Z are different? \n- \"since the set of all possible π values is countable and finite, optimization can be performed by an exhaustive search.\" it seems for large sample size N, this is not truly feasible with exhaustive search. Even with Hungarian complexity O(N^3), it is still quite a big number, in particular in large datasets popular nowadays. \n- Hence, it would be interesting to see the time comparison between different objectives in the experiments. \n- \" we emphasize our focus on learning useful z|x instead of simplifying the model with factorized x|z conditional\" can authors elaborate on this point? what if x|z is not factorized?\n- \"To reduce the high variance of CoLLike unsupervised classification\": authors should have discussed this point earlier of this work. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "SQYaPhMZsZN", "reviewer": "Reviewer_V5Wu", "summary": "", "questions": "", "limitations": "", "rating": 1, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The authors propose a new learning objective for fitting latent variable models. The basic idea is to sample a set of latent variables from the prior and find a permutation of those that aligns well with the observed data. The method provides a point estimate of the latent variables, but the paper is written from a perspective where it is suggested as an alternative to marginalising over the latent variables rather than focusing on the effect of the constraint imposed over standard joint likelihood. The approach is demonstrated primarily in context of supervised classification problems, where it is shown to find latent variables that align well with known classes.", "summary_of_the_review": "A bold attempt of re-defining the fundamental learning objective of latent variable models, but there are too many open questions regarding the theoretical foundations and also the empirical side. The biggest issue is that the approach is compared against marginal likelihood when it actually corresponds to a specific way of constraining a point estimate for the complete likelihood, and because of this mis-positioning the paper actually avoids discussing what the effect of sampling the set of possible latent variables from the prior is in practical cases or even in easy examples. From the perspective of general latent variable models the assumptions made for deriving the new objective are either shaky or even wrong, but it may be that for specific types of models they provide reasonable inductive biases. However, for that the paper would need to be re-written from a very different perspective.", "strength_and_weaknesses": "Strengths:\n- The paper focuses on a fundamental learning problem and hence has potential for significant impact if the contribution is considered valuable by the community\n- The core idea of constraining the possible values for the latent variables based on a sample from the prior is interesting\n- The learning algorithm seems to work well, in the sense that we observe empirically the properties we should expect from the objective (higher MI, lower KL between the prior and the conditional etc)\n\nWeaknesses:\n- The objective is poorly motivated, using only high-level argumentation like posterior collapse in some specific latent-variable models as justification for limitations of marginal likelihood, and it relies on possibly questionable assumptions (sharp distribution for the latents). The main story focusing on its relationship with marginal likelihood is odd.\n- The presentation is on a too high level. For instance, the main method is described only in terms of pseudo-code and the experiments are near impossible to parse from the main paper, for instance to figure out how the proposed method was used in these.\n- The computational cost of operating with permutations is dismissed lightly; the authors mention it as a challenge and outline possible solutions on a high level, but do not transparently communicate what it means in practice. The experiments completely omit all discussion on the computational cost or even whether the proposed speedups were used.\n- The notation and terminology are confusing and non-standard, making the paper difficult to read. For example, the phrases \"unknown $z|x$ target conditional\" or \"target $z|x$ distribution\" used in Introduction do not mean anything at that stage. Similarly, Section 3 is written in rather confusing terminology and makes standard concepts sound unjustified choices; the choice of conducting learning in a manner that is faithful to the modelling assumptions is called \"keeping the model posterior unchanged\" etc.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "1: The main claims of the paper are incorrect or not at all supported by theory or empirical results.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "Q7PER-SgQ4", "reviewer": "Reviewer_wzh3", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposes Complete Latent Likelihood (CoLLike) which utilizes permutation for matching in the latent domain. The authors argues that permuting the latents and then matching to inputs leads to expressive models with exact likelihood and no posterior collapsing, compared to MaL. Also, the suggested CoLLike can link likelihood and optimal transport.\n", "summary_of_the_review": "The paper is fairly-written in general. However, the authors are missing a critical point, which is an analysis on the computation complexity. Following from the authors, the provided algorithm has O(N^3) complexity, where N indicates the number of data instances. My major concern for this paper is that the utility of the proposed method may vary with computational time. Also, the clarity of the paper needs be improved. Therefore, at this moment, I'm leaning to the negative side on this paper.\n", "strength_and_weaknesses": "The experimental results show the efficacy of the suggested CoLLike.\n\nHowever, it seems that the computational complexity of O(N^3) can be problematic, where N indicates the number of data instances.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject