Original: U NDERSTANDING
Revised: U NDERSTANDING

Original: T HE R OBUSTNESS OF S ELFSUPERVISED L EARNING T HROUGH T OPIC M ODELING
Revised: T HE R OBUSTNESS OF S ELFSUPERVISED L EARNING T HROUGH T OPIC M ODELING

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Self-supervised learning has signiﬁcantly improved the performance of many NLP tasks.
Revised: Self-supervised learning has signiﬁcantly improved the performance of many NLP tasks.

Original: However, how can self-supervised learning discover useful representations, and why is it better than traditional approaches such as probabilistic models are still largely unknown.
Revised: However, how can self-supervised learning discover useful representations, and why is it better than traditional approaches such as probabilistic models are still largely unknown.

Original: In this paper, we focus on the context of topic modeling and highlight a key advantage of self-supervised learning - when applied to data generated by topic models, self-supervised learning can be oblivious to the speciﬁc model, and hence is less susceptible to model misspeciﬁcation.
Revised: In this paper, we focus on the context of topic modeling and highlight a key advantage of self-supervised learning - when applied to data generated by topic models, self-supervised learning can be oblivious to the speciﬁc model, and hence is less susceptible to model misspeciﬁcation.

Original: 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.
Revised: 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.

Original: Empirically, we show that the same objectives can perform on par with posterior inference using the correct model, while outperforming posterior inference using misspeciﬁed models.
Revised: Empirically, we show that the same objectives can perform on par with posterior inference using the correct model, while outperforming posterior inference using misspeciﬁed models.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: 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.
Revised: 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.

Original: ; Jaiswal et al.
Revised: ; Jaiswal et al.

Original: and references therein).
Revised: and references therein).

Original: 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 predict the unknown words.
Revised: 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 predict the unknown words.

Original: This approach is known as self-supervised learning (SSL).
Revised: This approach is known as self-supervised learning (SSL).

Original: Why can self-supervised approaches learn useful representations?
Revised: Why can self-supervised approaches learn useful representations?

Original: To understand this we ﬁrst need to deﬁne what are “useful representations”.
Revised: To understand this we ﬁrst need to deﬁne what are “useful representations”.

Original: A recent line of work (Tosh et al., 2021a; Wei et al., 2021) studied self-supervised learning in the context of probabilistic models: assuming the data is generated by a probabilistic model (such as a topic model or Hidden Markov Model), one can deﬁne representation of observed data as the corresponding hidden variables in the model (such as topic proportions in topic models or hidden states in Hidden Markov Model).
Revised: A recent line of work (Tosh et al., 2021a; Wei et al., 2021) studied self-supervised learning in the context of probabilistic models: assuming the data is generated by a probabilistic model (such as a topic model or Hidden Markov Model), one can deﬁne representation of observed data as the corresponding hidden variables in the model (such as topic proportions in topic models or hidden states in Hidden Markov Model).

Original: These works show that self-supervised learning approach is as good as explicitly doing inference using such models.
Revised: These works show that self-supervised learning approach is as good as explicitly doing inference using such models.

Original: This approach naturally leads to the next question - why can self-supervised learning perform better than traditional inferencing based on probabilistic models?
Revised: This approach naturally leads to the next question - why can self-supervised learning perform better than traditional inferencing based on probabilistic models?

Original: In this paper we study this question in the context of topic modeling, and highlight one key advantage for self-supervised learning: robustness to model misspeciﬁcation.
Revised: In this paper we study this question in the context of topic modeling, and highlight one key advantage for self-supervised learning: robustness to model misspeciﬁcation.

Original: Many different models (such as Latent Dirichlet Allocation (LDA) (Blei et al., 2003), Correlated
Revised: Many different models (such as Latent Dirichlet Allocation (LDA) (Blei et al., 2003), Correlated

Original: Topic Model (CTM) (Blei & Lafferty, 2007), Pachinko Allocation Model (PAM) (Li & McCallum, 2006)) have been applied in practice.
Revised: Topic Model (CTM) (Blei & Lafferty, 2007), Pachinko Allocation Model (PAM) (Li & McCallum, 2006)) have been applied in practice.

Original: Traditional approaches would require different ways of doing inference depending on which model is used to generate the data.
Revised: Traditional approaches would require different ways of doing inference depending on which model is used to generate the data.

Original: On the other hand, we show that no matter which topic model is used to generate the data, if standard self-supervised learning objectives such as the reconstruction-based objective 1 or the contrastive objective 2 can be minimized, then they will generate representations that contain useful information about the topic proportions of a document.
Revised: On the other hand, we show that no matter which topic model is used to generate the data, if standard self-supervised learning objectives such as the reconstruction-based objective 1 or the contrastive objective 2 can be minimized, then they will generate representations that contain useful information about the topic proportions of a document.

Original: Self-supervised learning is oblivious to the choice of the probabilistic model, while the traditional approach of probabilistic modeling depends highly on the speciﬁc model.
Revised: Self-supervised learning is oblivious to the choice of the probabilistic model, while the traditional approach of probabilistic modeling depends highly on the speciﬁc model.

Original: Therefore, one would expect self-supervised learning to perform similarly to inferencing with the correct model, and outperforms inferencing with misspeciﬁed model.
Revised: Therefore, one would expect self-supervised learning to perform similarly to inferencing with the correct model, and outperforms inferencing with misspeciﬁed model.

Original: To verify our theory, we run synthetic experiments to show that self-supervised learning indeed outperforms inferencing with misspeciﬁed models.
Revised: To verify our theory, we run synthetic experiments to show that self-supervised learning indeed outperforms inferencing with misspeciﬁed models.

Original: Unlike large-scale models, our self-supervised learning is applied in the much simpler context of topic models, but we also demonstrate that even this simple application can improve over simple baselines on real data.
Revised: Unlike large-scale models, our self-supervised learning is applied in the much simpler context of topic models, but we also demonstrate that even this simple application can improve over simple baselines on real data.

Original: 1. R ELATED W ORKS
Revised: 1. R ELATED W ORKS

Original: Self-Supervised Learning Self-supervised learning recently has been shown to be able to learn useful representation, which is later used for downstream tasks.
Revised: Self-Supervised Learning Self-supervised learning recently has been shown to be able to learn useful representation, which is later used for downstream tasks.

Original: See for example Bachman et al.
Revised: See for example Bachman et al.

Original: ; Caron et al.
Revised: ; Caron et al.

Original: ; Chen et al.
Revised: ; Chen et al.

Original: (2020a;b;c);
Revised: (2020a;b;c);

Original: Grill et al.
Revised: Grill et al.

Original: ; Chen & He (2021); Tian et al.
Revised: ; Chen & He (2021); Tian et al.

Original: (2020a)
Revised: (2020a)

Original: ; He et al.
Revised: ; He et al.

Original: and references therein.
Revised: and references therein.

Original: In particular, Devlin et al.
Revised: In particular, Devlin et al.

Original: (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.
Revised: (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.

Original: 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., 2021a; 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).
Revised: 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., 2021a; 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).

Original: Also, several papers considered the problem from a multi-view perspective (Tsai et al., 2020; Tosh et al., 2021b), which covers both contrastive and reconstruction-based learning.
Revised: Also, several papers considered the problem from a multi-view perspective (Tsai et al., 2020; Tosh et al., 2021b), which covers both contrastive and reconstruction-based learning.

Original: Moreover, Wei et al.
Revised: Moreover, Wei et al.

Original: (2020) and Tian et al.
Revised: (2020) and Tian et al.

Original: (2021) studied the theoretical properties of self-training and the contrastive learning without the negative pairs respectively.
Revised: (2021) studied the theoretical properties of self-training and the contrastive learning without the negative pairs respectively.

Original: Saunshi et al.
Revised: Saunshi et al.

Original: (2020) investigated the beneﬁts of pre-trained language models for downstream tasks.
Revised: (2020) investigated the beneﬁts of pre-trained language models for downstream tasks.

Original: Most relevant to our paper, Tosh et al.
Revised: Most relevant to our paper, Tosh et al.

Original: (2021a) considered the contrastive learning in the topic models setting.
Revised: (2021a) considered the contrastive learning in the topic models setting.

Original: Our theoretical results extend their theory to reconstruction-based objective (while also removing some assumptions), and our empirical results show that the reconstruction-based objective can be effectively minimized.
Revised: Our theoretical results extend their theory to reconstruction-based objective (while also removing some assumptions for the contrastive objective), and our empirical results show that the reconstruction-based objective can be effectively minimized.

Original: Theoretical Analysis of Topic Models
Revised: Theoretical Analysis of Topic Models

Original: Many works have proposed provable algorithms to learn topic models, such as methods 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).
Revised: Many works have proposed provable algorithms to learn topic models, such as methods 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).

Original: Much less is known about provable inference for topic models.
Revised: Much less is known about provable inference for topic models.

Original: Sontag & Roy (2011) showed that MAP estimation can be NP-hard even for LDA model.
Revised: Sontag & Roy (2011) showed that MAP estimation can be NP-hard even for LDA model.

Original: Arora et al.
Revised: Arora et al.

Original: (2016b) considered approximate inference algorithms.
Revised: (2016b) considered approximate inference algorithms.

Original: 1. O UTLINE
Revised: 1. O UTLINE

Original: We ﬁrst introduce the basic concepts of topic models and our objectives in Section 2.
Revised: We ﬁrst introduce the basic concepts of topic models and our objectives in Section 2.

Original: Then in
Revised: Then in

Original: Section 3 we prove guarantees for the reconstruction-based objective.
Revised: Section 3 we prove guarantees for the reconstruction-based objective.

Original: Section 4 connects the contrastive objective to reconstruction-based objective which allows us to prove a stronger guarantee for the former.
Revised: Section 4 connects the contrastive objective to reconstruction-based objective which allows us to prove a stronger guarantee for the former.

Original: We then demonstrate the ability of self-supervised learning to adapt to different models by synthetic experiments in Section 5.
Revised: We then demonstrate the ability of self-supervised learning to adapt to different models by synthetic experiments in Section 5.

Original: Finally, we also evaluate the reconstruction-based objective on real-data to show that despite the simplicity of the topic modeling context, it extracts reasonable representations in Section 6.
Revised: Finally, we also evaluate the reconstruction-based objective on real-data to show that despite the simplicity of the topic modeling context, it extracts reasonable representations in Section 6.

Original:  P RELIMINARIES
Revised:  P RELIMINARIES

Original: In this section we ﬁrst introduce some general notations.
Revised: In this section we ﬁrst introduce some general notations.

Original: Then we brieﬂy describe the topic models we consider in Section 2.1.
Revised: Then we brieﬂy describe the topic models we consider in Section 2.1.

Original: Finally we deﬁne the self-supervised learning objectives in Section 2. and give our main results.
Revised: Finally we deﬁne the self-supervised learning objectives in Section 2. and give our main results.

Original: Notation We use [ n ] to denote set { 1 , 2 , . . . , n } .
Revised: Notation We use [ n ] to denote set { 1 , 2 , . . . , n } .

Original: For vector x ∈ R d , denote (cid:107) x (cid:107) as its (cid:96) 2 norm and (cid:107) x (cid:107) 1 as its (cid:96) 1 norm.
Revised: For vector x ∈ R d , denote (cid:107) x (cid:107) as its (cid:96) 2 norm and (cid:107) x (cid:107) 1 as its (cid:96) 1 norm.

Original: For matrix A ∈ R m × n , we use A i ∈ R m to denote its i -th column.
Revised: For matrix A ∈ R m × n , we use A i ∈ R m to denote its i -th column.

Original: When matrix A has full column rank, denote its left pseudo-inverse as A † =
Revised: When matrix A has full column rank, denote its left pseudo-inverse as A † =

Original: ( A (cid:62) A ) − 1 A (cid:62) .
Revised: ( A (cid:62) A ) − 1 A (cid:62) .

Original: For matrix or general tensor T ∈ R
Revised: For matrix or general tensor T ∈ R

Original: d 1 ×···× d l , we use vector vec ( T ) ∈ R d 1 ··· d l to represent its vectorization.
Revised: d 1 ×···× d l , we use vector vec ( T ) ∈ R d 1 ··· d l to represent its vectorization.

Original: Let
Revised: Let

Original: S k
Revised: S k

Original: = { x ∈ R k | (cid:80) ki
Revised: = { x ∈ R k | (cid:80) ki

Original: =1 x i = 1 , x i ≥ 0 } to denote k − 1 dimensional probability simplex.
Revised: =1 x i = 1 , x i ≥ 0 } to denote k − 1 dimensional probability simplex.

Original: For two probability vectors p, q , deﬁne their total variation (TV) distance as TV ( p, q )
Revised: For two probability vectors p, q , deﬁne their total variation (TV) distance as TV ( p, q )

Original: = (cid:107) p − q (cid:107) 1 / 2 .
Revised: = (cid:107) p − q (cid:107) 1 / 2 .

Original: 2. T OPIC M ODELS
Revised: 2. T OPIC M ODELS

Original: Many topic models treat documents as a bag-of-words and use a two-step procedure of generating a document: ﬁrst 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.
Revised: Many topic models treat documents as a bag-of-words and use a two-step procedure of generating a document: ﬁrst 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.

Original: To generate a word, one ﬁrst selects a topic based on the topic proportions w , and then samples a word from that particular topic.
Revised: To generate a word, one ﬁrst selects a topic based on the topic proportions w , and then samples a word from that particular topic.

Original: Since the topic distributions are shared across the entire corpus, and the topic proportion vector w is speciﬁc to documents, it makes sense to deﬁne w as the representation for the document.
Revised: Since the topic distributions are shared across the entire corpus, and the topic proportion vector w is speciﬁc to documents, it makes sense to deﬁne w as the representation for the document.

Original: More precisely, let V be a ﬁnite vocabulary with size V and K be a set of K topics, where each topic is a distribution over V .
Revised: More precisely, let V be a ﬁnite vocabulary with size V and K be a set of K topics, where each topic is a distribution over V .

Original: We denote the topic-word matrix as A with A ij =
Revised: We denote the topic-word matrix as A with A ij =

Original: P ( word i | topic j ) , so that each column in A represents the word distribution of a topic.
Revised: P ( word i | topic j ) , so that each column in A represents the word distribution of a topic.

Original: Each document corresponds to a convex combination of different topics with proportions w .
Revised: Each document corresponds to a convex combination of different topics with proportions w .

Original: For each word in the document, ﬁrst 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 ).
Revised: For each word in the document, ﬁrst 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 ).

Original: Different topic models differ in how they generate w , which we formulate in the following deﬁnition:
Revised: Different topic models differ in how they generate w , which we formulate in the following deﬁnition:

Original: Deﬁnition 2.1 (General Topic Model) .
Revised: Deﬁnition 2.1 (General Topic Model) .

Original: A general topic model speciﬁes a distribution ∆( K ) for each number of topics K .
Revised: A general topic model speciﬁes a distribution ∆( K ) for each number of topics K .

Original: Given a topic-word matrix A and ∆( K ) , to generate a document, one ﬁrst sample w ∼ ∆( K ) and then sample each word in the document from the distribution
Revised: Given a topic-word matrix A and ∆( K ) , to generate a document, one ﬁrst sample w ∼ ∆( K ) and then sample each word in the document from the distribution

Original: Aw .
Revised: Aw .

Original: Here ∆( K ) is a prior distribution of w which is crucial when trying to infer the true topic proportions w given the words in a document.
Revised: Here ∆( K ) is a prior distribution of w which is crucial when trying to infer the true topic proportions w given the words in a document.

Original: Of course, different topic models may also specify different priors for the topic-word matrix A .
Revised: Of course, different topic models may also specify different priors for the topic-word matrix A .

Original: However, given a large number of documents generated from the topic model, in many settings one can hope to learn the topic-word matrix A and the prior distribution ∆( K ) (see e.g., Arora et al.
Revised: However, given a large number of documents generated from the topic model, in many settings one can hope to learn the topic-word matrix A and the prior distribution ∆( K ) (see e.g., Arora et al.

Original: (2012; 2016a)), therefore we consider the following inference problem:
Revised: (2012; 2016a)), therefore we consider the following inference problem:

Original: Deﬁnition 2.2 (Topic Inference) .
Revised: Deﬁnition 2.2 (Topic Inference) .

Original: Given a topic-word 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 ) .
Revised: Given a topic-word 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 ) .

Original: Note that our general topic model can capture many standard topic models, including pure topic model, LDA, CTM and PAM.
Revised: Note that our general topic model can capture many standard topic models, including pure topic model, LDA, CTM and PAM.

Original: 2. S EMI - SUPERVISED L EARNING S ETUP
Revised: 2. S EMI - SUPERVISED L EARNING S ETUP

Original: Our goal is to understand the representation learning aspect of self-supervised learning, where we are given large number of unlabeled documents and a small number of labeled documents.
Revised: Our goal is to understand the representation learning aspect of self-supervised learning, where we are given large number of unlabeled documents and a small number of labeled documents.

Original: In the unsupervised learning stage, the algorithm only has access to the unlabeled documents, and the goal is to learn a representation (a mapping from documents to a feature vector).
Revised: In the unsupervised learning stage, the algorithm only has access to the unlabeled documents, and the goal is to learn a representation (a mapping from documents to a feature vector).

Original: In the supervised learning stage, the algorithm would use the labeled documents to train a simple classiﬁer that maps the feature vector to the actual label.
Revised: In the supervised learning stage, the algorithm would use the labeled documents to train a simple classiﬁer that maps the feature vector to the actual label.

Original: We often refer to the unsupervised learning stage as “feature learning” and the supervised learning stage as “downstream application”.
Revised: We often refer to the unsupervised learning stage as “feature learning” and the supervised learning stage as “downstream application”.

Original: In the theoretical model, we assume that the label y is a simple function of the topic proportions w .
Revised: In the theoretical model, we assume that the label y is a simple function of the topic proportions w .

Original: 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 .
Revised: 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 .

Original: Of course, if one knows the actual parameters of the model, the best predictor for y would be to ﬁrst estimate the posterior distribution w and then apply the function that maps w to label y .
Revised: Of course, if one knows the actual parameters of the model, the best predictor for y would be to ﬁrst estimate the posterior distribution w and then apply the function that maps w to label y .

Original: We will show that for functions that are approximable by low degree polynomials self-supervised learning can always provide a good representation.
Revised: We will show that for functions that are approximable by low degree polynomials self-supervised learning can always provide a good representation.

Original: 2. S ELF - SUPERVISED LEARNING
Revised: 2. S ELF - SUPERVISED LEARNING

Original: In general, self-supervised learning tries to turn unlabeled data into a supervised learning problem by hiding some known information.
Revised: In general, self-supervised learning tries to turn unlabeled data into a supervised learning problem by hiding some known information.

Original: There are many ways to achieve this goal (see e.g., Liu et al.
Revised: There are many ways to achieve this goal (see e.g., Liu et al.

Original: Jaiswal et al.
Revised: Jaiswal et al.

Original: In this paper, we focus on two different approaches for self-supervised learning:
Revised: In this paper, we focus on two different approaches for self-supervised learning:

Original: reconstruction-based objective and
Revised: reconstruction-based objective and

Original: contrastive
Revised: contrastive

Original: objective
Revised: objective

Original: We ﬁrst formally deﬁne the objectives and then discuss our corresponding result.
Revised: We ﬁrst formally deﬁne the objectives and then discuss our corresponding result.

Original: Reconstruction-Based Objective One common approach of self-supervised learning is to ﬁrst
Revised: Reconstruction-Based Objective One common approach of self-supervised learning is to ﬁrst

Original: mask part of the data and then try to ﬁnd a function f to reconstruct the missing part given the unmasked part of input.
Revised: mask part of the data and then try to ﬁnd a function f to reconstruct the missing part given the unmasked part of input.

Original: This is commonly used for language modeling (Devlin et al., 2018;
Revised: This is commonly used for language modeling (Devlin et al., 2018;

Original: Radford et al.,
Revised: Radford et al.,

Original: In the context of topic modelling, suppose all the documents are generated from an underlying topic model with A and ∆( K ) .
Revised: In the context of topic modelling, suppose all the documents are generated from an underlying topic model with A and ∆( K ) .

Original: Then for any given document x unsup , since each word is i.i.d.
Revised: Then for any given document x unsup , since each word is i.i.d.

Original: sampled (so that the order of words is not important), we pick t random words from the document and mark them as unknown, then we ask the learner to predict these t words given the remaining words in the document.
Revised: sampled (so that the order of words is not important), we pick t random words from the document and mark them as unknown, then we ask the learner to predict these t words given the remaining words in the document.

Original: Speciﬁcally, we split x unsup into x and y where y is the t words that we select and let x is the document with these t words removed.
Revised: Speciﬁcally, we split x unsup into x and y where y is the t words that we select and let x is the document with these t words removed.

Original: We aim to select a predictor f that minimizes the following reconstruction objective:
Revised: We aim to select a predictor f that minimizes the following reconstruction objective:

Original: [Equation]
Revised: [Equation]

Original: where (cid:96) (ˆ y, y )
Revised: where (cid:96) (ˆ y, y )

Original: =
Revised: =

Original: (cid:80) k − y k log ˆ y k is the cross entropy loss.
Revised: (cid:80) k − y k log ˆ y k is the cross entropy loss.

Original: Here we slightly abuse the notation to use y as an one-hot vector for V t classes and f also outputs a probability vector on V t .
Revised: Here we slightly abuse the notation to use y as an one-hot vector for V t classes and f also outputs a probability vector on V t .

Original: Depending on the context, we will use y to denote either the actual next t words or its corresponding one-hot label.
Revised: Depending on the context, we will use y to denote either the actual next t words or its corresponding one-hot label.

Original: Now we are ready to present our result for the reconstruction-based objective:
Revised: Now we are ready to present our result for the reconstruction-based objective:

Original: Theorem 2.3.
Revised: Theorem 2.3.

Original: (Informal) Consider the general topic model setting as Deﬁnition 2.1, suppose function f minimizes the reconstruction-based objective (1) .
Revised: (Informal) Consider the general topic model setting as Deﬁnition 2.1, suppose function f minimizes the reconstruction-based objective (1) .

Original: 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 ) .
Revised: 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 ) .

Original: The theorem shows that if we want to get basic information about posterior distribution (such as mean, variance), then it sufﬁces to predict a small constant number of words (1 for mean and 2 for variance).
Revised: The theorem shows that if we want to get basic information about posterior distribution (such as mean, variance), then it sufﬁces to predict a small constant number of words (1 for mean and 2 for variance).

Original: See Theorem 3.1 for the formal statement.
Revised: See Theorem 3.1 for the formal statement.

Original: Contrastive Objective
Revised: Contrastive Objective

Original: Another common approach in self-supervised learning is the contrastive learning (He et al., 2020; Chen et al., 2020a).
Revised: Another common approach in self-supervised learning is the contrastive learning (He et al., 2020; Chen et al., 2020a).

Original: In contrastive learning, the training data is usually a pair of data ( x, x (cid:48) ) with a label y ∈ { 0 , 1 } , where label 1 means ( x, x (cid:48) ) is a positive sample ( x and x (cid:48) are similar) and label 0 means ( x, x (cid:48) ) is a negative sample ( x and x (cid:48) are dissimilar).
Revised: In contrastive learning, the training data is usually a pair of data ( x, x (cid:48) ) with a label y ∈ { 0 , 1 } , where label 1 means ( x, x (cid:48) ) is a positive sample ( x and x (cid:48) are similar) and label 0 means ( x, x (cid:48) ) is a negative sample ( x and x (cid:48) are dissimilar).

Original: The task is to ﬁnd a function f such that it can distinguish the positive sample and negative sample, i.e., f ( x, x (cid:48) )
Revised: The task is to ﬁnd a function f such that it can distinguish the positive sample and negative sample, i.e., f ( x, x (cid:48) )

Original: =
Revised: =

Original: y .
Revised: y .

Original: Formally, we want to select a predictor f such that it minimizes the following contrastive objective:
Revised: Formally, we want to select a predictor f such that it minimizes the following contrastive objective:

Original: [Equation]
Revised: [Equation]

Original: In the context of topic models, following previous work (Tosh et al., 2021a), we generate the data ( x, x (cid:48) , y ) as follows.
Revised: In the context of topic models, following previous work (Tosh et al., 2021a), we generate the data ( x, x (cid:48) , y ) as follows.

Original: Suppose all documents are generated from an underlying topic model with A and ∆( K ) .
Revised: Suppose all documents are generated from an underlying topic model with A and ∆( K ) .

Original: We ﬁrst generate a document x from the word distribution
Revised: We ﬁrst generate a document x from the word distribution

Original: Aw where w is sampled from ∆( K ) .
Revised: Aw where w is sampled from ∆( K ) .

Original: Then, (i) with half probability we generate t words from the same distribution
Revised: Then, (i) with half probability we generate t words from the same distribution

Original: Aw to form the document x (cid:48) and set y = 1 ; (ii) with half probability we generate t words from a different word distribution
Revised: Aw to form the document x (cid:48) and set y = 1 ; (ii) with half probability we generate t words from a different word distribution

Original: Aw (cid:48) with w (cid:48) ∼ ∆( K ) (
Revised: Aw (cid:48) with w (cid:48) ∼ ∆( K ) (

Original: so that w (cid:54) = w (cid:48) ) and set y = 0 .
Revised: so that w (cid:54) = w (cid:48) ) and set y = 0 .

Original: We consider the square loss (cid:96) (ˆ y, y ) =
Revised: We consider the square loss (cid:96) (ˆ y, y ) =

Original: (ˆ y − y ) 2 for simplicity.
Revised: (ˆ y − y ) 2 for simplicity.

Original: We now give our informal result on contrastive objective.
Revised: We now give our informal result on contrastive objective.

Original: See Theorem 4.1 for the formal statement.
Revised: See Theorem 4.1 for the formal statement.

Original: Note that this theorem generalizes Theorem 3 in Tosh et al.
Revised: Note that this theorem generalizes Theorem 3 in Tosh et al.

Original: (2021a).
Revised: (2021a).

Original: Theorem 2.4.
Revised: Theorem 2.4.

Original: (Informal) Consider the general topic model setting as Deﬁnition 2.1, suppose function f minimizes the contrastive objective (2) .
Revised: (Informal) Consider the general topic model setting as Deﬁnition 2.1, suppose function f minimizes the contrastive objective (2) .

Original: 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 ) .
Revised: 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 ) .

Original: Advantage of SSL approach Note that neither objectives (1) or (2) depend on the prior distribution ∆( K ) , so it is possible to optimize these without specifying a speciﬁc model; on the other hand, even the deﬁnition of topic inference (Deﬁnition 2.2) relies heavily on ∆( K ) .
Revised: Advantage of SSL approach Note that neither objectives (1) or (2) depend on the prior distribution ∆( K ) , so it is possible to optimize these without specifying a speciﬁc model; on the other hand, even the deﬁnition of topic inference (Deﬁnition 2.2) relies heavily on ∆( K ) .

Original:  G UARANTEES FOR THE R ECONSTRUCTION -B ASED O BJECTIVE
Revised:  G UARANTEES FOR THE R ECONSTRUCTION -B ASED O BJECTIVE

Original: In this section, we consider the reconstruction-based objective (1) and provide theoretical guarantees for its performance.
Revised: In this section, we consider the reconstruction-based objective (1) and provide theoretical guarantees for its performance.

Original: We ﬁrst 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.
Revised: We ﬁrst 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.

Original: Theorem 3.1 (Main Result) .
Revised: Theorem 3.1 (Main Result) .

Original: Consider the general topic model setting as Deﬁnition 2.1, suppose topic-word matrix A satisﬁes rank ( A ) =
Revised: Consider the general topic model setting as Deﬁnition 2.1, suppose topic-word matrix A satisﬁes rank ( A ) =

Original: K and function f minimizes the reconstruction-based objective (1) .
Revised: K and function f minimizes the reconstruction-based objective (1) .

Original: Then, for any document x and its posterior w | ( x, A ) , the posterior mean of any polynomial P ( w ) with degree at most t is linear in f ( x ) .
Revised: Then, for any document x and its posterior w | ( x, A ) , the posterior mean of any polynomial P ( w ) with degree at most t is linear in f ( x ) .

Original: That is, there exists a θ ∈ R V t such that for all documents x
Revised: That is, there exists a θ ∈ R V t such that for all documents x

Original: [Equation]
Revised: [Equation]

Original: To understand this theorem, ﬁrst consider a warm-up example where t = 1 (see details in Section A. in appendix).
Revised: To understand this theorem, ﬁrst consider a warm-up example where t = 1 (see details in Section A. in appendix).

Original: Intuitively, in this case the best way to predict the missing word for a given document x is to estimate its topic proportion w , and then predict the word using Aw where A is the topic-word matrix.
Revised: Intuitively, in this case the best way to predict the missing word for a given document x is to estimate its topic proportion w , and then predict the word using Aw where A is the topic-word matrix.

Original: 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 ) =
Revised: 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 ) =

Original: A E [ w | x, A ] .
Revised: A E [ w | x, A ] .

Original: When A is full rank multiplying by the pseudo-inverse of A recovers the expectation of the posterior w | ( x, A ) .
Revised: When A is full rank multiplying by the pseudo-inverse of A recovers the expectation of the posterior w | ( x, A ) .

Original: The proof for the general case of predicting t -words requires more careful characterization of the optimal prediction and its relationship to E w [ P ( w ) | x, A ] , which we defer to Section A.2.
Revised: The proof for the general case of predicting t -words requires more careful characterization of the optimal prediction and its relationship to E w [ P ( w ) | x, A ] , which we defer to Section A.2.

Revised: The above discussion also suggests that though the representation dimension may seem to be V t , its “effective" dimension is in fact K t which is much smaller than V t .

Revised: See more details in Appendix A.4.

Original: Robustness for an approximate minimizer In Theorem 3.1 we focus on the case when function f is exactly the minimizer of the reconstruction-based objective (1), i.e., L reconst ( f ) =
Revised: Robustness for an approximate minimizer In Theorem 3.1 we focus on the case when function f is exactly the minimizer of the reconstruction-based objective (1), i.e., L reconst ( f ) =

Original: L ∗ reconst (cid:44) min f L reconst ( f ) .
Revised: L ∗ reconst (cid:44) min f L reconst ( f ) .

Original: However, in practice one cannot hope to ﬁnd such a function exactly.
Revised: However, in practice one cannot hope to ﬁnd such a function exactly.

Original: In the following, we provide a robust version of Theorem 3.1 such that it allows us to ﬁnd an approximate solution instead of the exact optimal solution.
Revised: In the following, we provide a robust version of Theorem 3.1 such that it allows us to ﬁnd an approximate solution instead of the exact optimal solution.

Original: To present our result, we need to ﬁrst introduce the following notion of condition number, which was used in many previous works, such as collaborative ﬁltering systems (Kleinberg & Sandler, 2008) and topic models (Arora et al., 2016b).
Revised: To present our result, we need to ﬁrst introduce the following notion of condition number, which was used in many previous works, such as collaborative ﬁltering systems (Kleinberg & Sandler, 2008) and topic models (Arora et al., 2016b).

Original: Intuitively, κ ( B ) measures how large a vector would change after multiplying with B in the (cid:96) 1 norm sense.
Revised: Intuitively, κ ( B ) measures how large a vector would change after multiplying with B in the (cid:96) 1 norm sense.

Original: Deﬁnition 3.2 ( (cid:96) 1 Condition Number) .
Revised: Deﬁnition 3.2 ( (cid:96) 1 Condition Number) .

Original: For matrix B ∈ R m × n , deﬁne its (cid:96) 1 condition number κ ( B ) as κ ( B ) (cid:44) min x ∈ R n : x (cid:54) =0 (cid:107) Bx (cid:107) 1 / (cid:107) x
Revised: For matrix B ∈ R m × n , deﬁne its (cid:96) 1 condition number κ ( B ) as κ ( B ) (cid:44) min x ∈ R n : x (cid:54) =0 (cid:107) Bx (cid:107) 1 / (cid:107) x

Original: (cid:107) 1 = max i ∈
Revised: (cid:107) 1 = max i ∈

Original: [ n ] (cid:107) B
Revised: [ n ] (cid:107) B

Original: i (cid:107) 1 , where B i is the i -th column of B .
Revised: i (cid:107) 1 , where B i is the i -th column of B .

Original: Let W post ∈
Revised: Let W post ∈

Original: R K ×···× K be the topic posterior tensor for t unknown words y
Revised: R K ×···× K be the topic posterior tensor for t unknown words y

Original: = ( y 1 , . . .
Revised: = ( y 1 , . . .

Original: , y t ) given remaining document x .
Revised: , y t ) given remaining document x .

Original: That is, for each entry [ W post ] z 1 ,...,z k =
Revised: That is, for each entry [ W post ] z 1 ,...,z k =

Original: P ( for all i , z i is the topic of word y i
Revised: P ( for all i , z i is the topic of word y i

Original: | x, A ) =
Revised: | x, A ) =

Original: E w [ w z 1 · · · w z t | x, A ] and W post =
Revised: E w [ w z 1 · · · w z t | x, A ] and W post =

Original: E w [ w ⊗ t | x, A ] .
Revised: E w [ w ⊗ t | x, A ] .

Original: Thus, any polynomial P ( w ) of degree at most t can be represented as E w [ P ( w ) | x, A ] =
Revised: Thus, any polynomial P ( w ) of degree at most t can be represented as E w [ P ( w ) | x, A ] =

Original: β (cid:62) vec ( W post ) for some β , where vec ( W post ) ∈ R K t is the vectorization of W post .
Revised: β (cid:62) vec ( W post ) for some β , where vec ( W post ) ∈ R K t is the vectorization of W post .

Original: We now are ready to present the robust version of Theorem 3.1.
Revised: We now are ready to present the robust version of Theorem 3.1.

Original: It shows if we only ﬁnd a function f whose loss is at most (cid:15) 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 ( (cid:15) ) error.
Revised: It shows if we only ﬁnd a function f whose loss is at most (cid:15) 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 ( (cid:15) ) error.

Original: Theorem 3.3 (Robust Version) .
Revised: Theorem 3.3 (Robust Version) .

Original: Consider the general topic model setting as Deﬁnition 2.1, suppose topic-word matrix A satisﬁes rank ( A ) =
Revised: Consider the general topic model setting as Deﬁnition 2.1, suppose topic-word matrix A satisﬁes rank ( A ) =

Original: K and function f satisﬁes L reconst ( f ) ≤ L ∗ reconst + (cid:15) for some (cid:15) > 0 .
Revised: K and function f satisﬁes L reconst ( f ) ≤ L ∗ reconst + (cid:15) for some (cid:15) > 0 .

Original: Then, for any document x and its posterior w | ( x, A ) , the posterior mean of any polynomial
Revised: Then, for any document x and its posterior w | ( x, A ) , the posterior mean of any polynomial

Original: P ( w ) with degree at most t is approximately linear in f ( x ) .
Revised: P ( w ) with degree at most t is approximately linear in f ( x ) .

Original: That is, there exists a θ ∈ R V t such that
Revised: That is, there exists a θ ∈ R V t such that

Original: [Equation]
Revised: [Equation]

Original: A [Equation] ) (cid:15), where E w [ P ( w ) | x, A ] = β (cid:62) vec ( W post ) .
Revised: where E w [ P ( w ) | x, A ] = β (cid:62) vec ( W post ) .

Original: Note that the dependency on (cid:107) β (cid:107) 2 κ 2 t ( A † ) is expected, since this is the norm (cid:107) θ (cid:107) 2 that we would have if (cid:15) = 0 , i.e., the θ we would have in Theorem 3.1.
Revised: Note that the dependency on (cid:107) β (cid:107) 2 κ 2 t ( A † ) is expected, since this is the norm (cid:107) θ (cid:107) 2 that we would have if (cid:15) = 0 , i.e., the θ we would have in Theorem 3.1.

Original: Thus, this quantity should be understood as the complexity of the target function for the downstream task.
Revised: Thus, this quantity should be understood as the complexity of the target function for the downstream task.

Original: Empirically we show that κ ( A † ) is small in Section C.6.
Revised: Empirically we show that κ ( A † ) is small in Section C.6.

Original: The proof is deferred to Section A.3.
Revised: The proof is deferred to Section A.3.

Original:  G UARANTEES FOR THE C ONTRASTIVE O BJECTIVE
Revised:  G UARANTEES FOR THE C ONTRASTIVE O BJECTIVE

Original: 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.
Revised: 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.

Original: We use the same approach as Tosh et al.
Revised: We use the same approach as Tosh et al.

Original: (2021a) to construct a representation g ( x ) .
Revised: (2021a) to construct a representation g ( x ) .

Original: Given a set of landmark documents { l i } mi =1 with length | l i | = t as references, the representation is deﬁned as
Revised: Given a set of landmark documents { l i } mi =1 with length | l i | = t as references, the representation is deﬁned as

Original: [Equation]
Revised: [Equation]

Original: The following theorem gives the theoretical guarantee for this representation.
Revised: The following theorem gives the theoretical guarantee for this representation.

Original: Similar to Theorem 3.
Revised: Similar to Theorem 3.

Original: 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.
Revised: 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.

Original: Our proof here relies on the observation that having g ( x, { l i } mi =1 ) for all short landmark documents of length t gives similar information as f for the reconstruction objective.
Revised: Our proof here relies on the observation that having g ( x, { l i } mi =1 ) for all short landmark documents of length t gives similar information as f for the reconstruction objective.

Original: This observation allows us to remove the anchor words assumption needed in (Tosh et al., 2021a).
Revised: This observation allows us to remove the anchor words assumption needed in (Tosh et al., 2021a).

Original: Theorem 4.1.
Revised: Theorem 4.1.

Original: Consider the general topic model setting as Deﬁnition 2.1, suppose topic-word matrix
Revised: Consider the general topic model setting as Deﬁnition 2.1, suppose topic-word matrix

Original: A satisﬁes rank ( A ) =
Revised: A satisﬁes rank ( A ) =

Original: K and function f minimizes the contrastive objective (2) .
Revised: K and function f minimizes the contrastive objective (2) .

Original: Assume we randomly sampled m =
Revised: Assume we randomly sampled m =

Original: K t different landmark documents { l i } mi =1 and construct g ( x, { l i } mi =1 ) as (3) .
Revised: K t different landmark documents { l i } mi =1 and construct g ( x, { l i } mi =1 ) as (3) .

Original: Then, for any document x and its posterior w | ( x, A ) , the posterior mean of any polynomial P ( w ) with degree at most t is linear in g ( x, { l i } mi =1 ) .
Revised: Then, for any document x and its posterior w | ( x, A ) , the posterior mean of any polynomial P ( w ) with degree at most t is linear in g ( x, { l i } mi =1 ) .

Original: That is, there exists θ ∈ R m such that for all documents x
Revised: That is, there exists θ ∈ R m such that for all documents x

Original: [Equation]
Revised: [Equation]

Original: 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.
Revised: 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.

Original:  S YNTHETIC E XPERIMENTS
Revised:  S YNTHETIC E XPERIMENTS

Original: In this section, we optimize the reconstruction-based objective (1) on data generated by several topic models to show that self-supervised learning performs on par with inference using correct model, and both of them outperform inference with misspeciﬁed models.
Revised: In this section, we optimize the reconstruction-based objective (1) on data generated by several topic models to show that self-supervised learning performs on par with inference using correct model, and both of them outperform inference with misspeciﬁed models.

Original: 5. T OPIC M ODELS
Revised: 5. T OPIC M ODELS

Original: We consider four types of topic models in our experiments.
Revised: We consider four types of topic models in our experiments.

Original: Our ﬁrst topic model is the puretopic model, where each document’s topic comes from a discrete uniform distribution over the
Revised: Our ﬁrst topic model is the puretopic model, where each document’s topic comes from a discrete uniform distribution over the

Original: K -dimensional linear basis, namely { e 1 , e 2 , ..., e K } .
Revised: K -dimensional linear basis, namely { e 1 , e 2 , ..., e K } .

Original: Our second topic model is the Latent Dirichlet
Revised: Our second topic model is the Latent Dirichlet

Original: Allocation (LDA) model, where ∆( K ) is a symmetric Dirichlet distribution Dir (1 /K ) .
Revised: Allocation (LDA) model, where ∆( K ) is a symmetric Dirichlet distribution Dir (1 /K ) .

Original: We are also interested in topic models that involve more subtle topic correlations and similarities between different topic’s word distributions.
Revised: We are also interested in topic models that involve more subtle topic correlations and similarities between different topic’s word distributions.

Original: To this end, we consider the Correlated Topic Model (CTM) (Blei & Lafferty, 2007) and the Pachinko Allocation Model (PAM) (Li & McCallum, 2006).
Revised: To this end, we consider the Correlated Topic Model (CTM) (Blei & Lafferty, 2007) and the Pachinko Allocation Model (PAM) (Li & McCallum, 2006).

Original: Our goal is to construct settings where the correlation between topics provide useful information for inference.
Revised: Our goal is to construct settings where the correlation between topics provide useful information for inference.

Original: To achieve this goal, in CTM, we construct groups of 4 topics.
Revised: To achieve this goal, in CTM, we construct groups of 4 topics.

Original: Within the group, topic pairs (0,2) and (1,3) are highly correlated (as speciﬁed by the prior), while topics 0 and 1 share many words (see Figure 1 for an illustration of the setting).
Revised: Within the group, topic pairs (0,2) and (1,3) are highly correlated (as speciﬁed by the prior), while topics 0 and 1 share many words (see Figure 1 for an illustration of the setting).

Original: 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 ﬁrst set of words are likely from topic 0, not topic 1.
Revised: 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 ﬁrst set of words are likely from topic 0, not topic 1.

Original: We construct such correlations in CTM 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.
Revised: We construct such correlations in CTM 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.

Original: We construct similar examples for PAM, see Appendix C for details.
Revised: We construct similar examples for PAM, see Appendix C for details.

Original: 5. S IMULATION S ETUP
Revised: 5. S IMULATION S ETUP

Original: Document generation Our documents are synthetically generated through the following steps: we ﬁrst construct a V × K topic-word matrix A .
Revised: Document generation Our documents are synthetically generated through the following steps: we ﬁrst construct a V × K topic-word matrix A .

Original: 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.
Revised: 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.

Original: from the word distribution given by Aw .
Revised: from the word distribution given by Aw .

Original: In our simulation, we set K = 20 , V = 5000 , λ ∈ { 30 , 60 } .
Revised: In our simulation, we set K = 20 , V = 5000 , λ ∈ { 30 , 60 } .

Original: Often A will be drawn from a Dirichlet distribution Dir ( α/K ) , and we take α = 1 , 3 , 5 , 7 , 9 to vary the difﬁculty level of the inference problem, where a larger α introduces more similarity between topics.
Revised: Often A will be drawn from a Dirichlet distribution Dir ( α/K ) (see Appendix C for details), and we take α = 1 , 3 , 5 , 7 , 9 to vary the difﬁculty level of the inference problem, where a larger α introduces more similarity between topics.

Original: Neural network models
Revised: Neural network models

Original: We consider two different network architectures as f in the experiments: (1) Fully-connected network with residual connection, with bag-of-words document representation as the input since the order does not matter in topic models; (2) Attention-based architecture.
Revised: We consider two different network architectures as f in the experiments: (1) Fully-connected network with residual connection, with bag-of-words document representation as the input since the order does not matter in topic models; (2) Attention-based architecture.

Original: Due to the architecture of attention, the input needs to be a sequence of words instead of bag-of-words vector.
Revised: Due to the architecture of attention, the input needs to be a sequence of words instead of bag-of-words vector.

Original: Thus, we transform bag-of-words representation into full document by repeating each word by its frequency and concatenating them in random order, and then use the full document as the input.
Revised: Thus, we transform bag-of-words representation into full document by repeating each word by its frequency and concatenating them in random order, and then use the full document as the input.

Original: More speciﬁcally, the attention-based architecture 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.
Revised: More speciﬁcally, the attention-based architecture 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.

Original: The network’s second to last layer averages over the outputs of the last transformer block and the ﬁnal layer projects it to a V -dimensional word distribution.
Revised: The network’s second to last layer averages over the outputs of the last transformer block and the ﬁnal layer projects it to a V -dimensional word distribution.

Original: See more details in Section C.5.
Revised: See more details in Section C.5.

Original: We ﬁnd 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.
Revised: We ﬁnd 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.

Original: 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.
Revised: 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.

Original: The loss function is the reconstruction-based objective (1) with t = 1 , i.e., we want the model to predict one missing word.
Revised: The loss function is the reconstruction-based objective (1) with t = 1 , i.e., we want the model to predict one missing word.

Original: To reduce the variance during the training, in each training document, we sampled 6 words as the prediction target, hide them from the model and ask the model to predict each word separately.
Revised: To reduce the variance during the training, in each training document, we sampled 6 words as the prediction target, hide them from the model and ask the model to predict each word separately.

Original: The trained model is evaluated on test documents generated from the same topic prior as the training documents.
Revised: The trained model is evaluated on test documents generated from the same topic prior as the training documents.

Original: We use 5,000 test documents for the pure-topic prior and 200 test documents for the remaining priors.
Revised: We use 5,000 test documents for the pure-topic prior and 200 test documents for the remaining priors.

Original: 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 truthexpected topic posterior distribution.
Revised: For each test document, we use MCMC (see Appendix C.9) assuming the document’s correct topic prior to approximate the ground truth topic posterior.

Original: To get the posterior estimations from SSL approach, we multiply the pseudo-inverse of the topic matrix A with the model output f ( x ) to get the estimated posterior mean vector A † f ( x ) as we explained in Section 3 (and also Section A.1).
Revised: To get the posterior estimations from SSL approach, we multiply the pseudo-inverse of the topic matrix A with the model output f ( x ) to get the estimated posterior mean vector A † f ( x ) as we explained in Section 3 (and also Section A.1).

Original: Then, we measure topic posterior recovery loss as the Total Variation (TV) distance between the recovered topic posterior mean vector and the ground truth topic posterior mean vector.
Revised: Then, we measure topic posterior recovery loss as the Total Variation (TV) distance between the recovered topic posterior mean vector and the ground truth topic posterior mean vector.

Original: 5. E XPERIMENT R ESULTS Topic posterior recovery
Revised: 5. E XPERIMENT R ESULTS Topic posterior recovery

Original: As illustrated in Table 1, after 200 training epochs, our model can accurately recover the topic posterior mean vector.
Revised: As illustrated in Table 1, after 200 training epochs, our model can accurately recover the topic posterior mean vector.

Original: Meanwhile, it can be observed that for larger values of the Dirichlet hyperparameter α , the recovery loss gets higher.
Revised: Meanwhile, it can be observed that for larger values of the Dirichlet hyperparameter α , the recovery loss gets higher.

Original: This is expected because higher α leads to more similar topics, which makes the learning problem more difﬁcult.
Revised: This is expected because higher α leads to more similar topics, which makes the learning problem more difﬁcult.

Original: This effect is also captured by Deﬁnition 3.2 and we show the computed condition numbers in Section C.6.
Revised: This effect is also captured by Deﬁnition 3.2 and we show the computed condition numbers in Section C.6.

Original: Major topic recovery
Revised: Major topic recovery

Original: We examine the extent our recovered topic posterior captures the major topics in a given document.
Revised: We examine the extent our recovered topic posterior captures the major topics in a given document.

Original: For pure-topic and LDA documents, we measure major topic recovery accuracy of correctly estimating the topic with largest proportion.
Revised: For pure-topic and LDA documents, we measure major topic recovery accuracy of correctly estimating the topic with largest proportion.

Original: Considering that topics in CTM and PAM documents are correlated in pairs, we measure the major topic recovery accuracy for CTM and PAM as the top-2 topic overlap rate on their test documents.
Revised: Considering that topics in CTM and PAM documents are correlated in pairs, we measure the major topic recovery accuracy for CTM and PAM as the top-2 topic overlap rate on their test documents.

Original: Table 1 shows that our algorithm is successful throughout all settings we consider.
Revised: Table 1 shows that our algorithm is successful throughout all settings we consider.

Original: Robustness of self-supervised learning We compare the self-supervised approach to traditional topic inference (MCMC with a specifed topic prior, the results on variational inference are reported in Section C.3).
Revised: Robustness of self-supervised learning We compare the self-supervised approach to traditional topic inference (MCMC with a specifed topic prior, the results on variational inference are reported [Figure]

Original: Speciﬁcally, for each α value, we take 200 test documents from each category of documents, and we run posterior inference assuming a speciﬁc topic model and calculate the TV
Revised: [Table]

Original: [Table] [Table]
Revised: [Table] [Table]

Original:  distance between this posterior mean vector and the ground truth posterior mean vector.
Revised: in Section C.3). Speciﬁcally, for each α value, we take 200 test documents from each category of documents, and we run posterior inference assuming a speciﬁc topic model and calculate the TV distance between this posterior mean vector and the ground truth posterior mean vector.

Original: We exclude assuming pure-topic prior from our comparison because it gives invalid results for documents with mixed topics.
Revised: We exclude assuming pure-topic prior from our comparison because it gives invalid results for documents with mixed topics.

Original: For α = 1 , as shown in Table 2, the topic posterior recovered from SSL approach is closer to the ground truth topic posterior than that recovered from a misspeciﬁed topic prior.
Revised: For α = 1 , as shown in Table 2, the topic posterior recovered from SSL approach is closer to the ground truth topic posterior than that recovered from a misspeciﬁed topic prior.

Original: We also report the major topic recovery accuracy for SSL approach against topic inference using both correct and incorrect topic model in Table 2.
Revised: We also report the major topic recovery accuracy for SSL approach against topic inference using both correct and incorrect topic model in Table 2.

Original: The major topic recovery rate of SSL approach is similar to that of posterior inference assuming the correct prior.
Revised: The major topic recovery rate of SSL approach is similar to that of posterior inference assuming the correct prior.

Original: For pure topic model, all four methods get 100% major topics recovery rate on pure-topic documents.
Revised: For pure topic model, all four methods get 100% major topics recovery rate on pure-topic documents.

Original: For LDA, again all four methods perform similarly well since the instance is not difﬁcult.
Revised: For LDA, again all four methods perform similarly well since the instance is not difﬁcult.

Original: However we observe signiﬁcant difference for more complicated CTM and PAM models.
Revised: However we observe signiﬁcant difference for more complicated CTM and PAM models.

Original: For these models, the SSL approach performs similarly to posterior inference using the correct model, and both of them are signiﬁcantly better than posterior inference using misspeciﬁed models.
Revised: For these models, the SSL approach performs similarly to posterior inference using the correct model, and both of them are signiﬁcantly better than posterior inference using misspeciﬁed models.

Original: The results for α = 3 , 5 , 7 , 9 are presented in Section C.
Revised: The results for α = 3 , 5 , 7 , 9 are presented in Section C.

Original: The comparison further reveals the robustness of self-superivised learning.
Revised: The comparison further reveals the robustness of self-superivised learning.

Original:  E XPERIMENTS ON R EAL D ATA
Revised:  E XPERIMENTS ON R EAL D ATA

Original: The self-supervised learning approach studied in this paper is simpliﬁed for theoretical analysis as it uses a bag-of-words approach.
Revised: The self-supervised learning approach studied in this paper is simpliﬁed for theoretical analysis as it uses a bag-of-words approach.

Original: In this section, we show that although this simpliﬁed approach is not state-of-art, it has reasonable performance compared to several baselines.
Revised: In this section, we show that although this simpliﬁed approach is not state-of-art, it has reasonable performance compared to several baselines.

Original: Experiment Setup We use 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.
Revised: Experiment Setup We use 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.

Original: Each category has 30,000 samples in the training set and 19,000 samples in the testing set.
Revised: Each category has 30,000 samples in the training set and 19,000 samples in the testing set.

Original: The number of words in the vocabulary is about 16,700.
Revised: The number of words in the vocabulary is about 16,700.

Original: We generally follow the experiment setup done by Tosh et al.
Revised: We generally follow the experiment setup done by Tosh et al.

Original: (2021a)
Revised: (2021a)

Original: (see more details in Section D.1), but the representation we used is generated by our reconstruction-based objective.
Revised: (see more details in Section D.1), but the representation we used is generated by our reconstruction-based objective.

Original: Experiment results on two additional datasets are provided in Appendix D.
Revised: Experiment results on two additional datasets are provided in Appendix D.

Original: Our training follows the semi-supervised setup that involves two major phases: unsupervised phase and supervised phase.
Revised: Our training follows the semi-supervised setup that involves two major phases: unsupervised phase and supervised phase.

Original: In the unsupervised phase, we train the self-supervised model with reconstruction-based objective on most of training data.
Revised: In the unsupervised phase, we train the self-supervised model with reconstruction-based objective on most of training data.

Original: In the supervised phase, using the document representation generated by the SSL model as input, we train a linear classiﬁer to classify the category of the documents using multi-class logistics regression (by Scikit-learn (Pedregosa et al., 2011)) with 3-fold cross validation for parameter tuning ( l 2 regularization term and solver).
Revised: In the supervised phase, using the document representation generated by the SSL model as input, we train a linear classiﬁer to classify the category of the documents using multi-class logistics regression (by Scikit-learn (Pedregosa et al., 2011)) with 3-fold cross validation for parameter tuning ( l 2 regularization term and solver).

Original: We measure the topic recovery accuracy on test set to compare the performance of different methods.
Revised: We measure the topic recovery accuracy on test set to compare the performance of different methods.

Original: We choose residual blocks as the basic building block for our neural network architecture, with varying width and depth.
Revised: We choose residual blocks as the basic building block for our neural network architecture, with varying width and depth.

Original: The network used for Figure 2 has 3 residual blocks and a width of 4096 neurons per layer.
Revised: The network used for Figure 2 has 3 residual blocks and a width of 4096 neurons per layer.

Original: Other detailed hyperparameters and ablation studies can be found in Appendix D.3.
Revised: Other detailed hyperparameters and ablation studies can be found in Appendix D.3.

Original: We train the model using reconstruction-based objective (1) with t = 1 using the same variance-reduction technique in synthetic experiments.
Revised: We train the model using reconstruction-based objective (1) with t = 1 using the same variance-reduction technique in synthetic experiments.

Original: To extract our representation, we use an identity matrix to replace the last layer.
Revised: To extract our representation, we use an identity matrix to replace the last layer.

Original: That is, we take a softmax function on top of the second-to-last layer of the neural network.
Revised: That is, we take a softmax function on top of the second-to-last layer of the neural network.

Original: This is different from our theory but we found that it effectively reduces the high output dimension and improves performance.
Revised: This is different from our theory but we found that it effectively reduces the high output dimension and improves performance.

Original: We include more details in Section D.2 in Appendix.
Revised: We include more details in Section D.2 in Appendix.

Original: Experiment Results The performance of 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 classiﬁer in supervised learning phase.
Revised: Experiment Results The performance of 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 classiﬁer in supervised learning phase.

Original: The representation RBL performs better than both bag-of-words (BOW),Word2vec and LDA baselines (see more details about baselines in Section D.2).
Revised: The representation RBL performs better than both bag-of-words (BOW),Word2vec and LDA baselines (see more details about baselines in Section D.2).

Original: Notably, Word2vec representation is inferior only by a small margin, and works well in general even when limited training samples are provided.
Revised: Notably, Word2vec representation is inferior only by a small margin, and works well in general even when limited training samples are provided.

Original: On the contrary, although BOW representation has a decent test accuracy when training samples are abundant, it performs signiﬁcantly poorly on smaller training set.
Revised: On the contrary, although BOW representation has a decent test accuracy when training samples are abundant, it performs signiﬁcantly poorly on smaller training set.

Original: We also compare our result to previous noise contrastive estimation (NCE) representation benchmark achieved by Tosh et al.
Revised: We also compare our result to previous noise contrastive estimation (NCE) representation benchmark achieved by Tosh et al.

Original: (2021a) using contrastive objective, where their reported best accuracy is around 87 .
Revised: (2021a) using contrastive objective, where their reported best accuracy is around 87 .

Original: 5% when full 4000 training samples are used, slightly higher than our RBL corresponding test accuracy of 87 .
Revised: 5% when full 4000 training samples are used, slightly higher than our RBL corresponding test accuracy of 87 .

Original: In our attempt to reproduce their result, parameter tuning yields the best accuracy of 86% when full training samples were used.
Revised: In our attempt to reproduce their result, parameter tuning yields the best accuracy of 86% when full training samples were used.

Original: We plot our own reproduced results on in
Revised: We plot our own reproduced results on in

Original: Figure 2 since parameter tuning does not achieve a close benchmark to their result.
Revised: Figure 2 since parameter tuning does not achieve a close benchmark to their result.

Original:  C ONCLUSION AND L IMITATIONS
Revised:  C ONCLUSION AND L IMITATIONS

Original: “All models are wrong but some are useful.”
Revised: “All models are wrong but some are useful.”

Original: If one self-supervised objective can capture all models, then it would be able to extract useful information.
Revised: If one self-supervised objective can capture all models, then it would be able to extract useful information.

Original: 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.
Revised: 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.

Original: Our results generalized previous work (Tosh et al., 2021a) to both contrastive learning and reconstruction-based learning and our techniques allow us to depend on weaker assumptions.
Revised: Our results generalized previous work (Tosh et al., 2021a) to both contrastive learning and reconstruction-based learning and our techniques allow us to depend on weaker assumptions.

Original: We also empirically showed that the reconstruction-based learning performs better than the posterior inference under misspeciﬁed models, and it can provide useful representation for the topic inference problem.
Revised: We also empirically showed that the reconstruction-based learning performs better than the posterior inference under misspeciﬁed models, and it can provide useful representation for the topic inference problem.

Original: Our theoretical analysis is limited to a bag-of-words setting, which also limits its empirical performance.
Revised: Our theoretical analysis is limited to a bag-of-words setting, which also limits its empirical performance.

Original: Extending the work to more complicated models is an immediate open problem.
Revised: Extending the work to more complicated models is an immediate open problem.

Original: A O MITTED P ROOFS IN S ECTION
Revised: A O MITTED P ROOFS IN S ECTION

Original: In this section, we give the omitted proofs in Section 3.
Revised: In this section, we give the omitted proofs in Section 3.

Original: We ﬁrst give a warmup example to illustrate our proof idea in Section A.1.
Revised: We ﬁrst give a warmup example to illustrate our proof idea in Section A.1.

Original: Then we give the proof of our main result (Theorem 3.1) in Section A. and the proof of robust version (Theorem 3.3) in Section A.3.
Revised: Then we give the proof of our main result (Theorem 3.1) in Section A. and the proof of robust version (Theorem 3.3) in Section A.3.

Original: A. W ARM -U P E XAMPLE : RECONSTRUCT THE UNKNOWN WORD
Revised: A. W ARM -U P E XAMPLE : RECONSTRUCT THE UNKNOWN WORD

Original: 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).
Revised: 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).

Original: The following result is the special case of Theorem 3.1 with t = 1 , which shows that the learned representation is able to give any linear function of the topic posterior.
Revised: The following result is the special case of Theorem 3.1 with t = 1 , which shows that the learned representation is able to give any linear function of the topic posterior.

Original: Theorem A.1.
Revised: Theorem A.1.

Original: Consider the general topic model setting as Deﬁnition 2.1, suppose topic-word matrix
Revised: Consider the general topic model setting as Deﬁnition 2.1, suppose topic-word matrix

Original: A satisﬁes rank ( A ) =
Revised: A satisﬁes rank ( A ) =

Original: K and function f minimizes the reconstruction-based objective (1) .
Revised: K and function f minimizes the reconstruction-based objective (1) .

Original: Then, for any document x and its posterior w | ( x, A ) , any linear function P ( w ) is linear in f ( x ) , that is for any P ( w )
Revised: Then, for any document x and its posterior w | ( x, A ) , any linear function P ( w ) is linear in f ( x ) , that is for any P ( w )

Original: =
Revised: =

Original: β (cid:62) w , there exists a θ ∈ R V such that for all documents x
Revised: β (cid:62) w , there exists a θ ∈ R V such that for all documents x

Original: [Equation]
Revised: [Equation]

Original: Following the proof idea described in Section 3, we ﬁrst give a characterization of the optimal function f .
Revised: Following the proof idea described in Section 3, we ﬁrst give a characterization of the optimal function f .

Original: The following lemma shows that f ( x ) must be the word posterior vector given the document x .
Revised: The following lemma shows that f ( x ) must be the word posterior vector given the document x .

Original: 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.
Revised: 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.

Original: Recall that our vocabulary set is V = { v 1 , . .
Revised: Recall that our vocabulary set is V = { v 1 , . .

Original: . , v V } and x unsup =
Revised: . , v V } and x unsup =

Original: ( x, y ) is the given document, where x is unmasked part and y is word marked as unknown.
Revised: ( x, y ) is the given document, where x is unmasked part and y is word marked as unknown.

Original: Lemma A.2.
Revised: Lemma A.2.

Original: If f minimizes the reconstruction-based objective (1) , then we have for all document x
Revised: If f minimizes the reconstruction-based objective (1) , then we have for all document x

Original: [Equation]
Revised: [Equation]

Original: Based on the above lemma, we are able to prove Theorem A.1.
Revised: Based on the above lemma, we are able to prove Theorem A.1.

Original: It is easy to see that the word posterior distribution is A E w [ w | x ] , so we have f ( x ) =
Revised: It is easy to see that the word posterior distribution is A E w [ w | x ] , so we have f ( x ) =

Original: A E w [ w | x ] .
Revised: A E w [ w | x ] .

Original: Since the columns of A are linearly independent, we know E w [ w | x ]
Revised: Since the columns of A are linearly independent, we know E w [ w | x ]

Original: = A † f ( x ) .
Revised: = A † f ( x ) .

Original: Thus, for any linear function β (cid:62) w , there exists θ =
Revised: Thus, for any linear function β (cid:62) w , there exists θ =

Original: ( A † ) (cid:62) β such that E w [ β (cid:62) w | x ] = θ (cid:62) f
Revised: ( A † ) (cid:62) β such that E w [ β (cid:62) w | x ] = θ (cid:62) f

Original: ( x ) .
Revised: ( x ) .

Original: The formal proof of the general case is given in
Revised: The formal proof of the general case is given in

Original: Section A.2.
Revised: Section A.2.

Original: A.1. P ROOF OF L EMMA A. 
Revised: A.1. P ROOF OF L EMMA A.

Original: Instead of focusing on the t = 1 case, we directly give the corresponding lemma of Lemma A.2 for general t .
Revised: Instead of focusing on the t = 1 case, we directly give the corresponding lemma of Lemma A.2 for general t .

Original: Recall that x unsup =
Revised: Recall that x unsup =

Original: ( x, y ) is the given document, x is the unmasked part, y =
Revised: ( x, y ) is the given document, x is the unmasked part, y =

Original: ( y 1 , . . .
Revised: ( y 1 , . . .

Original: , y t ) is the t unknown words that we want to predict and V = { v 1 , . .
Revised: , y t ) is the t unknown words that we want to predict and V = { v 1 , . .

Original: . , v V } is the set of vocabulary.
Revised: . , v V } is the set of vocabulary.

Original: It shows that the optimal f is the t words posterior distribution.
Revised: It shows that the optimal f is the t words posterior distribution.

Original: Note that the words posterior distribution is linear in the topic posterior of t -th moment, which is useful for the later analysis.
Revised: Note that the words posterior distribution is linear in the topic posterior of t -th moment, which is useful for the later analysis.

Original: Lemma A.3 (Lemma A.2, General Case) .
Revised: Lemma A.3 (Lemma A.2, General Case) .

Original: If f minimizes the reconstruction-based objective (1) , then we have for all document x
Revised: If f minimizes the reconstruction-based objective (1) , then we have for all document x

Original: [Equation]
Revised: [Equation]

Original: Proof.
Revised: Proof.

Original: Denote the word posterior distribution as p ∗ ( x ) .
Revised: Denote the word posterior distribution as p ∗ ( x ) .

Original: We will show f ( x ) =
Revised: We will show f ( x ) =

Original: p ∗ ( x ) .
Revised: p ∗ ( x ) .

Original: By the law of total expectation, we have
Revised: By the law of total expectation, we have

Original: [Equation]
Revised: [Equation]

Original: We know the probability of y given x is p ∗ ( x ) .
Revised: We know the probability of y given x is p ∗ ( x ) .

Original: Since (cid:96) is cross-entropy loss, we have
Revised: Since (cid:96) is cross-entropy loss, we have

Original: [Equation]
Revised: [Equation]

Original: where the equality is obtained when f ( x )
Revised: where the equality is obtained when f ( x )

Original: = p ∗ ( x ) .
Revised: = p ∗ ( x ) .

Original: Thus, when f minimizes the reconstruction-based objective (1), we have f ( x ) =
Revised: Thus, when f minimizes the reconstruction-based objective (1), we have f ( x ) =

Original: p ∗ ( x ) .
Revised: p ∗ ( x ) .

Original: A. P ROOF OF T HEOREM 3. 
Revised: A. P ROOF OF T HEOREM 3.

Original: In this section, we give the proof of Theorem 3.1.
Revised: In this section, we give the proof of Theorem 3.1.

Original: Theorem A.1 is covered by Theorem 3.1 as the special case of t = 1 .
Revised: Theorem A.1 is covered by Theorem 3.1 as the special case of t = 1 .

Original: The proof follows the same idea as for Theorem A.1.
Revised: The proof follows the same idea as for Theorem A.1.

Original: We ﬁrst show that f ( x ) is the t -words posterior (Lemma A.3).
Revised: We ﬁrst show that f ( x ) is the t -words posterior (Lemma A.3).

Original: 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 .
Revised: 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 .

Original: The result follows by the observation that for a polynomial P ( w ) with degree at most t , E w [ P ( w ) | x ] is a linear function of the topic posterior of t -th moment.
Revised: The result follows by the observation that for a polynomial P ( w ) with degree at most t , E w [ P ( w ) | x ] is a linear function of the topic posterior of t -th moment.

Original: Theorem 3.1 (Main Result) .
Revised: Theorem 3.1 (Main Result) .

Original: Consider the general topic model setting as Deﬁnition 2.1, suppose topic-word matrix A satisﬁes rank ( A ) =
Revised: Consider the general topic model setting as Deﬁnition 2.1, suppose topic-word matrix A satisﬁes rank ( A ) =

Original: K and function f minimizes the reconstruction-based objective (1) .
Revised: K and function f minimizes the reconstruction-based objective (1) .

Original: Then, for any document x and its posterior w | ( x, A ) , the posterior mean of any polynomial P ( w ) with degree at most t is linear in f ( x ) .
Revised: Then, for any document x and its posterior w | ( x, A ) , the posterior mean of any polynomial P ( w ) with degree at most t is linear in f ( x ) .

Original: That is, there exists a θ ∈ R V t such that for all documents x
Revised: That is, there exists a θ ∈ R V t such that for all documents x

Original: [Equation]
Revised: [Equation]

Original: Proof.
Revised: Proof.

Original: Since f is the minimizer of reconstruction-based objective (1), by Lemma A.3 we know given an input document x ,
Revised: Since f is the minimizer of reconstruction-based objective (1), by Lemma A.3 we know given an input document x ,

Original: [Equation]
Revised: [Equation]

Original: We will show that f ( x ) is linear in the topic posterior.
Revised: We will show that f ( x ) is linear in the topic posterior.

Original: In the following, we focus on [ f ( x )] y 1 ,...,y t , which is word posterior probability of the t unknown words being y 1 , . . .
Revised: In the following, we focus on [ f ( x )] y 1 ,...,y t , which is word posterior probability of the t unknown words being y 1 , . . .

Original: , y t ∈ V given document x .
Revised: , y t ∈ V given document x .

Original: Recall S K be the K − 1 dimensional probability simplex.
Revised: Recall S K be the K − 1 dimensional probability simplex.

Original: By the law of total probability, we have [ f ( x )] y 1 ,y 2 ,...,y t
Revised: By the law of total probability, we have [ f ( x )] y 1 ,y 2 ,...,y t

Original: [Equation]
Revised: [Equation]

Original: P
Revised: P

Original: [Equation]
Revised: [Equation]

Original: y 1 , y
Revised: y 1 , y

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: y t | x )
Revised: y t | x )

Original: [Equation]
Revised: [Equation]

Original: Note that z
Revised: Note that z

Original: i is the topic of word y i .
Revised: i is the topic of word y i .

Original: Since we consider the general topic model as Deﬁnition 2.1, we know
Revised: Since we consider the general topic model as Deﬁnition 2.1, we know

Original: P ( y 1 , y 2 , . . .
Revised: P ( y 1 , y 2 , . . .

Original: , y t | z 1 , z 2 , ..., z t , w ) =
Revised: , y t | z 1 , z 2 , ..., z t , w ) =

Original: P ( y 1 , y 2 , . .
Revised: P ( y 1 , y 2 , . .

Original: , y t | z 1 , z 2 , ...
Revised: , y t | z 1 , z 2 , ...

Original: [Equation]
Revised: [Equation]

Original: z t ) =
Revised: z t ) =

Original: t (cid:89) i =
Revised: t (cid:89) i =

Original: P ( y i |
Revised: P ( y i |

Original: z i )
Revised: z i )

Original: [Equation]
Revised: [Equation]

Original: t (cid:89) i =
Revised: t (cid:89) i =

Original: A y i ,z i , where A is the topic-word matrix.
Revised: A y i ,z i , where A is the topic-word matrix.

Original: Hence,
Revised: Hence,

Original: [Equation]
Revised: [Equation]

Original: Recall that the topic posterior tensor is W post
Revised: Recall that the topic posterior tensor is W post

Original: = E w [ w ⊗ t | x ] ∈ R K × ...
Revised: = E w [ w ⊗ t | x ] ∈ R K × ...

Original: × K , where each entry [ W post ] z 1 ,...,z k = P
Revised: × K , where each entry [ W post ] z 1 ,...,z k = P

Original: ( z i is the topic of word y i for i ∈
Revised: ( z i is the topic of word y i for i ∈

Original: [ t ] | x ) =
Revised: [ t ] | x ) =

Original: E w [ w z 1 . . .
Revised: E w [ w z 1 . . .

Original: w z t | x ] .
Revised: w z t | x ] .

Original: Therefore, [ f ( x )] y 1 ,y 2 ,...,y t
Revised: Therefore, [ f ( x )] y 1 ,y 2 ,...,y t

Original: [Equation]
Revised: [Equation]

Original: where ⊗ is the Kronecker product, A =
Revised: where ⊗ is the Kronecker product, A =

Original: A ⊗ A ⊗ · · · ⊗ A ∈ R V t × K t and vec ( W post ) ∈
Revised: A ⊗ A ⊗ · · · ⊗ A ∈ R V t × K t and vec ( W post ) ∈

Original: R K t is the vectorization of W post .
Revised: R K t is the vectorization of W post .

Original: Note that A † =
Revised: Note that A † =

Original: A † ⊗ · · · ⊗ A † , so we have vec ( W post ) =
Revised: A † ⊗ · · · ⊗ A † , so we have vec ( W post ) =

Original: A † f ( x ) .
Revised: A † f ( x ) .

Original: Since P ( w ) is a polynomial of degree at most t , we know there exists β such that E w [ P ( w ) | x ] = β (cid:62) vec ( W post ) .
Revised: Since P ( w ) is a polynomial of degree at most t , we know there exists β such that E w [ P ( w ) | x ] = β (cid:62) vec ( W post ) .

Original: Thus, let θ =
Revised: Thus, let θ =

Original: ( A † ) (cid:62) β , we have E w [ P ( w ) | x ] = β (cid:62) vec ( W post )
Revised: ( A † ) (cid:62) β , we have E w [ P ( w ) | x ] = β (cid:62) vec ( W post )

Original: = θ (cid:62) f
Revised: = θ (cid:62) f

Original: ( x ) .
Revised: ( x ) .

Original: A. P ROOF OF T HEOREM 3. 
Revised: A. P ROOF OF T HEOREM 3.

Original: Similar to Lemma A.2, we can show that the representation f ( x ) is close to t words posterior.
Revised: Similar to Lemma A.2, we can show that the representation f ( x ) is close to t words posterior.

Original: Then, the key lemma to obtain Theorem 3.3 is the following result.
Revised: Then, the key lemma to obtain Theorem 3.3 is the following result.

Original: It suggests that if function f is (cid:15) -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 ( √ (cid:15) ) error.
Revised: It suggests that if function f is (cid:15) -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 ( √ (cid:15) ) error.

Original: For example, we will choose B =
Revised: For example, we will choose B =

Original: A † when t = 1 .
Revised: A † when t = 1 .

Original: See the proof of Theorem 3.3 for the choice of B in the general case.
Revised: See the proof of Theorem 3.3 for the choice of B in the general case.

Original: Lemma A.4.
Revised: Lemma A.4.

Original: For cross-entropy loss (cid:96) and any two probability vectors p, p ∗ , if (cid:96) ( p, p ∗ ) ≤ min p : (cid:80) i p i =1 ,p ≥ 0 (cid:96) ( p, p ∗ )
Revised: For cross-entropy loss (cid:96) and any two probability vectors p, p ∗ , if (cid:96) ( p, p ∗ ) ≤ min p : (cid:80) i p i =1 ,p ≥ 0 (cid:96) ( p, p ∗ )

Original: + (cid:15) , then for any matrix B we have (cid:107)
Revised: + (cid:15) , then for any matrix B we have (cid:107)

Original: Bp − Bp ∗ (cid:107) 1 ≤ κ ( B ) √ 2 (cid:15) .
Revised: Bp − Bp ∗ (cid:107) 1 ≤ κ ( B ) √ 2 (cid:15) .

Original: Proof.
Revised: Proof.

Original: Recall that (cid:96) is cross-entropy loss, we know
Revised: Recall that (cid:96) is cross-entropy loss, we know

Original: [Equation]
Revised: [Equation]

Original: By Pinsker’s Inequality, we have that
Revised: By Pinsker’s Inequality, we have that

Original: [Equation]
Revised: [Equation]

Original: Then, for any matrix B we have
Revised: Then, for any matrix B we have

Original: [Equation]
Revised: [Equation]

Original: where we use the deﬁnition of (cid:96) 1 condition number (Deﬁnition 3.2).
Revised: where we use the deﬁnition of (cid:96) 1 condition number (Deﬁnition 3.2).

Original: We now are ready to give the proof of Theorem 3.3, which is based on Lemma A.4 and the proof of
Revised: We now are ready to give the proof of Theorem 3.3, which is based on Lemma A.4 and the proof of

Original: Theorem 3.1.
Revised: Theorem 3.1.

Original: Theorem 3.3 (Robust Version) .
Revised: Theorem 3.3 (Robust Version) .

Original: Consider the general topic model setting as Deﬁnition 2.1, suppose topic-word matrix A satisﬁes rank ( A ) =
Revised: Consider the general topic model setting as Deﬁnition 2.1, suppose topic-word matrix A satisﬁes rank ( A ) =

Original: K and function f satisﬁes L reconst ( f ) ≤ L ∗ reconst + (cid:15) for some (cid:15) > 0 .
Revised: K and function f satisﬁes L reconst ( f ) ≤ L ∗ reconst + (cid:15) for some (cid:15) > 0 .

Original: Then, for any document x and its posterior w | ( x, A ) , the posterior mean of any polynomial
Revised: Then, for any document x and its posterior w | ( x, A ) , the posterior mean of any polynomial

Original: P ( w ) with degree at most t is approximately linear in f ( x ) .
Revised: P ( w ) with degree at most t is approximately linear in f ( x ) .

Original: That is, there exists a θ ∈ R V t such that
Revised: That is, there exists a θ ∈ R V t such that

Original: [Equation]
Revised: [Equation]

Original: where E w [ P ( w ) | x, A ] = β (cid:62) vec ( W post ) .
Revised: where E w [ P ( w ) | x, A ] = β (cid:62) vec ( W post ) .

Original: Proof.
Revised: Proof.

Original: Recall from the proof of Theorem 3.1 that E w [ P ( w ) | x ] =
Revised: Recall from the proof of Theorem 3.1 that E w [ P ( w ) | x ] =

Original: β (cid:62) vec ( W post )
Revised: β (cid:62) vec ( W post )

Original: =
Revised: =

Original: β (cid:62) A † f ∗ ( x ) , where A † =
Revised: β (cid:62) A † f ∗ ( x ) , where A † =

Original: A † ⊗ · · · ⊗ A † and κ ( A † )
Revised: A † ⊗ · · · ⊗ A † and κ ( A † )

Original: =
Revised: =

Original: κ t ( A † ) .
Revised: κ t ( A † ) .

Original: Same as in the proof of Theorem 3.1, let θ =
Revised: Same as in the proof of Theorem 3.1, let θ =

Original: ( A † ) (cid:62) β .
Revised: ( A † ) (cid:62) β .

Original: We have
Revised: We have

Original: [Equation]
Revised: [Equation]

Original: We are going to use Lemma A.4 to bound (cid:13)(cid:13)
Revised: We are going to use Lemma A.4 to bound (cid:13)(cid:13)

Original: A † f ∗ ( x ) − A † f ( x ) (cid:13)(cid:13) 1 .
Revised: A † f ∗ ( x ) − A † f ( x ) (cid:13)(cid:13) 1 .

Original: By Lemma A.3 we know f ∗ ( x ) is the word posterior, so f ∗ ( x ) =
Revised: By Lemma A.3 we know f ∗ ( x ) is the word posterior, so f ∗ ( x ) =

Original: E y | x [ y | x ] .
Revised: E y | x [ y | x ] .

Original: Since (cid:96) is cross-entropy loss, we know
Revised: Since (cid:96) is cross-entropy loss, we know

Original: E y | x [ (cid:96) ( f ( x ) , y ) − (cid:96) ( f ∗ ( x ) , y )
Revised: E y | x [ (cid:96) ( f ( x ) , y ) − (cid:96) ( f ∗ ( x ) , y )

Original: | x ] = [Equation] y | x  
Revised: | x ] = E y | x   (cid:88)

Original: [Equation] [Equation] [Equation]
Revised: k ∈

Original:  log (cid:18)[ [Equation]
Revised: [ V t ] y k log (cid:18)

Original:  ( x [Equation]
Revised: [ f ∗ ( x )]

Original: [Equation] [ f ( x )] [Equation] (cid:19) (cid:12)(cid:12)(cid:12)(cid:12) x  
Revised: k [ f ( x )] k (cid:19) (cid:12)(cid:12)(cid:12)(cid:12) x   =

Original: [Equation]
Revised: [Equation]

Original: ( f
Revised: ( f

Original: [Equation]
Revised: [Equation]

Original: ( x [Equation] || f ( x )) , which implies (cid:96) ( f ( x ) , f ∗ ( x ))
Revised: ( x ) || f ( x )) , which implies (cid:96) ( f ( x ) , f ∗ ( x ))

Original: − min p (cid:96) ( p, f ∗ ( x )) ≤ (cid:15) .
Revised: − min p (cid:96) ( p, f ∗ ( x )) ≤ (cid:15) .

Original: Therefore, by Lemma A.4 with B = A † , we know
Revised: Therefore, by Lemma A.4 with B = A † , we know

Original: [Equation]
Revised: [Equation]

Original: Given the desired bound, we have
Revised: Given the desired bound, we have

Original: [Equation]
Revised: [Equation]

Original: B O MITTED P ROOFS IN S ECTION
Revised: A. D ISCUSSION ON THE EFFECTIVE DIMENSION OF THE REPRESENTATION

Revised: Although the dimension of f ( x ) is V t , its “effective dimension" is actually K t , which is much smaller than V t (the number of topic K usually is much smaller than the size of vocabulary V ). To see this, let’s consider the simple case of t = 1 . As we discussed in the paragraph below Theorem 3.1, we have f ( x ) =

Revised: A E [ w | x, A ] , where A ∈ R V × K is the topic-word matrix and E [ w | x, A ] is the posterior mean of w ∈ R K .

Revised: In words, this means that f ( x ) is completely determined by a linear transformation of a K -dimensional vector.

Revised: Therefore, when we consider the feature of n documents as a matrix [ f ( x (1) ) , ..., f ( x ( n ) )] ∈ R V × n , the rank of this matrix is at most K and one can get these

Revised: K -dimensional features by running standard dimension reduction method (e.g., SVD) on the original feature matrix.

Revised: B O MITTED P ROOFS IN S ECTION

Original: In this section, we give the omitted proofs in Section 4.
Revised: In this section, we give the omitted proofs in Section 4.

Original: We give the proof of Theorem 4.1 in
Revised: We give the proof of Theorem 4.1 in

Original: Section B.1 and proof of Lemma B.1 in Section B.2.
Revised: Section B.1 and proof of Lemma B.1 in Section B.2.

Original: B. P ROOF OF T HEOREM 4. 
Revised: B. P ROOF OF T HEOREM 4.

Original: Before presenting the proof of Theorem 4.1, we ﬁrst give a characterization of the representation g ( x, x (cid:48) ) , which would be useful in the later analysis.
Revised: Before presenting the proof of Theorem 4.1, we ﬁrst give a characterization of the representation g ( x, x (cid:48) ) , which would be useful in the later analysis.

Original: Note that this the same as the one shown in
Revised: Note that this the same as the one shown in

Original: Tosh et al.
Revised: Tosh et al.

Original: (2021a).
Revised: (2021a).

Original: We provide its proof for completeness in Section B.2.
Revised: We provide its proof for completeness in Section B.2.

Original: Lemma B.1.
Revised: Lemma B.1.

Original: If f minimizes the contrastive objective (2) , then we have
Revised: If f minimizes the contrastive objective (2) , then we have

Original: [Equation]
Revised: [Equation]

Original: Now we are ready to proof Theorem 4.1.
Revised: Now we are ready to proof Theorem 4.1.

Original: Theorem 4.1.
Revised: Theorem 4.1.

Original: Consider the general topic model setting as Deﬁnition 2.1, suppose topic-word matrix
Revised: Consider the general topic model setting as Deﬁnition 2.1, suppose topic-word matrix

Original: A satisﬁes rank ( A ) =
Revised: A satisﬁes rank ( A ) =

Original: K and function f minimizes the contrastive objective (2) .
Revised: K and function f minimizes the contrastive objective (2) .

Original: Assume we randomly sampled m =
Revised: Assume we randomly sampled m =

Original: K t different landmark documents { l i } mi =1 and construct g ( x, { l i } mi =1 ) as (3) .
Revised: K t different landmark documents { l i } mi =1 and construct g ( x, { l i } mi =1 ) as (3) .

Original: Then, for any document x and its posterior w | ( x, A ) , the posterior mean of any polynomial P ( w ) with degree at most t is linear in g ( x, { l i } mi =1 ) .
Revised: Then, for any document x and its posterior w | ( x, A ) , the posterior mean of any polynomial P ( w ) with degree at most t is linear in g ( x, { l i } mi =1 ) .

Original: That is, there exists θ ∈ R m such that for all documents x
Revised: That is, there exists θ ∈ R m such that for all documents x

Original: [Equation]
Revised: [Equation]

Original: Proof.
Revised: Proof.

Original: Since f is the minimizer of contrastive objective (2), by Lemma B.1 we know
Revised: Since f is the minimizer of contrastive objective (2), by Lemma B.1 we know

Original: [Equation]
Revised: [Equation]

Original: Similar to the proof of Theorem 3.1, we will show that g ( x, x (cid:48) ) is linear in the topic posterior when x (cid:48) is ﬁxed.
Revised: Similar to the proof of Theorem 3.1, we will show that g ( x, x (cid:48) ) is linear in the topic posterior when x (cid:48) is ﬁxed.

Original: By Bayes’ rule, we have
Revised: By Bayes’ rule, we have

Original: [Equation]
Revised: [Equation]

Original: where we use P ( y = 0)
Revised: where we use P ( y = 0)

Original: =
Revised: =

Original: P ( y = 1) = 1 / 2 .
Revised: P ( y = 1) = 1 / 2 .

Original: Denote the t words in x (cid:48) as x (cid:48) 1 , . . .
Revised: Denote the t words in x (cid:48) as x (cid:48) 1 , . . .

Original: , x (cid:48) t and their corresponding topics as z 1 , . . .
Revised: , x (cid:48) t and their corresponding topics as z 1 , . . .

Original: , z t .
Revised: , z t .

Original: Then, by our way of generating x, x (cid:48) , we know
Revised: Then, by our way of generating x, x (cid:48) , we know

Original: [Equation]
Revised: [Equation]

Original: which implies
Revised: which implies

Original: [Equation]
Revised: [Equation]

Original: Note that
Revised: Note that

Original: [Equation]
Revised: [Equation]

Original: where A is the topic-word matrix.
Revised: where A is the topic-word matrix.

Original: Hence,
Revised: Hence,

Original: [Equation]
Revised: [Equation]

Original: Recall that the topic posterior tensor is W post
Revised: Recall that the topic posterior tensor is W post

Original: = E w [ w ⊗ t | x ] ∈ R K × ...
Revised: = E w [ w ⊗ t | x ] ∈ R K × ...

Original: × K , where each entry [ W post ] z 1 ,...,z k = E w [ w z 1 . . .
Revised: × K , where each entry [ W post ] z 1 ,...,z k = E w [ w z 1 . . .

Original: w z t | x ] .
Revised: w z t | x ] .

Original: Therefore,
Revised: Therefore,

Original: [Equation]
Revised: [Equation]

Original: where A =
Revised: where A =

Original: A ⊗ A ⊗ · · · ⊗ A ∈ R V t × K t ,
Revised: A ⊗ A ⊗ · · · ⊗ A ∈ R V t × K t ,

Original: A [ x (cid:48) ]
Revised: A [ x (cid:48) ]

Original: ∈ R K t is the row of A that correspond to x (cid:48) 1 , . . .
Revised: ∈ R K t is the row of A that correspond to x (cid:48) 1 , . . .

Original: , x (cid:48) t , and vec ( W post ) ∈
Revised: , x (cid:48) t , and vec ( W post ) ∈

Original: R K t is the vectorization of W post .
Revised: R K t is the vectorization of W post .

Original: Recall we have m randomly sampled different documents { l i } ni =1 .
Revised: Recall we have m randomly sampled different documents { l i } ni =1 .

Original: Denote D ∈ R m × m as a diagonal matrix such that D i,i = P ( l i ) , and ˜ A ∈ R m × K t as a submatrix of A such that each row of ˜
Revised: Denote D ∈ R m × m as a diagonal matrix such that D i,i = P ( l i ) , and ˜ A ∈ R m × K t as a submatrix of A such that each row of ˜

Original: A is A [ l i ] .
Revised: A is A [ l i ] .

Original: Since l i are randomly sampled, so D i,i > 0 .
Revised: Since l i are randomly sampled, so D i,i > 0 .

Original: Thus, we know
Revised: Thus, we know

Original: [Equation]
Revised: [Equation]

Original: , which implies vec ( W post )
Revised: , which implies vec ( W post )

Original: =
Revised: =

Original: ˜
Revised: ˜

Original: A † Dg ( x, { l i } ni =1 ) .
Revised: A † Dg ( x, { l i } ni =1 ) .

Original: Note that we need to show ˜
Revised: Note that we need to show ˜

Original: A † is well-deﬁned.
Revised: A † is well-deﬁned.

Original: Since A has full column rank, we know A also has full column rank.
Revised: Since A has full column rank, we know A also has full column rank.

Original: Thus, the submatrix ˜
Revised: Thus, the submatrix ˜

Original: A also has full column rank since m =
Revised: A also has full column rank since m =

Original: K t .
Revised: K t .

Original: This implies ˜
Revised: This implies ˜

Original: A † is well-deﬁned.
Revised: A † is well-deﬁned.

Original: Since P ( w ) is a polynomial of degree at most t , we know there exists β such that E w [ P ( w ) | x ] = β (cid:62) vec ( W post ) .
Revised: Since P ( w ) is a polynomial of degree at most t , we know there exists β such that E w [ P ( w ) | x ] = β (cid:62) vec ( W post ) .

Original: Thus, let θ = D
Revised: Thus, let θ = D

Original: ( ˜ A † ) (cid:62) β , we have
Revised: ( ˜ A † ) (cid:62) β , we have

Original: [Equation]
Revised: [Equation]

Original: B. P ROOF OF L EMMA B. 
Revised: B. P ROOF OF L EMMA B.

Original: Lemma B.1.
Revised: Lemma B.1.

Original: If f minimizes the contrastive objective (2) , then we have
Revised: If f minimizes the contrastive objective (2) , then we have

Original: [Equation]
Revised: [Equation]

Original: Proof.
Revised: Proof.

Original: Since f is the minimizer of contrastive objective (2) and
Revised: Since f is the minimizer of contrastive objective (2) and

Original: [Equation]
Revised: [Equation]

Original: it is easy to see that f is the Bayes optimal predictor P ( y = 1 | x, x (cid:48) ) .
Revised: it is easy to see that f is the Bayes optimal predictor P ( y = 1 | x, x (cid:48) ) .

Original: Therefore, we know
Revised: Therefore, we know

Original: [Equation]
Revised: [Equation]

Original: B. D ISCUSSIONS ON THE S ELF - REFERENCING
Revised: B. D ISCUSSIONS ON THE S ELF - REFERENCING

Original: 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 &
Revised: 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 &

Original: Recht, 2007)).
Revised: Recht, 2007)).

Original: Corollary B.2.
Revised: Corollary B.2.

Original: Denote { ( x i , ˜ y i ) } mi =1 as the downstream dataset, where ˜ y
Revised: Denote { ( x i , ˜ y i ) } mi =1 as the downstream dataset, where ˜ y

Original: i
Revised: i

Original: =
Revised: =

Original: E w [ P ( w ) | x i ] is the target for document x i .
Revised: E w [ P ( w ) | x i ] is the target for document x i .

Original: If we set landmarks { l i } mi =1 to be { x i } mi =1 , then solve the downstream task is equivalent to a kernel regression, that is
Revised: If we set landmarks { l i } mi =1 to be { x i } mi =1 , then solve the downstream task is equivalent to a kernel regression, that is

Original: [Equation]
Revised: [Equation]

Original: where ˜
Revised: where ˜

Original: y =
Revised: y =

Original: (˜ y 1 , . .
Revised: (˜ y 1 , . .

Original: ˜ y m ) (cid:62) and G ∈ R m × m is a kernel matrix such that G ij = g
Revised: ˜ y m ) (cid:62) and G ∈ R m × m is a kernel matrix such that G ij = g

Original: ( x i , x j ) .
Revised: ( x i , x j ) .

Original: As discussed in Corollary B.2, we show that one can use the same set of documents for both representation and downstream task.
Revised: As discussed in Corollary B.2, we show that one can use the same set of documents for both representation and downstream task.

Original: Recall that in Theorem 4.1 we need a large number of landmark documents to generate the representation and then use it for the downstream task (e.g., ﬁt a polynomial of topic posterior).
Revised: Recall that in Theorem 4.1 we need a large number of landmark documents to generate the representation and then use it for the downstream task (e.g., ﬁt a polynomial of topic posterior).

Original: However, one may not have enough additional documents to be used as landmarks.
Revised: However, one may not have enough additional documents to be used as landmarks.

Original: Therefore, the beneﬁt of the self-referencing is that we do not need additional landmark documents to generate the representation.
Revised: Therefore, the beneﬁt of the self-referencing is that we do not need additional landmark documents to generate the representation.

Original: The key observation in Corollary B.2 is that the learned representation g ( x, x (cid:48) ) can be viewed as a kernel.
Revised: The key observation in Corollary B.2 is that the learned representation g ( x, x (cid:48) ) can be viewed as a kernel.

Original: Suppose the documents for the downstream task are x 1 , . . .
Revised: Suppose the documents for the downstream task are x 1 , . . .

Original: , x m .
Revised: , x m .

Original: Construct a kernel matrix
Revised: Construct a kernel matrix

Original: G ∈ R m × m such that G ij = g
Revised: G ∈ R m × m such that G ij = g

Original: ( x i , x j ) .
Revised: ( x i , x j ) .

Original: Then, it is easy to see θ (cid:62) g
Revised: Then, it is easy to see θ (cid:62) g

Original: ( x i , { x j } mj =1 )
Revised: ( x i , { x j } mj =1 )

Original: = θ
Revised: = θ

Original: [Equation]
Revised: [Equation]

Original: where G i is the i -th column of G .
Revised: where G i is the i -th column of G .

Original: Thus, θ can be found by the following kernel regression:
Revised: Thus, θ can be found by the following kernel regression:

Original: [Equation]
Revised: [Equation]

Original: where ˜
Revised: where ˜

Original: y =
Revised: y =

Original: (˜ y 1 , . .
Revised: (˜ y 1 , . .

Original: ˜ y m ) (cid:62) and ˜ y
Revised: ˜ y m ) (cid:62) and ˜ y

Original: i
Revised: i

Original: =
Revised: =

Original: E w [ P ( w ) | x i ] .
Revised: E w [ P ( w ) | x i ] .

Original: C S YNTHETIC E XPERIMENT A DDITIONAL D ETAILS
Revised: C S YNTHETIC E XPERIMENT A DDITIONAL D ETAILS

Original: In this section, we include more details about our synthetic experiments in Section 5.
Revised: In this section, we include more details about our synthetic experiments in Section 5.

Original: In Section C.1, we describe the document generation process for PAM model.
Revised: In Section C.1, we describe the document generation process for PAM model.

Original: Additional results on topic posterior recovery loss and major topic recovery are included in Section C.2 and Section C.3.
Revised: Additional results on topic posterior recovery loss and major topic recovery are included in Section C.2 and Section C.3.

Original: In Section C. and Section C.5, we report the results about different training epochs and hyperparameters tuning.
Revised: In Section C. and Section C.5, we report the results about different training epochs and hyperparameters tuning.

Original: (cid:96) 1 condition number is reported in Section C.

Revised: (cid:96)condition number is reported in Section C.

Original: and we give some visualizations for topic correlations for the PAM model in Section C.7.
Revised: and we give some visualizations for topic correlations for the PAM model in Section C.7.

Original: Moreover, in Section C.8 we report the performance of SSL on recovering the topic posterior mean vector when t = 2 .
Revised: Moreover, in Section C.8 we report the performance of SSL on recovering the topic posterior mean vector when t = 2 .

Original: We run our experiments on 2080RTXTi/V100s.
Revised: We run our experiments on 2080RTXTi/V100s.

Revised: The conﬁdence interval reported in the paper is calculated as mean ± 1 .

Revised: 96 ∗ s testdocs / √ N testdocs , where mean is the mean of TV distance (or accuracy)

Revised: , s testdocs is the standard deviation of TV distance (or accuracy) and N testdocs is the number of test documents.

Revised: Details of generating topic-word matrix A For pure topic and LDA experiments, each column of

Original:  C. P ACHINKO A LLOCATION M ODEL T OPIC P RIOR G ENERATION
Revised: A (i.e., each topic’s word distribution) is generated from Dir ( α/K ) . For CTM and PAM, our goal is to construct K/ 4 groups of 4 topics such that within each group, topic 0 and topic 1 have similar word distribution, and topic 2 and topic 3 have different word distribution (and different from the word distribution of topic 0 and 1 as well). To achieve this, we ﬁrst sample K/ 4 word distribution vectors from Dir ( α/K ) , and for each word distribution vector, we permute its entries to get the word distribution of topic 0,1,2,3 within each group and intentionally align the large entries in topic 0 and 1 to ensure high similarity between them. C. P ACHINKO A LLOCATION M ODEL T OPIC P RIOR G ENERATION

Original: The topic proportion of the Pachinko Allocation Model (PAM) described in Section 5 is generated from the following process: for each document, we ﬁrst sample a “super-topic” proportion from a symmetric Dirichlet distribution Dir (1 /K s ) and a super-topic-to-topic proportion from a symmetric
Revised: The topic proportion of the Pachinko Allocation Model (PAM) described in Section 5 is generated from the following process: for each document, we ﬁrst sample a “super-topic” proportion from a symmetric Dirichlet distribution Dir (1 /K s ) and a super-topic-to-topic proportion from a symmetric

Original: Dirichlet distribution Dir (30) .
Revised: Dirichlet distribution Dir (30) .

Original: Then, we sample each word by ﬁrst sample a super-topic according to the super-topic proportion and then sample the actual topic from the super-topic-to-topic proportion.
Revised: Then, we sample each word by ﬁrst sample a super-topic according to the super-topic proportion and then sample the actual topic from the super-topic-to-topic proportion.

Original: In our experiment, we set K s = 10 .
Revised: In our experiment, we set K s = 10 .

Original: Figure C.1 offers a few examples of PAM topic proportion generated from the PAM model.[Figure]
Revised: Figure C.1 offers a few examples of PAM topic proportion generated from the PAM model.

Original:  M AJOR T OPIC R ECOVERY A DDITIONAL R ESULTS
Revised: C. M AJOR T OPIC R ECOVERY A DDITIONAL R ESULTS

Original: In Section 5.3, we measure the major topics recovery of CTM and PAM documents as the toptopic overlap rate.
Revised: In Section 5.3, we measure the major topics recovery of CTM and PAM documents as the toptopic overlap rate.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: C. S ELF - SUPERVISED L EARNING VERSUS A SSUMING S PECIFIC T OPIC P RIOR
Revised: C. S ELF - SUPERVISED L EARNING VERSUS A SSUMING S PECIFIC T OPIC P RIOR

Original: We present in Section 5.3 a comparison between the performance of our model and Markov Chain
Revised: We present in Section 5.3 a comparison between the performance of our model and Markov Chain

Original: Monte Carlo assuming a speciﬁc topic model for α = 1 , on the task of TV distance and major topic[Table] recovery.
Revised: Monte Carlo assuming a speciﬁc topic model for α = 1 , on the task of TV distance and major topic recovery.

Original: In this subsection, we supplement the experiment results for larger values of α along with α = 1 .
Revised: In this subsection, we supplement the experiment results for larger values of α along with α = 1 .

Original: 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.1 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.
Revised: In particular, we consider self-supervised learning with 6 million training documents and [Figure] [Table] 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.1 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.

Original: Additionally, when training with 120 thousand documents, we randomly shufﬂe 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.
Revised: Additionally, when training with 120 thousand documents, we randomly shufﬂe 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.

Original: We report our results in the form of 95% conﬁdence interval.
Revised: We report our results in the form of 95% conﬁdence interval.

Original: Table C.1 shows that our approach with 6 million training documents consistently outperforms misspeciﬁed topic priors by yielding a much lower TV distance from the true topic posterior.
Revised: Table C.1 shows that our approach with 6 million training documents consistently outperforms misspeciﬁed topic priors by yielding a much lower TV distance from the true topic posterior.

Original: Note that the entries with correct topic model are omitted, because in our experiment the Markov Chain
Revised: Note that the entries with correct topic model are omitted, because in our experiment the Markov Chain

Original: Monte Carlo recovered topic posterior assuming the correct topic model is exactly what we use as  our ground truth topic posterior.
Revised: Monte Carlo recovered topic posterior assuming the correct topic model is exactly what we use as [Table] our ground truth topic posterior.

Original: Meanwhile, our approach with 120 thousand training documents outperforms misspeciﬁed prior in almost every scenario, except when compared to LDA prior on pure topic documents.
Revised: Meanwhile, our approach with 120 thousand training documents outperforms misspeciﬁed prior in almost every scenario, except when compared to LDA prior on pure topic documents.

Original: This is is likely because the LDA topic prior is relatively concentrated on one speciﬁc topic and thus similar to the pure topic prior.
Revised: This is is likely because the LDA topic prior is relatively concentrated on one speciﬁc topic and thus similar to the pure topic prior.

Original: We also present a comparison between our self-supervised learning approach and posterior inference assuming a speciﬁc 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.
Revised: We also present a comparison between our self-supervised learning approach and posterior inference assuming a speciﬁc 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.

Original: 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 misspeciﬁed topic prior.[Table]
Revised: 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 misspeciﬁed topic prior.

Original: To further investigate the performance of self-supervised learning, we compare our approach with
Revised: To further investigate the performance of self-supervised learning, we compare our approach with

Original: Variational Inference assuming a speciﬁc topic prior on the task of recovering major topics in test document’s topic proportion.
Revised: Variational Inference assuming a speciﬁc topic prior on the task of recovering major topics in test document’s topic proportion.

Original: Table C.3 shows that, similar to our ﬁndings 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 misspeciﬁed topic prior in almost every test scenario.
Revised: Table C.3 shows that, similar to our ﬁndings 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 misspeciﬁed topic prior in almost every test scenario.

Original: C. T HE E FFECT OF T RAINING E POCHS
Revised: C. T HE E FFECT OF T RAINING E POCHS

Original: We measure how the number of training epochs may inﬂuence topic posterior recovery loss by varying the number of training epochs.
Revised: We measure how the number of training epochs may inﬂuence topic posterior recovery loss by varying the number of training epochs.

Original: The results we present in Section 5 are based on models trained for  200 epochs for all types of documents.
Revised: The results we present in Section 5 are based on models trained for [Table] 200 epochs for all types of documents.

Original: Here, we present the topic recovery loss, measured as the Total
Revised: Here, we present the topic recovery loss, measured as the Total

Original: 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.
Revised: 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.

Original: Table C.4 shows that topic posterior recovery loss steadily decreases as the number of epochs gets larger.
Revised: Table C.4 shows that topic posterior recovery loss steadily decreases as the number of epochs gets larger.

Original: 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.[Table] [Table] [Figure]
Revised: 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.

Original:  H YPERPARAMETER T UNING E XPERIMENTS
Revised: C. H YPERPARAMETER T UNING E XPERIMENTS

Original: As described in Section [Table] the two types of neural network architecture we use are fully-connected neural networks and attention-based neural networks (see Figure C.3).
Revised: 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).

Original: In this section, we present some of our TV results of different combinations of hyperparameters for both types of models.
Revised: In this section, we present some of our TV results of different combinations of hyperparameters for both types of models.

Original: 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.
Revised: 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.

Original: We only use one attention head in each multi-head attention layer.
Revised: We only use one attention head in each multi-head attention layer.

Original: During training, we use AMSGrad optimizer and an initial learning rate of 0 . 0001 or 0 . 0002 .
Revised: During training, we use AMSGrad optimizer and an initial learning rate of 0 . 0001 or 0 . 0002 .

Original: We reduce the learning rate by 50% whenever validation loss does not decrease in 10 epochs.
Revised: We reduce the learning rate by 50% whenever validation loss does not decrease in 10 epochs.

Original: The main hyperparameters in our fully-connected neural networks include the hidden dimension, the number of layers [Table] and whether we apply residual connections.
Revised: The main hyperparameters in our fully-connected neural networks include the hidden dimension, the number of layers, and whether we apply residual connections.

Original: Table C.5 shows the TV distance results when we vary these hyperparameters for the LDA scenario when α = 1 .
Revised: Table C.5 shows the TV distance results when we vary these hyperparameters for the LDA scenario when α = 1 .

Original: However, fully-connectedneural networks perform poorly in CTM and PAM case.
Revised: However, fully-connected

Original:  For instance, when α = 1 and hidden dimension = 4096, we ﬁnd that hyperparameters that work well in the LDA case no longer yield satisfactory results (Table C.6)
Revised: [Figure] For instance, when α = 1 and hidden dimension = 4096, we ﬁnd that hyperparameters that work well in the LDA case no longer yield satisfactory results (Table C.6).

Original: [Table]
Revised: [Table]

Original: TV [Table] [Table] 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.
Revised: 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.

Original: 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 ﬁnal true CTM topic posterior.
Revised: 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 ﬁnal true CTM topic posterior.

Original: We ﬁnd that the model performs the best when attention dimension is 768 or 1024, and when the attention dimension increases tothe model gives a higher TV (see Table C.8).
Revised: We ﬁnd that the model performs the best when attention dimension is 768 or 1024, and when the attention dimension increases tothe model gives a higher TV (see Table C.8).

Original: [Table]
Revised: [Table]

Revised: [Table]

Original: [Table]
Revised: [Table]

Original: C. T OPIC P OSTERIOR R ECOVERY L OSS AND C ONDITION N UMBER
Revised: C. T OPIC P OSTERIOR R ECOVERY L OSS AND C ONDITION N UMBER

Original: We have proved in Theorem 3.3 that the upper bound of topic posterior recovery loss depends on the (cid:96) 1 Condition Number κ ( A † ) .
Revised: We have proved in Theorem 3.3 that the upper bound of topic posterior recovery loss depends on the (cid:96) 1 Condition Number κ ( A † ) .

Original: In this section, we show that κ ( A † ) is small for every topic-word matrix
Revised: In this section, we show that κ ( A † ) is small for every topic-word matrix

Original: A in our experiments (Figure C.4).
Revised: A in our experiments (Figure C.4).

Original: Note that we use the same topic-word matrix A for pure-topic model and the LDA model.
Revised: Note that we use the same topic-word matrix A for pure-topic model and the LDA model.

Original: For CTM and PAM, we use the same topic-word matrix up to reordering of the topics that corresponds with pairwise topic correlations.
Revised: For CTM and PAM, we use the same topic-word matrix up to reordering of the topics that corresponds with pairwise topic correlations.

Original: Therefore, pure-topic model and LDA model share the same κ ( A † ) , and CTM and PAM share the same κ ( A † ) 
Revised: Therefore, pure-topic model and LDA model share the same κ ( A † ) , and CTM and PAM share the same κ ( A † ) .

Original: C. V ISUALIZING T OPIC C ORRELATIONS FOR PAM M ODEL

Original: 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.
Revised: 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.

Original: From this ﬁgure (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.
Revised: From this ﬁgure (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.

Original: C. T OPIC P OSTERIOR R ECOVERY FOR t =
Revised: C. T OPIC P OSTERIOR R ECOVERY FOR t =

Original: We have shown by experiment that self-supervised learning can accurately recover the topic posterior when t = 1 , that is, when the neural network model is trained to predict one missing word from a document.
Revised: We have shown by experiment that self-supervised learning can accurately recover the topic posterior when t = 1 , that is, when the neural network model is trained to predict one missing word from a document.

Original: In this subsection, we run experiments for t = 2 case in Theorem 3.3, which means using the self-supervised objective to train the neural network model to predict 2 missing word in a document.
Revised: In this subsection, we run experiments for t = 2 case in Theorem 3.3, which means using the self-supervised objective to train the neural network model to predict 2 missing word in a document.

Original: For the sake of reducing computational complexity, we change the setting from
Revised: For the sake of reducing computational complexity, we change the setting from

Original: V = 5000 , K = 20 , K s = 10 to V = 500 , K = 8 , and K s = 4 (where K s is a parameter in the
Revised: V = 5000 , K = 20 , K s = 10 to V = 500 , K = 8 , and K s = 4 (where K s is a parameter in the

Original: PAM model) while keeping document length, resample rate, training epoch, and the number of test documents the same as in Section 5.
Revised: PAM model) while keeping document length, resample rate, training epoch, and the number of test documents the same as in Section 5.

Original: Once a neural network f is trained, from the proof of Theorem 3.
Revised: Once a neural network f is trained, from the proof of Theorem 3.

Original: in Section A.2 we know that the training loss minimizer f satisﬁes f ( x ) =
Revised: in Section A.2 we know that the training loss minimizer f satisﬁes f ( x ) =

Original: ( A ⊗ A ) vec ( W post ) when t = 2 , so we obtain vec( W post ) by vec( W post ) =
Revised: ( A ⊗ A ) vec ( W post ) when t = 2 , so we obtain vec( W post ) by vec( W post ) =

Original: ( A ⊗ A ) † f ( x ) where we expect W post to be a symmetric matrix.
Revised: ( A ⊗ A ) † f ( x ) where we expect W post to be a symmetric matrix.

Original: But we ﬁnd that in our experiment W post is nearly symmetric but not perfectly symmetric, and to amend this, we use 1 2 ( W post + W Tpost ) as the SSL-recovered posterior tensor.
Revised: But we ﬁnd that in our experiment W post is nearly symmetric but not perfectly symmetric, and to amend this, we use 12 ( W post + W Tpost ) as the SSL-recovered posterior tensor.

Original: [Table]
Revised: [Table]

Original: We observe that the attention-based neural network architecture with 384 hidden dimension and varying number of attention blocks (shown in Table C.9) is performs better on learning the distribution of 2 missing words than the fully-connected neural network for all four types of documents, so all of the following results are based on the output of trained attention-based neural network.
Revised: We observe that the attention-based neural network architecture with 384 hidden dimension and varying number of attention blocks (shown in Table C.9) is performs better on learning the distribution of 2 missing words than the fully-connected neural network for all four types of documents, so all of the following results are based on the output of trained attention-based neural network.

Original: We use the posterior mean from MCMC assuming the correct prior as the ground-truth topic posterior, and use the TV distance to measure how far the SSL-recovered topic posterior deviates from it, as shown in
Revised: We use the posterior mean from MCMC assuming the correct prior as the ground-truth topic posterior, and use the TV distance to measure how far the SSL-recovered topic posterior deviates from it, as shown in

Original: Table C.10.
Revised: Table C.10.

Original: In most scenarios, the TV distance is less than 0.15, and TV distance gets slightly larger as α increases, i.e., the word distribution of each topic gets less distinguishable.
Revised: In most scenarios, the TV distance is less than 0.15, and TV distance gets slightly larger as α increases, i [Table] e., the word distribution of each topic gets less distinguishable.

Original: [Table]
Revised: [Table]

Original: We also measure how the largest entries in the SSL-recovered topic posterior overlap with the largest entries in the topic prior.
Revised: We also measure how the largest entries in the SSL-recovered topic posterior overlap with the largest entries in the topic prior.

Original: Note that for each document, since the underlying topic of each word is  drawn iid, the joint topic prior distribution for the 2 missing words is the Kronecker product of the document’s topic prior distribution with itself.
Revised: Note that for each document, since the underlying topic of each word is [Table] drawn iid, the joint topic prior distribution for the 2 missing words is the Kronecker product of the document’s topic prior distribution with itself.

Original: For the overlap rate, in the t = 1 case we measured the overlap of the top 1 topic for pure topic and LDA documents and the overlap of top 2 topics for CTM and PAM documents, because most pure topic documents and LDA documents are generated by one dominant topic whereas most CTM and PAM documents are generated by two correlated dominant topics.
Revised: For the overlap rate, in the t = 1 case we measured the overlap of the top 1 topic for pure topic and LDA documents and the overlap of top 2 topics for CTM and PAM documents, because most pure topic documents and LDA documents are generated by one dominant topic whereas most CTM and PAM documents are generated by two correlated dominant topics.

Original: Similarly, in the t = 2 case we measure the overlap of the top 1 entry for pure topic and LDA documents and the overlap of top 2 × 2 = 4 entries for CTM and PAM documents.
Revised: Similarly, in the t = 2 case we measure the overlap of the top 1 entry for pure topic and LDA documents and the overlap of top 2 × 2 = 4 entries for CTM and PAM documents.

Original: We will refer to this overlap rate as the “major topic pair recovery rate".
Revised: We will refer to this overlap rate as the “major topic pair recovery rate".

Original: The results are shown in Table C.11.
Revised: The results are shown in Table C.11.

Original: Overall,[Table] the overlap rate is very high, which indicates that for t = 2 our self-supervised learning approach can capture the most dominant topic(s) for the 2 missing words.
Revised: Overall, the overlap rate is very high, which indicates that for t = 2 our self-supervised learning approach can capture the most dominant topic(s) for the 2 missing words.

Original: D R EAL - DATA E XPERIMENT A DDITIONAL D ETAILS 
Revised: C. D ETAILS OF MCMC NUTS SAMPLING METHOD

Revised: In our experiments, we use PyMC3 package Salvatier et al.

Revised: (2016) to perform NUTS-based MCMC posterior sampling.

Revised: Speciﬁcally, we use 3000 tuning iterations to tune the sampler’s step size so that we approximate an acceptance rate of 0.9.

Revised: After tuning, we draw 2000 samples from the posterior distribution and use the mean of these 2000 samples as our best estimate for the topic posterior mean.

Revised: To further check the performance of NUTS, we look into the LDA documents.

Revised: As shown below, we observe that MCMC posterior assuming LDA model (bottom row) is also very similar to the ground truth w (top row) used to generate the documents.

Revised: D R EAL - DATA E XPERIMENT A DDITIONAL D ETAILS

Original: In this section, we give more details about our real data experiments in Section 6.
Revised: In this section, we give more details about our real data experiments in Section 6.

Original: In Section D.1, we describe the data processing.
Revised: In Section D.1, we describe the data processing.

Original: We give a more detailed description of baselines and how we extracted our representations in Section D.2.
Revised: We give a more detailed description of baselines and how we extracted our representations in Section D.2.

Original: In Section D.3, we report more results using different model architectures.
Revised: In Section D.3, we report more results using different model architectures.

Original: Moreover, in Section D.4, we offer details about experiments on other real-world datasets and report the results.
Revised: Moreover, in Section D.4, we offer details about experiments on other real-world datasets and report the results.

Original: D. D ATA P ROCESSING
Revised: D. D ATA P ROCESSING

Original: Here we detailed our usage of the AG news dataset by Zhang et al.
Revised: Here we detailed our usage of the AG news dataset by Zhang et al.

Original: Each category has 30, samples in the training set and 19,000 samples in the testing set.
Revised: Each category has 30, samples in the training set and 19,000 samples in the testing set.

Original: We ﬁrst 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.
Revised: We ﬁrst 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.

Original: (2021b).
Revised: (2021b).

Original: To split the data set into unsupervised dataset and supervised dataset, we selected a random sample of 1000 documents as labeled supervised dataset for each category of documents, while the remaining 116,000 documents fall into unsupervised dataset for representation learning.
Revised: To split the data set into unsupervised dataset and supervised dataset, we selected a random sample of 1000 documents as labeled supervised dataset for each category of documents, while the remaining 116,000 documents fall into unsupervised dataset for representation learning.

Original: D. E XTRACTING R EPRESENTATIONS
Revised: D. E XTRACTING R EPRESENTATIONS

Original: To give some details about two baseline representations we used, we described in details as follows:
Revised: To give some details about two baseline representations we used, we described in details as follows:

Original: • Bag of Words (BOW):
Revised: • Bag of Words (BOW):

Original: 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.
Revised: 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.

Original: • Word2vec: To generate Word2vec representation ﬁtted the Skip-gram word embedding model on unsupervised dataset Mikolov et al.
Revised: • Word2vec: To generate Word2vec representation ﬁtted the Skip-gram word embedding model on unsupervised dataset Mikolov et al.

Original: The implementation is done through the Gensim library Rehurek & Sojka (2011), where we used an embedding dimension of(after tuning, see Table D.1) and window size of 5.
Revised: The implementation is done through the Gensim library Rehurek & Sojka (2011), where we used an embedding dimension of(after tuning, see Table D.1) and window size of 5.

Original: Representation is taken as the average of all trained word embeddings in a single document.
Revised: Representation is taken as the average of all trained word embeddings in a single document.

Original: LDA: We construct a representation derived from LDA
Revised: LDA: We construct a representation derived from LDA

Original: by ﬁrst ﬁtting Latent Dirichlet
Revised: by ﬁrst ﬁtting Latent Dirichlet

Original: Allocation model Blei et al.
Revised: Allocation model Blei et al.

Original: (2003) on the unsupervised dataset and then use the posterior distribution over topics given a document as representation.
Revised: (2003) on the unsupervised dataset and then use the posterior distribution over topics given a document as representation.

Original: We implemented LDA using scikit-learn Pedregosa et al.
Revised: We implemented LDA using scikit-learn Pedregosa et al.

Original: (2011) with default parameters except for number of topics, which we use as representation size.
Revised: (2011) with default parameters except for number of topics, which we use as representation size.

Original: We use ﬁnal dimension of 100 (after tuning, see
Revised: We use ﬁnal dimension of 100 (after tuning, see

Original: Figure D.1).
Revised: Figure D.1).

Original: 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)
Revised: 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)

Original: 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).
Revised: 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).

Original: The dimension of representation with Word2vec is 300 while the softmax + second2 last layer has a dimension of 4096.
Revised: The dimension of representation with Word2vec is 300 while the softmax + second2 last layer has a dimension of 4096.

Original: 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.
Revised: 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.

Original: [Table]
Revised: [Table]

Original: We included a comparison of the two approaches with dimension reduction in Table D.
Revised: We included a comparison of the two approaches with dimension reduction in Table D.

Original: [Table]
Revised: [Table]

Original: from which we observe that a Softmax on the second-to-last layer output performed better across layers of 3, and 5.
Revised: from which we observe that a Softmax on the second-to-last layer output performed better across layers of 3, and 5.

Original: We reported the result using the softmax+second2last layer in Figure
Revised: We reported the result using the softmax+second2last layer in Figure

Original: [Table]
Revised: [Table]

Original: M ODEL A RCHITECTURE
Revised: M ODEL A RCHITECTURE

Original: We include here more details on our model architecture used in real data experiments.
Revised: We include here more details on our model architecture used in real data experiments.

Original: We used residual model, and tested two main hyperparameters: number of residual blocks and hidden dimension size.
Revised: We used residual model, and tested two main hyperparameters: number of residual blocks and hidden dimension size.

Original: We ﬁxed the rest of hyperparameters: we used a 5000 embedding dimension and trained forepochs, with 0.0002 learning rate, weight decay of 0.01 and resample rate of 2.
Revised: We ﬁxed the rest of hyperparameters: we used a 5000 embedding dimension and trained forepochs, with 0.0002 learning rate, weight decay of 0.01 and resample rate of 2.

Original: We sampled 4 words in labels for variance reduction purpose.
Revised: We sampled 4 words in labels for variance reduction purpose.

Original: Varying Depth and Width
Revised: Varying Depth and Width

Original: We investigated both the effect of model depth and width on the performance of RBL representation.
Revised: We investigated both the effect of model depth and width on the performance of RBL representation.

Original: 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.
Revised: 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.

Original: We varied the depth of the network by choosing to place 3, 4 and 5 of such block.
Revised: We varied the depth of the network by choosing to place 3, 4 and 5 of such block.

Original: As shown in Figure D.2, it appears that using wider models in the unsupervised phase leads to better performance when training a linear classiﬁer on the learned representations.
Revised: As shown in Figure D.2, it appears that using wider models in the unsupervised phase leads to better performance when training a linear classiﬁer on the learned representations.

Original: Number of residual block does not seem to have a clear relationship with test accuracy from limited experiments we ran.
Revised: Number of residual block does not seem to have a clear relationship with test accuracy from limited experiments we ran.

Original: 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.
Revised: 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.

Original: D. M ORE R EAL - DATA E XPERIMENTS
Revised: D. M ORE R EAL - DATA E XPERIMENTS

Original: In addition to the experiments on the AG news dataset, we also run experiments on the IMDB movie review sentiment classiﬁcation dataset by Maas et al.
Revised: In addition to the experiments on the AG news dataset, we also run experiments on the IMDB movie review sentiment classiﬁcation dataset by Maas et al.

Original: (2011) and the Small DBpedia ontology dataset, which is a subset of the DBpedia ontology by Zhang et al. (2015), following the preprocessing by
Revised: (2011) and the Small DBpedia ontology dataset, which is a subset of the DBpedia ontology by Zhang et al. (2015), following the preprocessing by

Original: Tosh et al.
Revised: Tosh et al.

Original: (2021b).
Revised: (2021b).

Original: In particular, the IMDB movie review dataset contains a unlabeled set of 50k movie reviews, a training set of 12.5k positive movie reviews and 12.5k negative movie reviews, and a test set of 12.5k positive movie reviews and 12.5k negative movie reviews.
Revised: In particular, the IMDB movie review dataset contains a unlabeled set of 50k movie reviews, a training set of 12.5k positive movie reviews and 12.5k negative movie reviews, and a test set of 12.5k positive movie reviews and 12.5k negative movie reviews.

Original: We randomly selected 1k documents from each class in the training set as our training document and selected 3.5k documents in the test set as our test documents, and used the remaining 91k documents as our unsupervised dataset for training our SSL model to learn document representation.
Revised: We randomly selected 1k documents from each class in the training set as our training document and selected 3.5k documents in the test set as our test documents, and used the remaining 91k documents as our unsupervised dataset for training our SSL model to learn document representation.

Original: After ﬁltering,
Revised: After ﬁltering,

Original: [Figure]
Revised: [Figure]

Original: In all these runs, we use Softmax on second-to-last layer to obtain our representations
Revised: In all these runs, we use Softmax on second-to-last layer to obtain our representations

Original: [Figure]
Revised: [Figure]

Original: the vocabulary size is about 35k.
Revised: the vocabulary size is about 35k.

Original: The Small DBpedia dataset consists of 4 non-overlapping classes, which are company, artist, athlete, and ofﬁce holder from the original DBpedia ontology dataset, and each class consists 40k training documents and 5k test documents.
Revised: The Small DBpedia dataset consists of 4 non-overlapping classes, which are company, artist, athlete, and ofﬁce holder from the original DBpedia ontology dataset, and each class consists 40k training documents and 5k test documents.

Original: We randomly chose 1k training documents each class as labeled training documents and used the rest of training documents as the unsupervised dataset, and used the original test documents as our test dataset.
Revised: We randomly chose 1k training documents each class as labeled training documents and used the rest of training documents as the unsupervised dataset, and used the original test documents as our test dataset.

Original: After we ﬁltered rare words in the original DBpedia dataset, the vocabulary size is about 10k.
Revised: After we ﬁltered rare words in the original DBpedia dataset, the vocabulary size is about 10k.

Original: In our experiments, the setup of the baselines follows Section D.2.
Revised: In our experiments, the setup of the baselines follows Section D.2.

Original: We plot the test accuracy of our method against the baselines in Figure D.3.We did not include NCE in representation comparison since limited tuning does not achieve comparable results as in Tosh et al.
Revised: We plot the test accuracy of our method against the baselines in Figure D.3.We did not include NCE in representation comparison since limited tuning does not achieve comparable results as in Tosh et al.

Original: (2021b).Parameter tuning result for LDA and Word2vec is documented in Figure D.4 and Table D.3 respectively.
Revised: (2021b).Parameter tuning result for LDA and Word2vec is documented in Figure D.4 and Table D.3 respectively.

Original: In the left panel of Figure D.3, we compare the performance of semi-supervised learning on IMDB dataset of our representation RBL as well as baselines.
Revised: In the left panel of Figure D.3, we compare the performance of semi-supervised learning on IMDB dataset of our representation RBL as well as baselines.

Original: We point out that in general, RBL has similar performance to Word2vec, outperforming LDA and BOW.
Revised: We point out that in general, RBL has similar performance to Word2vec, outperforming LDA and BOW.

Original: However,when labeled training data is less abundant, Word2vec
Revised: However,when labeled training data is less abundant, Word2vec

Original: and LDA have slightly stronger performance.
Revised: and LDA have slightly stronger performance.

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: In the right panel of Figure D. we perform a similar representation comparison on Small DBpedia dataset.
Revised: In the right panel of Figure D. we perform a similar representation comparison on Small DBpedia dataset.

Original: While RBL outperforms all other baselines when there are more than around 1000 training examples, it performs worse than Word2vec and LDA when labeled training data is less abundant.
Revised: While RBL outperforms all other baselines when there are more than around 1000 training examples, it performs worse than Word2vec and LDA when labeled training data is less abundant.