Neural Topic Model via Optimal Transport

1 INTRODUCTION . As an unsupervised approach , topic modelling has enjoyed great success in automatic text analysis . In general , a topic model aims to discover a set of latent topics from a collection of documents , each of which describes an interpretable semantic concept . Topic models like Latent Dirichlet Allocation ( LDA ) ( Blei et al. , 2003 ) and its hierarchical/Bayesian extensions , e.g. , in Blei et al . ( 2010 ) ; Paisley et al . ( 2015 ) ; Gan et al . ( 2015 ) ; Zhou et al . ( 2016 ) have achieved impressive performance for document analysis . Recently , the developments of Variational AutoEncoders ( VAEs ) and Autoencoding Variational Inference ( AVI ) ( Kingma & Welling , 2013 ; Rezende et al. , 2014 ) have facilitated the proposal of Neural Topic Models ( NTMs ) such as in Miao et al . ( 2016 ) ; Srivastava & Sutton ( 2017 ) ; Krishnan et al . ( 2018 ) ; Burkhardt & Kramer ( 2019 ) . Inspired by VAE , many NTMs use an encoder that takes the Bag-of-Words ( BoW ) representation of a document as input and approximates the posterior distribution of the latent topics . The posterior samples are further input into a decoder to reconstruct the BoW representation . Compared with conventional topic models , NTMs usually enjoy better flexibility and scalability , which are important for the applications on large-scale data . Despite the promising performance and recent popularity , there are several shortcomings for existing NTMs , which could hinder their usefulness and further extensions . i ) The training and inference processes of NTMs are typically complex due to the prior and posterior constructions of latent topics . To encourage topic sparsity and smoothness , Dirichlet ( Burkhardt & Kramer , 2019 ) or gamma ( Zhang et al. , 2018 ) distributions are usually used as the prior and posterior of topics , but reparameterisation is inapplicable to them , thus , complex sampling schemes or approximations have to be used , which could limit the model flexibility . ii ) A desideratum of a topic model is to generate better topical representations of documents with more coherent and diverse topics ; but for many existing NTMs , it is hard to achieve good document representation and coherent/diverse topics at the same time . This is because the objective of NTMs is to achieve lower reconstruction error , which usually means topics are less coherent and diverse , as observed and analysed in Srivastava & Sutton ( 2017 ) ; Burkhardt & Kramer ( 2019 ) . iii ) It is well-known that topic models degrade their performance severely on short documents such as tweets , news headlines and product reviews , as each individual document contains insufficient word co-occurrence information . This issue can be exacerbated for NTMs because of the use of the encoder and decoder networks , which are usually more vulnerable to data sparsity . To address the above shortcomings for NTMs , we in this paper propose a neural topic model , which is built upon a novel Optimal Transport ( OT ) framework derived from a new view of topic modelling . For a document , we consider its content to be encoded by two representations : the observed representation , x , a distribution over all the words in the vocabulary and the latent representation , z , a distribution over all the topics . x can be obtained by normalising a document ’ s word count vector while z needs to be learned by a model . For a document collection , the vocabulary size ( i.e. , the number of unique words ) can be very large but one individual document usually consists of a tiny subset of the words . Therefore , x is a sparse and low-level representation of the semantic information of a document . As the number of topics is much smaller than the vocabulary size , z is the relatively dense and high-level representation of the same content . Therefore , the learning of a topic model can be viewed as the process of learning the distribution z to be as close to the distribution x as possible . Accordingly , it is crucial to investigate how to measure the distance between two distributions with different supports ( i.e. , words to x and topics to z ) . As optimal transport is a powerful tool for measuring the distance travelled in transporting the mass in one distribution to match another given a specific cost function , and recent development on computational OT ( e.g. , in Cuturi ( 2013 ) ; Frogner et al . ( 2015 ) ; Seguy et al . ( 2018 ) ; Peyré et al . ( 2019 ) ) has shown the promising feasibility to efficiently compute OT for large-scale problems , it is natural for us to develop a new NTM based on the minimisation of OT . Specifically , our model leverages an encoder that outputs topic distribution z of a document by taking its word count vector as input like standard NTMs , but we minimise the OT distance between x and z , which are two discrete distributions on the support of words and topics , respectively . Notably , the cost function of the OT distance specifies the weights between topics and words , which we define as the distance in an embedding space . To represent their semantics , all the topics and words are embedded in this space . By leveraging the pretrained word embeddings , the cost function is then a function of topic embeddings , which will be learned jointly with the encoder . With the advanced properties of OT on modelling geometric structures on spaces of probability distributions , our model is able to achieve a better balance between obtaining good document representation and generating coherent/diverse topics . In addition , our model eases the burden of designing complex sampling schemes for the posterior of NTMs . More interestingly , our model is a natural way of incorporating pretrained word embeddings , which have been demonstrated to alleviate the issue of insufficient word co-occurrence information in short texts ( Zhao et al. , 2017 ; Dieng et al. , 2020 ) . With extensive experiments , our model can be shown to enjoy the state-of-the-art performance in terms of both topic quality and document representations for both regular and short texts . 2 BACKGROUND . In this section , we recap the essential background of neural topic models and optimal transport . 2.1 NEURAL TOPIC MODELS . Most of existing NTMs can be viewed as the extensions of the framework of VAEs where the latent variables can be interpreted as topics . Suppose the document collection to be analysed has V unique words ( i.e. , vocabulary size ) . Each document consists of a word count vector denoted as x ∈ NV and a latent distribution over K topics : z ∈ RK . An NTM assumes that z for a document is generated from a prior distribution p ( z ) and x is generated by the conditional distribution pφ ( x|z ) that is modelled by a decoder φ . The model ’ s goal is to infer the topic distribution given the word counts , i.e. , to calculate the posterior p ( z|x ) , which is approximated by the variational distribution qθ ( z|x ) modelled by an encoder θ . Similar to VAEs , the training objective of NTMs is the maximisation of the Evidence Lower BOund ( ELBO ) : max θ , φ ( Eqθ ( z|x ) [ log pφ ( x|z ) ] −KL [ qθ ( z|x ) ‖ p ( z ) ] ) . ( 1 ) The first term above is the expected log-likelihood or reconstruction error . As x is a count-valued vector , it is usually assumed to be generated from the multinomial distribution : pφ ( x|z ) : = Multi ( φ ( z ) ) , where φ ( z ) is a probability vector output from the decoder . Therefore , the expected log-likelihood is proportional to xT log φ ( z ) . The second term is the Kullback–Leibler ( KL ) divergence that regularises qθ ( z|x ) to be close to its prior p ( z ) . To interpret topics with words , φ ( z ) is usually constructed by a single-layer network ( Srivastava & Sutton , 2017 ) : φ ( z ) : = softmax ( Wz ) , where W ∈ RV×K indicates the weights between topics and words . Different NTMs may vary in the prior and the posterior of z , for example , the model in Miao et al . ( 2017 ) applies Gaussian distributions for them and Srivastava & Sutton ( 2017 ) ; Burkhardt & Kramer ( 2019 ) show that Dirichlet is a better choice . However , reparameterisation can not be directly applied to a Dirichlet , so various approximations and sampling schemes have been proposed . 2.2 OPTIMAL TRANSPORT . OT distances have been widely used for the comparison of probabilities . Here we limit our discussion to OT for discrete distributions , although it applies for continuous distributions as well . Specifically , let us consider two probability vectors r ∈ ∆Dr and c ∈ ∆Dc , where ∆D denotes a D − 1 simplex . The OT distance1 between the two probability vectors can be defined as : dM ( r , c ) : = min P∈U ( r , c ) 〈P , M〉 , ( 2 ) where 〈· , ·〉 denotes the Frobenius dot-product ; M ∈ RDr×Dc≥0 is the cost matrix/function of the transport ; P ∈ RDr×Dc > 0 is the transport matrix/plan ; U ( r , c ) denotes the transport polytope of r and c , which is the polyhedral set of Dr × Dc matrices : U ( r , c ) : = { P ∈ RDr×Dc > 0 |P1Dc = r , PT1Dr = c } ; and 1D is the D dimensional vector of ones . Intuitively , if we consider two discrete random variables X ∼ Categorical ( r ) and Y ∼ Categorical ( c ) , the transport matrix P is a joint probability of ( X , Y ) , i.e. , p ( X = i , Y = j ) = pij and U ( r , c ) is the set of all the joint probabilities . The above optimal transport distance can be computed by finding the optimal transport matrix P∗ . It is also noteworthy that the Wasserstein distance can be viewed as a specific case of the OT distances . As directly optimising Eq . ( 2 ) can be time-consuming for large-scale problems , a regularised optimal transport distance with an entropic constraint is introduced in Cuturi ( 2013 ) , named the Sinkhorn distance : dM , α ( r , c ) : = min P∈Uα ( r , c ) 〈P , M〉 , ( 3 ) where Uα ( r , c ) : = { P ∈ U ( r , c ) |h ( P ) ≥ h ( r ) + h ( c ) − α } , h ( · ) is the entropy function , and α ∈ [ 0 , ∞ ) . To compute the Sinkhorn distance , a Lagrange multiplier is introduced for the entropy constraint to minimise Eq . ( 3 ) , resulting in the Sinkhorn algorithm , widely-used for discrete OT problems .

Summary: The paper proposes a neural topic model derived from the perspective of optimal transport (OT). Topic embeddings are learned as part of the training process and is used to construct the cost matrix of the transport.  The cost function based on the OT distance is further improved by combining with the cross-entropy loss and by using the Sinkhorn distance to replace the OT distance.