Year: 2021

Abstract:
Self-supervised learning has significantly improved the performance of many NLP tasks. In this paper, we highlight a key advantage of self-supervised learning when applied to data generated by topic models, self-supervised learning can be oblivious to the specific model, and hence is less susceptible to model misspecification. In particular, we prove that commonly used self-supervised objectives based on reconstruction or contrastive samples can both recover useful posterior information for general topic models. Empirically, we show that the same objectives can perform competitively against posterior inference using the correct model, while outperforming posterior inference using mis-specified model.

heading: 1 INTRODUCTION
text: Recently researchers have successfully trained large-scale models like BERT (Devlin et al., 2018) and GPT (Radford et al., 2018), which offers extremely powerful representations for many NLP tasks (see e.g. Liu et al. (2021); Jaiswal et al. (2021) and references therein). To train these models, often one starts with sentences in a large text corpus, mark random words as “unknown” and ask the neural network to try to predict the unknown words. This approach is known as self-supervised learning. Despite many attempts, there is still no complete understanding of why self-supervised learning would lead to useful representations. A major bottleneck is that a “useful” representation often has no precise mathematical definition in the practical settings.
One approach to get a precise definition for “useful” representation is to apply self-supervised learning on synthetic data generated by probabilistic models. For example, many topic models (such as Latent Dirichlet Allocation (LDA) (Blei et al., 2003), Correlated Topic Model (CTM) (Blei & Lafferty, 2007), Pachinko Allocation Model (PAM) (Li & McCallum, 2006)) treat documents as a bag-of-words and use a two-step procedure of generating a document: first each document is viewed as a probability distribution of topics (often called the topic proportions w), and each topic is viewed as a probability distribution over words. If we use a matrixA to represent a topic-word matrix, where Aij represents the probability of using word i in topic j, then the words in this document will have probabilities according to Aw. Since the topic-word matrix A is shared across the entire corpus, and the topic distribution vector is specific to documents, it makes sense to definew as the representation for the document. In this setting, the questions we try to answer are: Can self-supervised learning learn a representation that contains information about w? Why is self-supervised learning better than the traditional approach of learning the probabilistic models and applying inference?
In this paper, we show that if the standard reconstruction-based objective (see Equation (1), similar to the objective used in Pathak et al. (2016); Devlin et al. (2018)) can be minimized, then the network will necessarily learn information about posterior of w (see Theorem 3). We also highlight one particular benefit of self-supervised learning – as long as the objective can be minimized, selfsupervised learning does not depend on what topic model was used to generate the data, while traditional inference methods rely heavily on having the correct model. Self-supervised learning is much more robust to model mis-specification.
Of course, the theory still needs to assume that either self-supervised learning objective can be successfully minimized (or at least approximately minimized). We demonstrate that this is indeed the case in practice in Section 5. In particular, we consider four different topic models (pure, LDA, CTM and PAM) and show that using the same reconstruction-based objective, it can extract information about the topic proportion w. Our experiments show that the performance of the self-supervised
model is competitive against posterior inference using the correct model, while outperforming posterior inference using mis-specified model. However, our experiments also show that it can be beneficial to adapt the architecture of the network to the complexity of the topic model (i.e., larger/more complex networks for CTM/PAM compared to pure and LDA models). Therefore if we take optimization into consideration, self-supervised learning is still more robust but not completely oblivious to the exact topic model. Finally in Section 6, we also demonstrate that minimizing this objective can lead to meaningful representations for real datasets (which are not generated by topic models).

heading: 1.1 RELATED WORKS
text: Self-Supervised Learning Self-supervised learning recently has been shown to be able to learn useful representation, which is later used for downstream tasks. See for example Bachman et al. (2019); Caron et al. (2020); Chen et al. (2020a;b;c); Grill et al. (2020); Chen & He (2021); Tian et al. (2020a); He et al. (2020) and references therein. In particular, Devlin et al. (2018) proposed BERT, which shows that self-supervised learning has the ability to train large-scale language models and could provide powerful representations for downstream natural language processing tasks.
Theoretical Understanding of Self-Supervised Learning Given the recent success of selfsupervise learning, many works have been tried to provide theoretical understanding on contrastive learning (Arora et al., 2019; Wang & Isola, 2020; Tosh et al., 2020; Tian et al., 2020b; HaoChen et al., 2021; Wen & Li, 2021; Zimmermann et al., 2021) and reconstruction-based learning (Lee et al., 2020; Saunshi et al., 2020; Teng & Huang, 2021). Also, several papers considered the problem from a multi-view perspective (Tsai et al., 2020; Tosh et al., 2021), which covers both contrastive and reconstruction-based learning. Moreover, Wei et al. (2020) and Tian et al. (2021) studied the theoretical properties of self-training and the contrastive learning without the negative pairs respectively. Saunshi et al. (2020) investigated the benefits of pre-trained language models for downstream tasks. They showed that if one can predict the next word accurately, then the learned representation can be used to solve downstream linear task. Our setting is different from theirs as we focus on the probabilistic models. Most relevant to our paper, Tosh et al. (2020) considered the contrastive learning in the topic models setting. Our theoretical results extend their theory to reconstructionbased objective (while also removing some assumptions), and our empirical results show that the reconstruction-based objective can be effectively minimized.
Theoretical Analysis of Topic Models Topic models are popular probabilistic models that has proved successful in many applications while still have clear structure to allow theoretical studies. Many works have proposed provable algorithms to learn topic models, such as method of moment based approaches (Anandkumar et al., 2012; 2013; 2014; 2015) and anchor word based approaches (Papadimitriou et al., 2000; Arora et al., 2012; 2016a; Gillis & Vavasis, 2013; Bittorf et al., 2012). Much less is known about provable inference for topic models. Sontag & Roy (2011) showed that MAP estimation can be NP-hard even for LDA model. Arora et al. (2016b) considered approximate inference algorithms.

heading: 1.2 OUTLINE
text: We first introduce the basic concepts of topic models and our objectives in Section 2. Then in Section 3 we prove guarantees for the reconstruction-based objective. Section 4 connects the contrastive objective to reconstruction-based objective which allows us to prove a stronger guarantee for the former. We then demonstrate the ability of self-supervised learning to adapt to different models in synthetic experiments in Section 5. Finally, we also evaluate the reconstruction-based objective on real-data to show that it extracts high-quality representations.

heading: 2 PRELIMINARIES
text: In this section we first introduce some general notations. Then we briefly describe the topic models we consider in Section 2.1. Finally we define the self-supervised learning objectives in Section 2.2 and give our main results.
Notation We use [n] to denote set {1, 2, . . . , n}. For vector x ∈ Rd, denote ‖x‖ as its `2 norm and ‖x‖1 as its `1 norm. For matrix A ∈ Rm×n, we use Ai ∈ Rm to denote its i-th column. When
matrix A has full column rank, denote its left pseudo-inverse as A† = (A>A)−1A>. For matrix or general tensor T ∈ Rd1×···×dl , we use vector vec(T ) ∈ Rd1···dl to represent its vectorization. Let ∆K to denote K − 1 dimensional probability simplex. For two probability vectors p, q, define their total variation (TV) distance as TV(p, q) = ‖p− q‖1 /2.

heading: 2.1 TOPIC MODELS
text: As we mentioned earlier, the generative process of many topic models can be viewed as a two-step procedure. We consider a general scenario where V is a finite set of vocabulary with size V and K is a set of K topics, where each topic is a distribution over V . We denote the topic-word matrix as A with Aij = P(word i | topic j), so that each column in A represents the word distribution of a topic. Each document corresponds to a convex combination of different topics with proportions w. For each word in the document, first sample a topic z according to the topic proportions w, and then sample the word from the corresponding topic (equivalently, one can also sample a word from distribution Aw). Different topic models differ in how they generate w, which we formulate in the following definition:
Definition 1 (General Topic Model). A general topic model specifies a distribution ∆(K) for each number of topics K. Given a topic matrix A and ∆(K), to generate a document, one first sample w ∼ ∆(K) and then sample each word in the document from the distribution Aw.
Here ∆(K) is a prior distribution of w which is crucial when we are trying to infer the true topic proportions w given the words in a document. Of course, different topic models may also specify different priors for the topic matrix A. However, given a large number of documents generated from the topic model, in many settings one can hope to learn the topic matrix A and the prior distribution ∆(K) (see e.g., Arora et al. (2012; 2016a)), therefore we consider the following inference problem:
Definition 2 (Topic Inference). Given a topic matrix A, prior distribution ∆(K) for topic proportions w, and a document x, the topic inference problem tries to compute the posterior distribution w|x,A given the prior w ∼ ∆(K).
Note that our general topic model can capture many standard topic models, including pure topic model, LDA, CTM and PAM.
Semi-supervised Learning Setup As our goal is to understand the representation learning aspect of self-supervised learning, we consider a semi-supervised setting where each document also has a label y which is determined as a function of w. We are given large number of unlabeled documents and a small number of labeled documents, the goal is to apply self-supervised learning on the unlabeled documents to get a good representation, and then use this representation to predict the label y. Of course, if one knows the actual parameters of the model, the best way predictor for y would be to first estimate the posterior distribution w and then apply the function that maps w to label y. We will show that for functions that are approximable by low degree polynomials self-supervised learning can always provide a good representation.

heading: 2.2 SELF-SUPERVISED LEARNING
text: In general, self-supervised learning tries to turn unlabeled data into a supervised learning problem by hiding some known information. There are many ways to achieve this goal (see e.g., Liu et al. (2021); Jaiswal et al. (2021)). In this paper, we focus on two different approaches for self-supervised learning: reconstruction-based objective and contrastive objective. We first formally define the objective and then discuss our corresponding result.
Reconstruction-Based Objective One common approach of self-supervised learning is to first mask part of the data and then try to find a function f to reconstruct the missing part given the unmasked part of input. This is commonly used for language modeling. In the context of topic modelling, since each word is i.i.d. sampled, given the document xunsup that generated from word distribution Aw, we pick t random words from the documents and mark them as unknown, then we ask the learner to predict these t words given the remaining words in the documents. Specifically, let y be the t words that we select and let x be the document with these t words removed, we aim to
select a predictor f that minimizes the following reconstruction objective:
min f Lreconst(f) , Ex,y[`(f(x), y)], (1) where `(ŷ, y) = ∑ k −yk log ŷk is the cross entropy loss. Here we slightly abuse the notation to use y as an one-hot vector for |V|t classes and f also outputs a probability distribution on Vt. Depending on the context, we will use y to denote either the actual next t words or its corresponding one-hot label. Now we are ready to present our result for the reconstruction-based objective: Theorem 1. (Informal) Consider the general topic model setting as Definition 1, suppose function f minimizes the reconstruction-based objective (1). Then, any polynomial P (w) of the posterior w|x,A with degree at most t can be represented by a linear function of f(x).
The theorem shows that if we want to get basic information about posterior distribution (such as mean, variance), then it suffices to predict a small constant number of words (1 for mean and 2 for variance). See Theorem 3 for the formal statement.
Contrastive Objective Another common approach in self-supervised learning is the contrastive learning. In contrastive learning, the training data is usually a pair of data (x, x′) with a label y ∈ {0, 1}, where label 1 means (x, x′) is a positive sample (x and x′ are similar) and label 0 means (x, x′) is a negative sample (x and x′ are dissimilar). The task is to find a function f such that it can distinguish the positive sample and negative sample, i.e., f(x, x′) = y. Formally, we want to select a predictor f such that it minimizes the following contrastive objective:
min f Lcontrast , Ex,x′,y[`(f(x, x′), y)]. (2)
In the context of topic models, following previous work (Tosh et al., 2020), we generate the data (x, x′, y) as follows. We first generate a document x from the word distribution Aw, then (i) with half probability we generate t words from the same distribution Aw to form the document x′ and set y = 1; (ii) with half probability we generate t words from a different word distribution Aw′ with w′ ∼ ∆(K) (so that w 6= w′) and set y = 0. For loss function `, we consider the square loss `(ŷ, y) = (ŷ − y)2. However, our results hold for other loss function such as logistic loss. We now give our informal result on contrastive objective. See Theorem 5 for the formal statement. Note that this theorem generalizes Theorem 3 in Tosh et al. (2020). Theorem 2. (Informal) Consider the general topic model setting as Definition 1, suppose function f minimizes the contrastive objective (2). Then we can use f and enough documents to construct a representation g(x) such that any polynomial P (w) of the posterior w|x,A with degree at most t can be represented by a linear function of g(x).

heading: 3 GUARANTEES FOR THE RECONSTRUCTION-BASED OBJECTIVE
text: In this section, we consider the reconstruction-based objective (1) and provide theoretical guarantees for its performance. We first show that if such objective with t unknown words can be minimized, then any polynomial of topic posterior w|x,A with degree at most t can be represented by a linear function of the learned representation. Theorem 3 (Main Result). Consider the general topic model setting as Definition 1, suppose topic matrix A satisfies rank(A) = K and function f minimizes the reconstruction-based objective (1). Then, any polynomial P (w) of the posterior w|x,A with degree at most t is linear in f(x), that is there exists a θ ∈ RV t such that for all documents x
Ew[P (w)|x] = θ>f(x).
To understand this theorem, we can first think about a warm-up example where t = 1 (see Section A.1 in appendix). Intuitively, in this case the best way to predict the missing word is to estimate the topic proportion w, and then predict the word using Aw where A is the topic matrix. Therefore, if the output of self-supervised learning f(x) (which is a V -dimensional vector indexed by words) minimizes the loss, then f(x) must have the form f(x) = AE[w|x]. When A is full rank multiplying by the pseudo-inverse of A recovers the expectation of the posterior. The proof for the general case of predicting t-words requires more careful characterization of the optimal prediction and its relationship to Ew[P (w)|x], which we defer to Section A.2.
Robustness for an approximate minimizer In Theorem 3 we focus on the case when function f is exactly the minimizer of reconstruction-based objective (1), i.e., Lreconst(f) = L∗reconst , minf Lreconst(f). However, in practice one cannot hope to find such a function exactly. In the following, we provide a robust version of Theorem 3 such that it allows us to find an approximate solution instead of the exact optimal solution.
To present our result, we need to first introduce the following notion of condition number, which was used in many previous works, such as collaborative filtering system (Kleinberg & Sandler, 2008) and topic models (Arora et al., 2016b). Intuitively, κ(B) measures how large a vector would change after multiplying with B in the `1 norm sense.
Definition 3 (`1 Condition Number). For matrix B ∈ Rm×n, define its `1 condition number κ(B) as κ(B) , minx∈Rn:x 6=0 ‖Bx‖1 / ‖x‖1 = maxi∈[n] ‖Bi‖1 , where Bi is the i-th column of B.
LetWpost ∈ RK×...×K be the topic posterior tensor for the t unknown words y = (y1, . . . , yt) given document x. That is, for each entry [Wpost]z1,...,zk = P(zi is the topic of word yi for i ∈ [t]|x) = Ew[wz1 . . . wzt |x] and Wpost = Ew[w⊗t|x].1 Therefore, any polynomial P (w) of degree at most t can be represented as Ew[P (w)|x] = β>vec(Wpost) for some β, where vec(Wpost) ∈ RK t
is the vectorization of Wpost.
We now are ready to present the robust version of Theorem 3. It shows that if we only find a function f whose loss is at most larger than the optimal loss, then a linear transformation of our learned representation can still give a good approximation of the target polynomial within a O( ) error.
Theorem 4 (Robust Version). Consider the general topic model setting as Definition 1, suppose topic matrix A satisfies rank(A) = K and function f satisfies Lreconst(f) ≤ L∗reconst + for some > 0. Then, any polynomial P (w) of the posterior w|x,A with degree at most t is approximately linear in f(x), that is there exists a θ ∈ RV t such that
Ex [( Ew[P (w)|x]− θ>f(x) )2] ≤ 2 ‖β‖2 κ2t(A†) , where Ew[P (w)|x] = β>vec(Wpost).
Note that the dependency on ‖β‖2 κ2t(A†) is expected, since this is the norm ‖θ‖2 that we would have if = 0, i.e., the θ we would have in Theorem 3. Thus, this quantity should be understood as the complexity of the target function for the downstream task. Empirically we show that κ(A†) is small in Section C.6. The proof is deferred to Section A.3.

heading: 4 GUARANTEES FOR THE CONTRASTIVE OBJECTIVE
text: In this section, we consider the contrastive objective (2) for self-supervised learning and provide similar provable guarantees on its performance as the reconstruction-based objective.
We use the same approach as Tosh et al. (2020) to construct a representation g(x) based on f(x, ·). Given a set of landmark documents {li}mi=1 with length |li| = t as references, the representation is defined as
g(x, {li}mi=1) = (g(x, l1), . . . , g(x, lm)) > , g(x, x′) =
f(x, x′)
1− f(x, x′) . (3)
The following theorem gives the theoretical guarantee for this representation. Similar to Theorem 3 for reconstruction-based objective, it shows that any polynomial P (w) of degree at most t can be representation by a linear function of the learned representation. Our proof here relies on the observation that having g(x, {li}mi=1) for all short landmark documents of length t gives similar information as f for the reconstruction objective. This observation allows us to remove the anchor words assumption needed in (Tosh et al., 2020).
1To simply the notations, WLOG we assume topics set K = [K] so that we can use topics zi ∈ [K] as indices.
Theorem 5. Consider the general topic model setting as Definition 1, suppose topic matrix A satisfies rank(A) = K and function f minimizes the contrastive objective (2). If we randomly sampled m = Kt different landmark documents {li}mi=1 and construct g(x, {li}mi=1) as (3), then any polynomial P (w) of the posterior w|x,A with degree at most t is linear in g(x, {li}mi=1), that is there exists a θ ∈ Rm such that for all documents x
Ew[P (w)|x] = θ>g(x, {li}mi=1).
In fact, one can use the same set of documents for both representation and downstream task so that we do not need additional landmark documents. We call this case as self-reference. The following corollary shows that solving the downstream task is equivalent to solve a kernel regression (in some sense the landmark documents are just random features for this kernel (Rahimi et al., 2007)). See more discussions in Section B.3.
Corollary 6. Denote {(xi, ỹi)}mi=1 as the downstream dataset, where ỹi = Ew[P (w)|xi] is the target for document xi. If we set landmarks {li}mi=1 to be {xi}mi=1, then solve the downstream task is equivalent to a kernel regression, that is
arg min θ m∑ i=1 ( ỹi − θ>g(x, {li}mi=1) )2 = arg min θ ‖ỹ −Gθ‖2 ,
where ỹ = (ỹ1, . . . , ỹm)> and G ∈ Rm×m is a kernel matrix such that Gij = g(xi, xj).

heading: 5 SYNTHETIC EXPERIMENTS
text: In this section, we optimize the reconstruction-based objective (1) on data generated by several topic models to show that the objective can adapt to different priors and recover topic posterior accurately.

heading: 5.1 TOPIC MODELS
text: We consider four types of topic models in our experiments. Our first topic model is the puretopic model, where each document’s topic comes from a discrete uniform distribution over the Kdimensional linear basis, namely {e1, e2, ..., eK}. Our second topic model is the Latent Dirichlet Allocation (LDA) model, where ∆(K) is a symmetric Dirichlet distribution Dir(1/K).
We are also interested in topic models that involve more subtle topic correlations and similarities between different topic’s word distributions. To this end, we consider the Correlated Topic Model (CTM) (Blei & Lafferty, 2007) and the Pachinko Allocation Model (PAM) (Li & McCallum, 2006). Our goal is to construct settings where the correlation between topics provide useful information for inference. To achieve this goal, in CTM, we construct groups of 4 topics. Within the group, topic pairs (0,2) and (1,3) are highly correlated (as specified by the prior), while topics 0 and 1 share many words (see Figure 1 for an illustration of the setting). In this case, if we observe a document with words that could either belong to topic 0 or 1, but this same document also has words that are associated with topic 2, we can infer that the first set of words are likely from topic 0, not topic 1. We construct such correlations in CTM model by setting the diagonal entries of its Gaussian covariance matrix to 15 and the covariance between correlated pairs of topics to 0.99 times the diagonal entries, with the remaining entries set to zero. We construct similar examples for Pachinko Allocation Model (PAM), see Section C in appendix for details.

heading: 5.2 SIMULATION SETUP
text: Document generation Our documents are synthetically generated through the following steps: we first construct a V ×K topic-word matrix A. Then, for each document, we determine the document length n from Poisson distribution Pois(λ), draw a topic distribution w from ∆(K), and draw n words i.i.d. from the word distribution given by Aw. In our simulation, we set K = 20, V = 5000, λ ∈ {30, 60}. Often A will be drawn from a dirichlet distribution Dir(α/K), and we take α = 1, 3, 5, 7, 9 to vary the difficulty level of the inference problem, where a larger α introduces more similarity between topics.
Neural network models We transform bag-of-word document representation into full document by repeating each word by its frequency and concatenating them in random order; afterwards, the full document will be our model’s input. We find that fully-connected neural network with residual connections performs the best on recovering topic posterior distribution for pure-topic and LDA documents, and attention-based architecture (Vaswani et al., 2017) performs the best for recovering topic posterior distribution for CTM and PAM documents. More specifically, the attention-based model contains 8 transformer blocks where each transformer block consists of a 768-dimensional attention layer and a feed-forward layer, with residual connection applied around every block. The model’s second to last layer averages over the outputs of the last transformer block and the final layer projects it to a V -dimensional word distribution.
Training setup During training, we resample 60,000 new documents after every 2 epochs, and the total amount of training data varies from 720K documents to 6M documents. The loss function is the reconstruction-based objective (1) with t = 1, i.e., we want the model to predict one missing word. To reduce the variance during the training, in each training document, the last 6 words are chosen as the prediction target and hidden from the model. The trained model is evaluated on test documents generated from the same topic prior as the training documents. We use 5,000 test documents for the pure-topic prior and 200 test documents for the remaining priors. For each test document, we use a Markov Chain Monte Carlo No U-Turn Sampler (Salvatier et al., 2016) assuming the document’s correct topic prior to approximate the ground truth expected topic posterior distribution. Then, we measure topic posterior recovery loss as the Total Variation (TV) distance between the recovered topic posterior and the ground truth topic posterior.

heading: 5.3 RESULTS
text: Topic posterior recovery loss As illustrated in Table 1, after 200 training epochs, our model can accurately recover the topic posterior distribution. Meanwhile, it can be observed that for larger values of the Dirichlet hyperparameter α, the recovery loss gets higher. This is expected because higher α leads to more similar topics, which makes the learning problem more difficult. This effect is also captured by Definition 3 and we show the computed condition numbers in Section C.6.
.
Major topic recovery We examine the extent our recovered topic posterior captures the major topics in a given document. For pure-topic and LDA documents, we measure major topic recovery probability of correctly estimating the topic with largest proportion. Considering that topics in CTM and PAM documents are correlated in pairs, we measure the major topic recovery rate for CTM and PAM as the top-2 topic overlap rate on their test documents. Table 1 shows that our algorithm is successful throughout all settings we’ve tried.
Robustness of self-supervised learning We compare the self-supervised approach to traditional topic inference. For each α value, we take 200 test documents from each category of documents, and we run posterior inference assuming a specific topic model and calculate the TV distance between this posterior and the ground truth posterior. We exclude assuming pure-topic prior from our comparison because it gives invalid results for documents with mixed topics. For α = 1, as shown in Table 2 Left, the topic posterior recovered from our approach is closer to the ground truth topic posterior than that recovered from a mis-specified topic prior.
We also test the major topic recovery rate for our model against topic inference using both correct and incorrect model (see Table 2 Right). The major topic recovery rate of our approach is similar to that of posterior inference assuming the correct prior. For pure topic model, all four methods get 100% major topics recovery rate on pure-topic documents; for LDA, again all four methods perform similarly well since the instance is not difficult. However we observe significant difference for more complicated CTM and PAM models. For these models, the self-supervised learning approach performs similarly to posterior inference using the correct model, and both of them are significantly better than posterior inference using mis-specified models. The results for α = 3, 5, 7, 9 are presented in Section C. The comparison further reveals the robustness of self-superivised learning.

heading: 6 EXPERIMENTS ON REAL DATA
text: Even though our theory suggests that self-supervised learning can learn useful representations for a general topic model, it is still unclear whether this simple approach learns any reasonable representation in practice as the documents are not really generated by any topic model. In this section, we show that the simple approach at least generates a better representation than several baselines.

heading: 6.1 EXPERIMENT SETUP
text: The dataset we used is the AG news dataset (Zhang et al., 2015), in which each document has a label of one out of four categories: world, sports, business, and sci/tech. Thus, the task falls into pure topic model. We generally follow experiment setup done by Tosh et al. (2020) (see more details in Section D.1), but the representation we used is generated by our reconstruction-based objective.
Our training involves two major phases: unsupervised phase and supervised phase. In the unsupervised phase, we trained our self-supervised model on most of our data that is later used for generating document representations; and then in the supervised phase, we trained a linear classifier
using multi-class Logistics Regression, implemented by Scikit-learn(Pedregosa et al. (2011)), with 3-fold Cross Validation for parameter tuning (l2 regularization term and solver).
We chose residual blocks as the basic building block for our neural network architecture, with varying width and depth. The network used for Figure 2 has 3 residual blocks and a width of 4096 neurons per layer. We explore other depths/weight combinations in Section D.3 in Appendix. We used AMSGrad optimizer (Reddi et al., 2019) with weight decay of 0.01. We resample our reconstructionbased document-label pair every 2 epochs. We trained the model using reconstruction-based objective (1) with t = 1. Using the same variance-reduction technique in synthetic experiments, in actual training, we sampled 4 words as our labels and the loss is an average of the cross entropy loss of each word evaluated separately.
To extract our representation, we use an identity matrix to replace the last layer. That is, we take a softmax function on top of the second-to-last layer of the neural network. This is different from our theory but we found that it effectively reduces the high output dimension (equal to vocabulary size, around 16,700 in our experiments) and improves performance. We included more details in Section D.2 in Appendix.

heading: 6.2 EXPERIMENT RESULT
text: The performance of our Reconstruction-based learning representation (RBL) is shown in Figure 2, where the test accuracy on representation is plotted against number of training samples used to train the classifier in supervised learning phase. Our representation RBL performs better than both BOW and Word2vec baselines. Notably, Word2vec representation is inferior only by a small margin, and works well in general even when limited training samples are provided. On the contrary, although BOW representation has a decent test accuracy when training samples are abundant, it performs significantly poorly on smaller training set.
We also compared our result to previous noise contrastive estimation (NCE) representation benchmark achieved by Tosh et al. (2020) using contrastive objective , where their reported best accuracy was around 87.5% when full 4000 training samples are used, slightly higher than our RBL corresponding test accuracy of 87.1%. In our attempt to reproduce their result, parameter tuning yields the best accuracy of 86% when full training samples were used. We plot our own reproduced results on in Figure 2 since parameter tuning did not achieve a close benchmark to their result.

heading: 7 CONCLUSION
text: “All models are wrong but some are useful.” If one self-supervised objective can capture all models, then it would be able to extract useful information. In this paper, we studied the self-supervised learning in the topic models setup and showed that it can provide useful information about the topic posterior no matter what topic model is used. Our results generalized previous work (Tosh et al., 2020) to both contrastive learning and reconstruction-based learning and our techniques allow us to depend on weaker assumptions. We also empirically showed that the reconstruction-based learning performs better than the posterior inference under mis-specified models, and it can provide useful representation for the topic inference problem. An immediate open problem is what other models (especially those that captures ordering of words) can be captured by the much more general selfsupervised learning approaches used in practice.

heading: REPRODUCIBILITY STATEMENT
text: For our theoretical results, we explain the detailed setup in Section 2, give the formal statement of the results in Section 3 and Section 4, and present the complete proofs in Section A and Section B. For our experiments, we report our results and general setup in Section 5 and Section 6 and explain detailed setup in Section C and Section D. We also include our codes that generate the results in the supplement material.

heading: A OMITTED PROOFS IN SECTION 3
text: In this section, we give the omitted proofs in Section 3. We first give a warmup example to illustrate our proof idea in Section A.1. Then we give the proof of our main result (Theorem 3) in Section A.2 and the proof of robust version (Theorem 4) in Section A.3.

heading: A.1 WARM-UP EXAMPLE: RECONSTRUCT THE UNKNOWN WORD
text: In this warm-up example, we consider the simple setting where we try to predict the only one unknown word of the document using reconstruction-based objective (1). The following result is the special case of Theorem 3 with t = 1, which shows that the learned representation is able to give any linear function of the topic posterior.
Theorem A.1. Consider the general topic model setting as Definition 1, suppose topic matrix A satisfies rank(A) = K and function f minimizes the reconstruction-based objective (1). Then, any linear function P (w) of w is linear in f(x), that is for any P (w) = β>w, there exists a θ ∈ RV such that for all documents x
Ew [ β>w|x ] = θ>f(x).
Following the proof idea described in Section 3, we first give a characterization of the optimal function f . The following lemma shows that f(x) must be the word posterior vector given the document x. The proof of this lemma is simply based on the property of the cross-entropy loss function and we defer it to Section A.1.1. Recall that our vocabulary set is V = {v1, . . . , vV } and xunsup = (x, y) is the given document, where x is unmasked part and y is word marked as unknown.
Lemma A.2. If f minimizes the reconstruction-based objective (1), then we have for all document x
f(x) = (P(y = v1|x), . . . ,P(y = vV |x))> .
Based on the above lemma, we are able to prove Theorem A.1. It is easy to see that the word posterior distribution is AEw[w|x], so we have f(x) = AEw[w|x]. Since the columns of A are linearly independent, we know Ew[w|x] = A†f(x). Thus, for any linear function β>w, there exists θ = (A†)>β such that Ew[β>w|x] = θ>f(x). The formal proof of the general case is given in Section A.2.

heading: A.1.1 PROOF OF LEMMA A.2
text: Instead of focusing on the t = 1 case, we directly give the corresponding lemma of Lemma A.2 for general t. Recall that xunsup = (x, y) is the given document, x is the unmasked part, y = (y1, . . . , yt) is the t unknown words that we want to predict and V = {v1, . . . , vV } is the set of vocabulary. It shows that the optimal f is the t words posterior distribution. Note that the words posterior distribution is linear in the topic posterior of t-th moment, which is useful for the later analysis.
Lemma A.3 (Lemma A.2, General Case). If f minimizes the reconstruction-based objective (1), then we have for all document x
f(x) = (P(y = (v1, v1, . . . , v1)|x),P(y = (v2, v1, . . . , v1)|x), . . . ,P(y = (vV , vV , . . . , vV )|x))> .
Proof. Denote the word posterior distribution as p∗(x). We will show f(x) = p∗(x). By the law of total expectation, we have
Lreconst = Ex,y[`(f(x), y)] = Ex[Ey|x[`(f(x), y)|x]].
We know the probability of y given x is p∗(x). Since ` is cross-entropy loss, we have
Lreconst ≥ Ex − ∑ k∈[V t] [p∗(x)]k log[p ∗(x)]k  , where the equality is obtained when f(x) = p∗(x). Thus, when f minimizes the reconstructionbased objective (1), we have f(x) = p∗(x).

heading: A.2 PROOF OF THEOREM 3
text: In this section, we give the proof of Theorem 3. Theorem A.1 is covered by Theorem 3 as the special case of t = 1. The proof follows the same idea as for Theorem A.1. We first show that f(x) is the t-words posterior (Lemma A.3). Then just as t = 1 case where we can recover topic posterior with A†f(x), we can also recover topic posterior of t-th moment with A†f(x) with some matrix A. The result follows by the observation that for a polynomial P (w) with degree at most t, Ew[P (w)|x] is a linear function of the topic posterior of t-th moment.
Theorem 3 (Main Result). Consider the general topic model setting as Definition 1, suppose topic matrix A satisfies rank(A) = K and function f minimizes the reconstruction-based objective (1). Then, any polynomial P (w) of the posterior w|x,A with degree at most t is linear in f(x), that is there exists a θ ∈ RV t such that for all documents x
Ew[P (w)|x] = θ>f(x).
Proof. Since f is the minimizer of reconstruction-based objective (1), by Lemma A.3 we know given an input document x,
f(x) = (P(y = (v1, v1, . . . , v1)|x),P(y = (v2, v1, . . . , v1)|x), . . . ,P(y = (vV , vV , . . . , vV )|x))> .
We will show that f(x) is linear in the topic posterior. In the following, we focus on [f(x)]y1,...,yt , which is word posterior probability of the t unknown words being y1, . . . , yt ∈ V given document x. Recall ∆K−1 be the K − 1 dimensional probability simplex. By the law of total probability, we have
[f(x)]y1,y2,...,yt = P(y1, y2, . . . , yt|x)
= ∫ w∈∆K−1 P(y1, y2, . . . , yt, w|x)dw
= ∫ w∈∆K−1 P(y1, y2, . . . , yt|w, x)P(w|x)dw
= ∫ w∈∆K−1 ∑ z1,...,zt∈[K] P(y1, y2, . . . , yt, z1, z2, ...zt|w)P(w|x)dw
= ∫ w∈∆K−1 ∑ z1,...,zt∈[K] P(z1, , . . . , zt|w)P(y1, . . . , yt|z1, . . . , zt, w)P(w|x)dw.
Note that zi is the topic of word yi. Since we consider the general topic model as Definition 1, we know
P(y1, y2, . . . , yt|z1, z2, ..., zt, w) = P(y1, y2, . . . , yt|z1, z2, ..., zt) = t∏ i=1 P(yi|zi) = t∏ i=1 Ayi,zi ,
where A is the topic matrix. Hence,
[f(x)]y1,y2,...,yt = ∑
z1,z2,...zt∈[K]
t∏ i=1 Ayi,zi ∫ w∈∆K−1 P(z1, z2, ..., zt|w)P(w|x)dw
= ∑
z1,z2,...zt∈[K]
t∏ i=1 Ayi,zi ∫ w∈∆K−1 t∏ i=1 wziP(w|x)dw
= ∑
z1,z2,...zt∈[K]
t∏ i=1 Axm+i,ziEw [ t∏ i=1 wzi ∣∣∣∣x ] .
Recall that the topic posterior tensor is Wpost = Ew[w⊗t|x] ∈ RK×...×K , where each entry [Wpost]z1,...,zk = P(zi is the topic of word yi for i ∈ [t]|x) = Ew[wz1 . . . wzt |x]. Therefore,
[f(x)]y1,y2,...,yt = ∑
z1,z2,...zt∈[K]
t∏ i=1 Ayi,zi [Wpost]z1,z2,...,zt ,
f(x) = (A⊗A⊗ · · · ⊗A︸ ︷︷ ︸ t times )vec(Wpost),
where ⊗ is the Kronecker product, A = A⊗A⊗ · · · ⊗A ∈ RV t×Kt and vec(Wpost) ∈ RK t
is the vectorization of Wpost. Note that A† = A† ⊗ · · · ⊗A†, so we have vec(Wpost) = A†f(x). Since P (w) is a polynomial of degree at most t, we know there exists β such that Ew[P (w)|x] = β>vec(Wpost). Thus, let θ = (A†)>β, we have Ew[P (w)|x] = β>vec(Wpost) = θ>f(x).

heading: A.3 PROOF OF THEOREM 4
text: Similar to Lemma A.2, we can show that the representation f(x) is close to t words posterior. Then, the key lemma to obtain Theorem 4 is the following result. It suggests that if function f is -close to the optimal function f∗ (word posterior distribution) under cross-entropy loss, then we can still recover the topic posterior with proper B up to O( √ ) error. For example, we will choose B = A† when t = 1. See the proof of Theorem 4 for the choice of B in the general case. Lemma A.4. For cross-entropy loss ` and any two probability vectors p, p∗, if `(p, p∗) ≤ minp: ∑ i pi=1,p≥0 `(p, p ∗) + , then for any matrix B we have ‖Bp−Bp∗‖1 ≤ κ(B) √ 2 .
Proof. Recall that ` is cross-entropy loss, we know
`(p, p∗)− min p: ∑ i pi=1,p≥0 `(p, p∗) = DKL(p || p∗) ≤ .
By Pinsker’s Inequality, we have that ‖p− p∗‖1 ≤ √ 2DKL(p∗ || p) ≤ √ 2 .
Then, for any matrix B we have
‖Bp−Bp∗‖1 ≤ κ(B)||p− p ∗||1 ≤ κ(B)
√ 2 ,
where we use the definition of `1 condition number (Definition 3).
We now are ready to give the proof of Theorem 4, which is based on Lemma A.4 and the proof of Theorem 3. Theorem 4 (Robust Version). Consider the general topic model setting as Definition 1, suppose topic matrix A satisfies rank(A) = K and function f satisfies Lreconst(f) ≤ L∗reconst + for some > 0. Then, any polynomial P (w) of the posterior w|x,A with degree at most t is approximately linear in f(x), that is there exists a θ ∈ RV t such that
Ex [( Ew[P (w)|x]− θ>f(x) )2] ≤ 2 ‖β‖2 κ2t(A†) , where Ew[P (w)|x] = β>vec(Wpost).
Proof. Recall from the proof of Theorem 3 that Ew[P (w)|x] = β>vec(Wpost) = β>A†f∗(x), where A† = A† ⊗ · · · ⊗ A† and κ(A†) = κt(A†). Same as in the proof of Theorem 3, let θ = (A†)>β. We have
Ex [( Ew[P (w)|x]− θ>f(x) )2] =Ex [( β>A†f∗(x)− β>A†f(x) )2] ≤Ex [ ‖β‖22 ∥∥A†f∗(x)−A†f(x)∥∥2 2
] ≤‖β‖22 Ex [∥∥A†f∗(x)−A†f(x)∥∥2 1 ] .
We are going to use Lemma A.4 to bound ∥∥A†f∗(x)−A†f(x)∥∥
1 . By Lemma A.3 we know f∗(x)
is the word posterior, so f∗(x) = Ey|x[y|x]. Since ` is cross-entropy loss, we know
Ey|x[`(f(x), y)− `(f∗(x), y)|x] = Ey|x  ∑ k∈[V t] yk log ( [f∗(x)]k [f(x)]k ) ∣∣∣∣x  = DKL(f∗(x) || f(x)),
which implies `(f(x), f∗(x))−minp `(p, f∗(x)) ≤ . Therefore, by Lemma A.4 with B = A†, we know ∥∥A†f(x)−A†f∗(x)∥∥
1 ≤ κ(A†)
√ 2 = κt(A†) √ 2 .
Given the desired bound, we have Ex [( Ew[P (w)|x]− θ>f(x) )2] ≤2 ‖β‖22 κ2t(A†) .

heading: B OMITTED PROOFS IN SECTION 4
text: In this section, we give the omitted proofs in Section 4. We give the proof of Theorem 5 in Section B.1 and proof of Lemma B.1 in Section B.2.

heading: B.1 PROOF OF THEOREM 5
text: Before presenting the proof of Theorem 5, we first give a characterization of the representation g(x, x′), which would be useful in the later analysis. Note that this the same as the one shown in Tosh et al. (2020). We provide its proof for completeness in Section B.2.
Lemma B.1. If f minimizes the contrastive objective (2), then we have
g(x, x′) , f(x, x′)
1− f(x, x′) = P(y = 1|x, x′) P(y = 0|x, x′) .
Now we are ready to proof Theorem 5.
Theorem 5. Consider the general topic model setting as Definition 1, suppose topic matrix A satisfies rank(A) = K and function f minimizes the contrastive objective (2). If we randomly sampled m = Kt different landmark documents {li}mi=1 and construct g(x, {li}mi=1) as (3), then any polynomial P (w) of the posterior w|x,A with degree at most t is linear in g(x, {li}mi=1), that is there exists a θ ∈ Rm such that for all documents x
Ew[P (w)|x] = θ>g(x, {li}mi=1).
Proof. Since f is the minimizer of contrastive objective (2), by Lemma B.1 we know
g(x, x′) = P(y = 1|x, x′) P(y = 0|x, x′) .
Similar to the proof of Theorem 3, we will show that g(x, x′) is linear in the topic posterior when x′ is fixed. By Bayes’ rule, we have
g(x, x′) = P(x, x′|y = 1)P(y = 1)/P(x, x′) P(x, x′|y = 0)P(y = 0)/P(x, x′) = P(x, x′|y = 1) P(x, x′|y = 0) ,
where we use P(y = 0) = P(y = 1) = 1/2.
Denote the t words in x′ as x′1, . . . , x ′ t and their corresponding topics as z1, . . . , zt. Then, by our way of generating x, x′, we know
P(x, x′|y = 0) = P(x)P(x′), P(x, x′|y = 1) = ∫ w P(x′1, . . . , x′t|w, x, y = 1)P(w, x|y = 1)dw
= ∫ w ∑ z1,...,zt∈[K] P(x′1, . . . , x′t|z1, . . . , zt, w)P(z1, . . . , zt|w)P(w, x)dw,
which implies
g(x, x′) = 1
P(x′) ∫ w ∑ z1,...,zt∈[K] P(x′1, . . . , x′t|z1, . . . , zt)P(z1, . . . , zt|w)P(w|x)dw
Note that
P(x′1, . . . , x′t|z1, . . . , zt) = t∏ i=1 P(x′i|zi) = t∏ i=1 Ax′1,zi ,
where A is the topic matrix. Hence,
g(x, x′) = 1 P(x′) ∑
z1,...,zt∈[K]
t∏ i=1 Ax′1,zi ∫ w P(z1, . . . , zt|w)P(w|x)dw
= 1 P(x′) ∑
z1,...,zt∈[K]
t∏ i=1 Ax′1,zi ∫ w t∏ i=1 wziP(w|x)dw
= 1 P(x′) ∑
z1,...,zt∈[K]
t∏ i=1 Ax′1,ziEw [ t∏ i=1 wzi ∣∣∣∣x ] .
Recall that the topic posterior tensor is Wpost = Ew[w⊗t|x] ∈ RK×...×K , where each entry [Wpost]z1,...,zk = Ew[wz1 . . . wzt |x]. Therefore,
g(x, x′) = 1 P(x′) ∑
z1,z2,...zt∈[K]
t∏ i=1 Ax′i,zi [Wpost]z1,z2,...,zt
= 1
P(x′) A[x′]>vec(Wpost),
whereA = A⊗A⊗· · ·⊗A ∈ RV t×Kt ,A[x′] ∈ RKt is the row ofA that correspond to x′1, . . . , x′t, and vec(Wpost) ∈ RK t is the vectorization of Wpost.
Recall we havem randomly sampled different documents {li}ni=1. DenoteD ∈ Rm×m as a diagonal matrix such that Di,i = P(li), and Ã ∈ Rm×K t
as a submatrix of A such that each row of Ã is A[li]. Since li are randomly sampled, so Di,i > 0. Thus, we know
g(x, {li}ni=1) = D−1Ãvec(Wpost),
which implies vec(Wpost) = Ã†Dg(x, {li}ni=1). Note that we need to show Ã† is well-defined. Since A has full column rank, we know A also has full column rank. Thus, the submatrix Ã also has full column rank since m = Kt. This implies Ã† is well-defined. Since P (w) is a polynomial of degree at most t, we know there exists β such that Ew[P (w)|x] = β>vec(Wpost). Thus, let θ = D(Ã†)>β, we have
Ew[P (w)|x] = β>vec(Wpost) = θ>g(x, {li}ni=1).

heading: B.2 PROOF OF LEMMA B.1
text: Lemma B.1. If f minimizes the contrastive objective (2), then we have
g(x, x′) , f(x, x′)
1− f(x, x′) = P(y = 1|x, x′) P(y = 0|x, x′) .
Proof. Since f is the minimizer of contrastive objective (2) and Lcontrast(f) = Ex,x′,y [ (f(x, x′)− y)2 ] = Ex,x′ [ Ey|x,x′ [ (f(x, x′)− y)2 ]] ,
it is easy to see that f is the Bayes optimal predictor P(y = 1|x, x′). Therefore, we know
g(x, x′) = P(y = 1|x, x′) P(y = 0|x, x′) .

heading: B.3 DISCUSSIONS ON THE SELF-REFERENCING
text: As discussed in Corollary 6, we show that one can use the same set of documents for both representation and downstream task. Recall that in Theorem 5 we need a large number of landmark documents to generate the representation and then use it for the downstream task (e.g., fit a polynomial of topic posterior). However, one may not have enough additional documents to be used as landmarks. Therefore, the benefit of the self-referencing is that we do not need additional landmark documents to generate the representation.
The key observation in Corollary 6 is that the learned representation g(x, x′) can be viewed as a kernel. Suppose the documents for the downstream task are x1, . . . , xm. Construct a kernel matrix G ∈ Rm×m such that Gij = g(xi, xj). Then, it is easy to see
θ>g(xi, {xj}mj=1) = θ>Gi,
where Gi is the i-th column of G. Thus, θ can be found by the following kernel regression:
min θ ‖ỹ −Gθ‖2 ,
where ỹ = (ỹ1, . . . , ỹm)> and ỹi = Ew[P (w)|xi].

heading: C SYNTHETIC EXPERIMENT ADDITIONAL DETAILS
text: In this section, we include more details about our synthetic experiments in Section 5. In Section C.1, we describe the document generation for PAM model. Additional results on topic posterior recovery loss and major topic recovery are included in Section C.2 and Section C.3. In Section C.4 and Section C.5, we report the results about different training epochs and hyperparameters tuning. `1 condition number is reported in Section C.6 and we give some visualizations for topic correlations between models in Section C.7.

heading: C.1 PACHINKO ALLOCATION MODEL TOPIC PRIOR GENERATION
text: The topic proportion of the Pachinko Allocation Model (PAM) described in Section 5 is generated from the following process: for each document, we first sample a “super-topic” proportion from a symmetric Dirichlet distribution Dir(1/Ks) and a super-topic-to-topic proportion from a symmetric Dirichlet distribution Dir(30). Then, we sample each word by first sample a super-topic according to the super-topic proportion and then sample the actual topic from the super-topic-to-topic proportion. In our experiment, we set Ks = 10. Figure C.1 offers a few examples of PAM topic proportion generated from the PAM model.

heading: C.2 MAJOR TOPIC RECOVERY ADDITIONAL RESULTS
text: In Section 5.3, we measure the major topics recovery of CTM and PAM documents as the top-2 topic overlap rate. Since topics are correlated in pairs in both of these topic models, we are also interested in the top-4 and top-6 topics overlap rate. As shown in Figure C.2, for both CTM and PAM documents, the average top-4 overlap rate is above 66% and the average top-6 overlap rate is above 62% on all α values we test on. These results show that our model can accurately capture more than just the major pair of correlated topics, but also some correlated pairs of less dominant topics.

heading: C.3 SELF-SUPERVISED LEARNING VERSUS ASSUMING SPECIFIC TOPIC PRIOR
text: We present in Section 5.3 a comparison between the performance of our model and Markov Chain Monte Carlo assuming a specific topic model for α = 1, on the task of TV distance and major topic
recovery. In this subsection, we supplement the experiment results for larger values of α along with α = 1. In particular, we consider self-supervised learning with 6 million training documents and with 120 thousand training documents after 200 training epochs, and we expect that the former can allow our model to approximate the optimal predictor f described in Theorem 3 and the latter may offer some insights on our approach’s performance in more practical settings, such as on a real world dataset with limited size. Additionally, when training with 120 thousand documents, we randomly shuffle each document and reassign the t target words in every epoch, to make sure our model learns enough information from the limited amount of training data. We report our results in the form of 95% confidence interval.
Table C.1 shows that our approach with 6 million training documents consistently outperforms misspecified topic priors by yielding a much lower TV distance from the true topic posterior. Note that the entries with correct topic model are omitted, because in our experiment the Markov Chain Monte Carlo recovered topic posterior assuming the correct topic model is exactly what we use as
our ground truth topic posterior. Meanwhile, our approach with 120 thousand training documents outperforms misspecified prior in almost every scenario, except when compared to LDA prior on pure topic documents. This is is likely because the LDA topic prior is relatively concentrated on one specific topic and thus similar to the pure topic prior.
We also present a comparison between our self-supervised learning approach and posterior inference assuming a specific topic prior on the task of recovering major topics in the topic proportion for α = 1, 3, 5, 7, 9, as a supplement to Section 5.3. As shown in Table C.2, in almost every scenario we test on, our model can perform competitively against the MCMC-inferred posterior assuming the correct topic prior and outperform misspecified topic prior.
To further investigate the performance of self-supervised learning, we compare our approach with Variational Inference assuming a specific topic prior on the task of recovering major topics in test document’s topic proportion. Table C.3 shows that, similar to our findings from Table C.2, our model can perform competitively against the Variational Inference posterior assuming the correct
topic prior, particularly for relative small α values, and can outperform misspecified topic prior in almost every test scenario.

heading: C.4 THE EFFECT OF TRAINING EPOCHS
text: We measure how the number of training epochs may influence topic posterior recovery loss by varying the number of training epochs. The results we present in Section 5 are based on models trained for 200 epochs for all types of documents. Here, we present the topic recovery loss, measured as the Total Variation distance between our recovered topic posterior and the true topic posterior, for epochs=25, 50, 100, 200 using the same model architecture and same sampling scheme as in Section 5.
Table C.4 shows that topic posterior recovery loss steadily decreases as the number of epochs gets larger. Interestingly, even with just training 50 epochs, our model can recover the topic posterior
within a Total Variation distance of less than 0.3 off the true topic posterior for documents generated from all four topic models.

heading: C.5 HYPERPARAMETER TUNING EXPERIMENTS
text: As described in Section 5, the two types of neural network architecture we use are fully-connected neural networks and attention-based neural networks (see Figure C.3). In this section, we present some of our TV results of different combinations of hyperparameters for both types of models. Note that our attention-based neural network slightly differs from typical transformers (Vaswani et al., 2017) in that we use batch normalization instead of layer normalization in every residual connection. We only use one attention head in each multi-head attention layer. During training, we use AMSGrad optimizer and an initial learning rate of 0.0001 or 0.0002. We reduce the learning rate by 50% whenever validation loss does not decrease in 10 epochs.
The main hyperparameters in our fully-connected neural networks include the hidden dimension, the number of layers, and whether we apply residual connections. Table C.5 shows the TV distance results when we vary these hyperparameters for the LDA scenario when α = 1. However, fullyconnected neural networks perform poorly in CTM and PAM case. For instance, when α = 1 and hidden dimension = 4096, we find that hyperparameters that work well in the LDA case no longer yield satisfactory results (Table C.6).

heading: 3 0.1509 0.1112 0.1095

heading: 6 0.7902 0.7893 0.7901
text: For attention-based neural networks applied to CTM and PAM scenarios, we mainly consider the number of layers and the hidden dimension per attention layer. Table C.7 presents the TV distance between recovered topic posterior and true topic posterior in the CTM setting, with varying number of layers and attention dimension for α = 1, 3, 5 and using a more coarse estimate to the true topic posterior than what we use as our final true CTM topic posterior. We find that the model performs
the best when attention dimension is 768 or 1024, and when the attention dimension increases to 2048 the model gives a higher TV (see Table C.8).

heading: 3 768 0.1471 0.1429 0.1398

heading: C.6 TOPIC POSTERIOR RECOVERY LOSS AND CONDITION NUMBER
text: We have proved in Theorem 4 that the upper bound of topic posterior recovery loss depends on the `1 Condition Number κ(A†). In this section, we show that κ(A†) is small for every topic matrix A in our experiments (Figure C.4). Note that we use the same topic matrix A for pure-topic model and the LDA model. For CTM and PAM, we use the same topic matrix up to reordering of the topics that corresponds with pairwise topic correlations. Therefore, pure-topic model and LDA model share the same κ(A†), and CTM and PAM share the same κ(A†).
C.7 VISUALIZING TOPIC CORRELATIONS BETWEEN MODELS
For documents generated by the PAM model, we plotted the estimated posteriors by the selfsupervised approach and posterior inference assuming different priors in Figure C.5. From this figure (especially in Documents 3 and 4), one can qualitatively see that the estimated posterior and the PAM MCMC posteriors are aware of the correlation between topics, while the LDA MCMC posterior fails to take the correlation into consideration and hence is more different from the ground truth topic proportions.

heading: D REAL-DATA EXPERIMENT ADDITIONAL DETAILS
text: In this section, we give more details about our real data experiments in Section 6. In Section D.1, we describe the data processing. We give a more detailed description of baselines and how we extracted our representations in Section D.2. In Section D.3, we report more results using different model architectures.

heading: D.1 DATA PROCESSING
text: Here we detailed our usage of the AG news dataset Zhang et al. (2015). Each category has 30,000 samples in the training set and 19,000 samples in the testing set. We first preprocessed the data by removing punctuation and words that occurred in fewer than 10 documents, obtaining in a vocabulary of around 16,700 words, in a similar fashion done by Tosh et al. (2020).
To split the data set into unsupervised dataset and supervised dataset. For each category of documents, we selected a random sample of 1000 documents as labeled supervised dataset while the remaining 116,000 documents fall into unsupervised dataset for representation learning.

heading: D.2 EXTRACTING REPRESENTATIONS
text: To give some details about two baseline representations we used, we described in details as follows:
• Bag of Words (BOW): For a single document, we constructs BOW embedding by creating a bag-of-words frequency vector of the dimension of vocabulary size, where each entry i represent the frequency of words with id i in that particular document.
• Word2vec: To generate word2vec representation fitted the Skip-gram word embedding model on unsupervised dataset (Mikolov et al. (2013)). The implementation is done through the Gensim library (Rehurek & Sojka (2011)), where we used an embedding dimension of 300 and window size of 5. Representation is taken as the average of all trained word embeddings in a single document.
For our own self-supervised method, to extract a representation, we attempted at 1) softmax+last layer: apply a Softmax directly to the last layer output, 2) word2vec+last layer: apply an word2vec embedding matrix to last layer output after Softmax to reduce dimension (the word2vec matrix is trained over the unsupervised dataset), and 3) softmax+second2last layer: apply a Softmax function on top of the second-to-last layer (equivalent to applying an identity matrix to replace the last layer). The dimension of representation with word2vec is 300 while the softmax + second2 last layer has a dimension of 4096. The original last layer representation has a dimension of the vocabulary size (around 16,700), and has shown inferior results than those with dimension reduction techniques.
We included a comparison of the two approaches with dimension reduction in Table D.1, from which we observe that a Softmax on the second-to-last layer output performed better across layers of 3,4 and 5. We reported the result using the softmax+second2last layer in Figure 2.

heading: D.3 MODEL ARCHITECTURE
text: We include here more details on our model architecture used in real data experiments. We used residual model, and tested two main hyperparameters: number of residual blocks and hidden dimension size. We fixed the rest of hyperparameters: we used a 5000 embedding dimension and trained
for 150 epochs, with 0.0002 learning rate, weight decay of 0.01 and resample rate of 2. We sampled 4 words in labels for variance reduction purpose.
Varying Depth and Width We investigated both the effect of model depth and width on the performance of RBL representation. We trained networks of width 2048, 3000, and 4096 nodes respectively, where the number of node refers to the number of node in the linear layer inside a residual block. We varied the depth of the network by choosing to place 3, 4 and 5 of such block.
As shown in Figure D.1, it appears that using wider models in the unsupervised phase leads to better performance when training a linear classifier on the learned representations. Number of residual block does not seem to have a clear relationship with test accuracy from limited experiments we ran. The best accuracy is achieved when we have 3 residual blocks with dimension of 4096, which is what we used to generate results in Figure 2.

Review summaries:
This paper introduces a new method for performing inference on topic models, based on self supervised learning. Their approach is generalizable to multiple types of topic models. Specifically, the authors show that the expectation of some polynomial function of the topic posterior will be a linear function of the output of function that optimizes the self-supervised learning objective. The authors explore 2 types of self-supervised learning objectives, one based on reconstruction and one based on contrastive learning that was explored in prior work. The authors provide proofs for their main result as well a result robust to approximate minimization of the self-supervised objective. In their experiments the authors show that their approach outperforms inference with a misspecified model. The authors also run semi-supervised learning experiments on their approach. 

This paper considers the problem of performing posterior inference for probabilistic topic models. A "general" topic model consists of two key parts:  - $A$ : the topic-word probability matrix - $\Delta$ : the prior distribution over doc-topic probability vector w  Given these parts and a new document x, the goal is to infer the posterior over w: $p(w | x)$.   It is well-known that this posterior can be estimated in several ways using standard approximate Bayesian inference methods: MCMC, variational, etc. The detailed steps of these methods are usually customized to the choice of prior.  This paper suggests a new possibility: that a learned transformation -- denoted f(x) -- trained using a so-called *reconstruction* loss (see Eq 1), can be then translated into the desired posterior via a linear function. Surprisingly, the paper suggests that learning f(x) does not depend on the prior \Delta, so the big idea is that this representation is robust to misspecification of the prior.  Concretely, the contributions of the paper seem to be:  1) Theorem 3, which claims that there exist linear weights \theta that, when applied to the vector produced by the ideal learned transformation f(x) that minimizes the reconstruction objective (Eq 1), can exactly equal the expected value of any polynomial summary function of the doc-topic probability vector w under the posterior.  2) Theorem 4, which claims that even if f is not an exact minimizer, if it is within an additive tolerance \epsilon of the ideal loss, then there exist again weights \theta such that the squared error between the true posterior summary and the linear function has bounded error for all documents.  3) Theorem 5, which looks at the contrastive objective in Eq 2, and shows that there exists a linear function that can recover the expected polynomical summaries of random variable w under the posterior.  Experiments in Sec. 5 look at toy data, with results in Fig 2 and Table 1. Essentially, these results show that the proposed SSL method produces lower error posterior estimates (in terms of  total-variation distance from ground truth) than using Bayesian inference with the wrong prior.   Experiments in Sec. 6 look at real data (the AG news dataset), comparing the representations learned by SSL at a downstream classification task to baselines of bag-of-words, word2vec, and the contrastive method of Tosh et al. Fig. 3 suggests that the proposed method offers a few percentage points absolute gain in accuracy over word2vec, and is similar to the previous NCE method (Tosh et al) in accuracy.

This paper shows theoretically and empirically that self-supervised objectives can be used to extract useful information about the posterior of the topic proportion vector given a document, regardless of the underlying models. It extends the findings of Tosh et al.2020 about contrastive objectives to reconstruction-based objectives. It is interesting to see that one simple reconstruction objective recover the posteriors generated from the different probabilistic topic models on the synthetic dataset.

The main strength of this paper lies in the thorough theoretical analysis and corresponding proofs. The authors not only prove that a new reconstruction-based objective can also extract posterior topic information for a general topic model, but also strengthen the guarantee for contrastive objective in Tosh et al. (2020) by removing some of their assumptions and the necessity of landmark documents.