Year: 2021

Abstract:
The discovery of large-scale discrete latent structures is crucial for understanding the fundamental generative processes of language. In this work, we use structured latent variables to study the representation space of contextualized embeddings and gain insight into the hidden topology of pretrained language models. However, existing methods are severely limited by issues of scalability and efficiency as working with large combinatorial spaces requires expensive memory consumption. We address this challenge by proposing a Randomized Dynamic Programming (RDP) algorithm for the approximate inference of structured models with DP-style exact computation (e.g., Forward-Backward). Our technique samples a subset of DP paths reducing memory complexity to as small as one percent. We use RDP to analyze the representation space of pretrained language models, discovering a large-scale latent network in a fully unsupervised way. The induced latent states not only serve as anchors marking the topology of the space (neighbors and connectivity), but also reveal linguistic properties related to syntax, morphology, and semantics. We also show that traversing this latent network yields unsupervised paraphrase generation.

heading: 1 INTRODUCTION
text: The discovery of large-scale discrete latent structures is crucial for understanding the fundamental generative processes of language, and has been shown useful to various NLP tasks ranging from data-to-text generation (Li & Rush, 2020), summarization (Angelidis et al., 2021), syntactic parsing (Kim et al., 2019), and knowledge graph reasoning (Qu et al., 2020). In this work, we use latent structures to analyze geometric properties of representation space of pretrained language models (PLMs). Despite the large volume of recent work analyzing PLMs and proposing various improvements (Rogers et al., 2020), little is known about the topological structure of their representation manifold. Since such structure cannot be easily observed, it is only natural to resort to latent variables. Yet scaling discrete combinatorial structures is extremely difficult with multiple modeling and computational challenges (Wainwright & Jordan, 2008).
In this work, we address the computational challenges arising from working with combinatorial structures. We consider linear-chain CRFs, a popular structured model family (Ma & Hovy, 2016; Sutton & McCallum, 2006) that uses dynamic programming for exact inference. Specifically, we focus on the forward algorithm (Rabiner, 1989), which is widely used to compute the partition function. Space complexity for this algorithm is O(TN2) where N is the number of latent states and T the length of the sequence. It is precisely the N2 term that becomes problematic when we construct the adjacent gradient graph with automatic differentiation. DP-based inference algorithms are not optimized for modern computational devices like GPUs and typically work under small-data regimes, with N in the range [10, 100] (Ma & Hovy, 2016; Wiseman et al., 2018). With larger N , inference becomes intractable since gradients do not easily fit into GPU memory (Sun et al., 2019).
Our algorithmic contribution is a randomization technique for dynamic programming which allows us to scale N to thousands (possibly more) latent states. Specifically, to approximate the partition function, instead of summing over all possible combinations of latent states, we only sum over paths with most probable states, and sample a subset of less likely paths to correct the bias according
to a reasonable proposal. Since we only calculate the sampled path, memory consumption can be reduced to a small controllable budget which is scale invariant. With a larger memory budget, our method becomes more accurate, and our estimation error smaller. We thus recast the memory complexity challenge into a tradeoff between memory budget, proposal accuracy, and estimation error. When applied to linear-chain CRFs, we show that RDP scales the model by two orders of magnitude with memory complexity as small as one percent of the full DP. Beyond linear-chains, RDP is applicable to any structured model with DP-style exact inference such as trees (Kim et al., 2019) and semi-Markov models (Li & Rush, 2020), and could also be extended to more general message passing algorithms (Wainwright & Jordan, 2008).
Our analytical contribution is a geometric study of the representation manifold of PLMs, using the proposed RDP algorithm. We hypothesize that there exist latent anchor embeddings (or landmarks) that describe the manifold topology. We also expect these anchor states to be informative enough to generate sentences, and their connections to be linguistically meaningful. We induce latent structures using a VAE with an inference model parameterized by a scaled CRF where state-word relations are modeled by the emission potential and state-state transitions are modeled by the transition matrix. The connections of words and states together form a latent network. We use the vector product between contextualized embeddings and state embeddings to parameterize the CRF potentials, bringing together the geometry of the representation space with graphical model inference. We further show that it is possible to generate paraphrases by traversing the induced network.
Our approach is fully unsupervised and the discovered latent network is intrinsic to the representation manifold, rather than imposed by external supervision, eschewing the criticism of much previous work on supervised probes (Hewitt & Liang, 2019; Chen et al., 2021). In experiments, we first verify the basic properties of RDP (bias-variance) and show its effectiveness for training latent variable models. We then visualize the discovered network based on BERT Devlin et al. (2019), demonstrating how states encode information pertaining to syntax, morphology, and semantics. Finally, we perform unsupervised paraphrase generation by latent network traversal.

heading: 2 RANDOMIZED DYNAMIC PROGRAMMING
text: Preliminaries: Speeding Summation by Randomization To motivate our randomized DP, we start with a simple setting, namely estimating the sum of a sorted list. Given a sorted list of positive numbers a, naive summation S = a1 + ...,+aN requires N − 1 addition operations, which is expensive whenN is large. Suppose we wish to reduce the number of addition operations toK1 << N , and we already know that the list is long-tailed (similar to how words in language follow a Zipfian distribution such that there are few very high-frequency words that account for most of the tokens in text and many low-frequency words). Then, we only need to sum over the top K1 values to get an efficient estimate:
Ŝ1 = a1 + ...+ aK1 where {ai}Ni=1 sorted, large to small (1)
Clearly, Ŝ1 underestimates S. When the summands are “dense”, i.e., not very different from each other, the bias is large because the top K1 terms do not contribute much to the sum (Fig. 1A). To correct this bias, we add samples aδ1 , ..., aδK2 from the remaining summands whose indices δi are sampled from proposal δi ∼ q = [qK1+1, ..., qN ]:
Ŝ2 = a1 + ...+ aK1 + 1
K2 (
1
qδ1 aδ1 + ...+
1
qδK2 aδK2 ) δi ∈ {K1 + 1, ..., N} (2)
where K1 + K2 = K. Note that this is an unbiased estimator as E[Ŝ2] = S, irrespective of how we choose q. Without any knowledge about a, the simplest proposal would be uniform, no matter what variance it induces. The more qi correlates with ai, the less variance Ŝ2 has. The oracle qi is proportional to ai, under which Ŝ2 becomes exact Ŝ2 ≡ S as qδi = aδi/(aK+1 + ...+ aN ) for all i. So, the strategy is to exploit our knowledge about a to construct a correlated proposal q. Given this estimator, we can also adjust the computation budget in order to reduce variance. When the distribution is long-tailed, we may increase K1 as an instance of Rao-Blackwellization (Liu et al., 2019). When the distribution is not long-tailed (enough), and top K1 summation underestimates significantly, we may increase K2 to reduce variance, provided we have a fairly accurate q, as an instance of importance sampling. This procedure is also discussed in Kool et al. (2020) for
A. Sampled Sum of a sorted list B. Sampled Forward recursion C. Model Architecture
Long-tail case ãt−1 ãt
BERT
x0 x1 x2 x3
r0 r1 r2 r3
z0 z1 z2 z3
Decoder
x0 x1 x2 x3
TopK summand Sampled summand Gap to oracle Dropped summand TopK state Sampled state
Representation Space
Dense case
Encoder
Latent structure in representation space
RDP Decoder Backprop. 1 2 3 Log Z
Paraphrasing4
D. Experiment Protocol
Figure 1: (A): Sampled summation of an array; in the dense case the proposal is important for variance reduction, while in the long-tailed case, topK summands are important; (B): core recursion step of the Randomized Forward algorithm. We get topK and sample from the proposal (black and grey bars); Errors stem from the difference (green bars) between the oracle proposal ã and constructed proposal q̃; (C): Inferring latent states within the BERT representation space. We parametrize the CRF factors with vector products; the relations between states and contextualized embeddings together form a latent network (Fig. 3 and 4); (D): Experimental protocol; we first study the basic properties of RDP (steps 1, 2) and then integrate RDP into a LVM for inferring the structure of the representation space (steps 3, 4). Best viewed in color.
gradient estimation. In fact, it is the underlying basis of many Monte Carlo estimators in various settings (Mohamed et al., 2020).
The Sampled Forward Algorithm Now we will show how estimator Ŝ2 can be used to scale summation in DP. Consider a linear chain CRF which defines a discrete state sequence z = [z1, ..., zT ], zt ∈ {1, ..., N} over an input sentence x = [x1, ..., xT ]. Later we will use this CRF to construct an inference model to discover latent network structures within contextualized representations. We are interested in the partition function Z which is commonly computed with the Forward algorithm, a dynamic programming algorithm that sums over the potentials of all possible state sequences. The core recursion steps are:
αt+1(i) = N∑ j=1 ãt+1(i, j) = N∑ j=1 αt(j)Φ(j, i)φ(xt+1, i) Z = N∑ j=1 αT (j) (3)
where αt(i) is the sum of all possible sequences up to step t and at state i, Φ(·, ·) is an N × N transition matrix, and φ(xt, i) is the emission potential that models how word xt generates state i. We assume all potentials are positive for simplicity. When implemented on GPUs, space complexity is O(TN2) (see number of edges in the DP graph in Figure 1B) and it is the squared term N2 that causes memory overflows under automatic differentiation (see Appendix B for engineering details).
Our key insight is to recursively use the memory-efficient randomization of Eq. 2 to estimate Eq. 3 at every step. Given a proposal q̃t for each step t that correlates with summands ãt (we discuss how to construct q̃t in the next section), we obtain its top K1 index and sample K2 from the rest:
[σt,1, ..., σt,K1 , ..., σt,N ] = arg sorti{q̃t(i)}Ni=1 (4) [δt,1, ..., δt,K2 ] ∼ Categorical{q̃t(σt,K1+1), ..., q̃t(σt,N )} (5)
where q̃t(·) are normalized to construct the categorical. Compared to Eq. 3, the key recursion of our Sampled Forward uses the top K1 index σt and sampled K2 index δt to substitute the full index:
α̂t+1(i) = K1∑ j=1 α̂t(σt,j)Φ(σt,j , i)φ(xt+1, i) + 1 K2 K2∑ j=1 Z̃t q̃t(δt,j) α̂t(δt,j)Φ(δt,j , i)φ(xt+1, i) (6)
Z̃t = N∑ j=K1+1 q̃t(σt,j) Ẑ = K1∑ j=1 α̂T (σT,j) + 1 K2 K2∑ j=1 Z̃T q̃T (δT,j) α̂T (δT,j) (7)
where the oracle proposal q∗t is proportional to the actual summand ãt (Eq. 3, a little bit algebra will show this is actually the backward sampling probability p(zt = i|zt+1 = j)) , which is only accessible with the full Forward. So, we use the proposal weight q̃t (Eq. 4) to move the computation outside the DP. In Fig. 1B, the top K1 summed terms correspond to black nodes. The proposal q̃t
corresponds to black and grey bars, and its distance from the oracle proposal ãt (which is the major source of variance) is highlighted in green. Sampled indices are shown as blue nodes. Essentially, our Sampled Forward algorithm restricts the DP computation from the full graph to subgraphs with top and sampled edges, reducing complexity to O(TK2) where K = K1 + K2. By varying K, memory complexity becomes a tradeoff between memory budget and estimation error. By induction, we can show that Ẑ (Eq. 7) is an unbiased estimator of Z since ∀t,E[α̂t] = αt. When implemented in log space, the expected log Ẑ is a lower bound of the exact logZ due to Jensen’s inequality, and the variance is (trivially) reduced by log(·). See Appendix for details on implementation (Section C), theoretical analysis (Section A), and extensions to general sum-product structures (Section D).

heading: 3 LATENT NETWORK TOPOLOGY IN PRETRAINED LANGUAGE MODELS
text: Latent States within Representation Space We now use the above technique to uncover hidden geometric structures in contextualized representations. In experiments we work with BERT (Devlin et al., 2019) and GPT2 (Radford et al., 2019), however, our method can be easily applied to other pretrained language models. Given sentence x = [x1, ..., xT ], we denote its contextualized representations as [r1, ..., rT ] = PLM(x). Representations r for all sentences lie in one manifold M, namely, the representation space of the language model. We hypothesize there exists a set of latent states s1, ..., sM that function as anchors and outline the space topology. We emphasize that all parameters of the PLM are fixed (i.e., no fine-tuning takes place), so all learned states are intrinsic toM. We focus on two topological relations: (a) state-word relations, which represent how word embeddings may be summarized by their states and how states can be explained by their corresponding words; and (b) state-state relations, which capture how states interact with each other and how their transitions denote meaningful word combinations. Taken together, these two relations form a latent network withinM (visualized in Fig. 3 and 4). We adopt a minimal parametrization of the inference network so as to respect the intrinsic structure of the representation manifold without imposing strong assumptions (e.g., via regularization). Specifically, for state-word relations, we associate each word embedding rt with a latent state indexed by zt ∈ {1, ..., N} (the corresponding embedding of zt is szt ). For state-state relations, we assume a transition weight Φ(i, j). Together we have a linear-chain CRF:
log φ(xt, zt) = r ᵀ t szt log Φ(zt−1, zt) = s ᵀ zt−1szt (8)
where the dot product follows the common practice of fine-tuning contextualized representations. We use log space for numerical stability. The probability of a state sequence given a sentence is:
qψ(z|x) = T∏ t=1 Φ(zt−1, zt)φ(xt, zt)/Z (9)
Here, the only learnable parameters are state embeddings: ψ = [s1, ..., sN ] as we try to be faithful to the representation manifold. Note how this parametrization reconciles space geometry with graphical models. As N is large, we estimate Z with the proposed Sampled Forward (Eq. 7).
Constructing the Proposal We now return to proposal q̃t (Eq. 4) which we construct based on a common observation that linguistic phenomena are long-tailed:
q̃t(i) ∝ Φ(i)φ(xt, i) Φ(i) = ||si||1 (10) where φ(xt, i) states that only a few states are likely to generate observation xt, which is often the case in NLP (e.g., there are only a few possible POS tags for each word); and Φ(i) models the prior probability of state i. This choice stems from the empirical observation that larger L1 norm correlates with larger dot product, and is thus more likely to be inferred. Essentially, our proposal combines local emissions φ and global prior Φ to approximate the ãt variables (Eq. 3) and bypass their expensive computation.
Inference and Learning We use amortized variational inference to learn s. We simply reuse the architecture from previous work Fu et al. (2020); Li & Rush (2020) and build a generative model:
pθ(x, z) = ∏ t p(xt|z1:t, x1:t−1) · p(zt|z1:t−1, x1:t−1) ht = Dec([szt−1 ;xt−1],ht−1) (11)
p(xt|z1:t, x1:t−1) = softmax(FF(ht)) p(zt|z1:t−1, x1:t−1) = softmax(FF([szt ;ht])) (12)
where θ denotes the decoder parameters, Dec(·) denotes the decoder (we use an LSTM), ht denotes decoder states, and FF(·) denotes a feed-forward network. This autoregressive formulation essentially encourages states to be “generative”, i.e., to generate sentences and themselves. We will show in experiments how this formulation lends itself to paraphrasing. We use qψ directly from Eq. 9 as our variational posterior, and optimize the following β-ELBO objective:
LELBO = Eqψ(z|x)[log pθ(x, z)]− βH(qψ(z|x)) (13)
where the β parameter modulates the topology of the latent structure and prevents posterior collapse. We follow Fu et al. (2020) and use their Gumbel reparameterization to optimize qψ , which is more stable than the REINFORCE gradient estimator (Li & Rush, 2020).
When integrating RDP with the Gumbel reparameterization, we noticed that the gradient will only pass through the top K1 and sampled K2 states, in other words, not all states receive gradients. In this case, trading K1 against K2 amounts to exploration versus exploitation. A large K1 means we give gradients to high-confidence states, i.e., we exploit large local emission and global transition potentials. While increasing K2 means we explore low-confidence states. So, by splitting the computation budget between top K1 and sampled K2 states, we not only reduce variance for estimating the partition, but also effectively introduce different strategies for searching over the latent space.

heading: 4 RELATED WORK
text: Efficient Inference for Structured Latent Variables There has been substantial interest recently in the application of deep latent variable models (LVMs) to various language related tasks (Wiseman et al., 2018; Li & Rush, 2020), which has also exposed scalability limitations. Earlier attempts to render CRF models efficient (Sokolovska et al., 2010; Lavergne et al., 2010) either make many stringent assumptions (e.g., sparsity), rely on handcrafted heuristics for bias correction (Jeong et al., 2009), or cannot be easily adapted to modern GPUs with tensorization and parallelization. Sun et al. (2019) are closest to our work, however they only consider topK summation and consistently underestimate the partition. Chiu & Rush (2020) scale HMMs but assume words are clustered beforehand. Our approach systematically trades computation with proposal accuracy and estimation error (rather than over-compromising for efficiency). Moreover, we do not impose any hard restrictions like sparsity (Correia et al., 2020), and can accommodate dense and long-tailed distributions. Our method is inspired by randomized automatic differentiation (RAD, Oktay et al., 2020), and can be viewed as RAD applied to the DP computation graph. Advantageously, our proposal is compatible with existing efficient implementations (like Rush, 2020) since it does not change the computation graph.
Interpretability of Contextualized Representations There has been a good deal of interest recently in analyzing contextualized representations and the information they encode. This line of research, collectively known as “Bertology” (Rogers et al., 2020; Hewitt & Manning, 2019), focuses mainly on supervised probing of linguistic properties (Tenney et al., 2019), while the geometric properties of the representations have been less studied (Cai et al., 2021). A major dilemma facing this work is whether supervised linguistic probes reveal properties intrinsic to the embeddings or imposed by the supervision signal itself (Hewitt & Liang, 2019; Hall Maudslay et al., 2020; Chen et al., 2021). In this work, we do not use any supervision to ensure that the discovered network is intrinsic to the representation space.

heading: 5 EXPERIMENTS
text: In this section, we present our experimental results aimed at analyzing RDP and showcasing its practical utility (see Fig. 1). Specifically, we (1) verify the basic properties of RDP by estimating the partition function and (2) using it to train the structured latent variable model introduced in Section 3; (3) we then turn our attention to pretrained language models and examine the network induced with our approach and whether it is meaningful; and (4) we generate sentence paraphrases by traversing this network. For experiments (1, 2, 4), we use (a). pretrained GPT2 as the encoder since they are more about autoregressive language modeling and generation; (b). the MSCOCO dataset, a common benchmark for paraphrasing (Fu et al., 2019). For experiment (2), we use (a). BERT since it has been the main focus of most previous analytical work (Rogers et al., 2020); (b). the 20News dataset, a popular benchmark for training latent variable models (Grisel et al.). Across all experiments, we
use an LSTM decoder with states identical to the encoder (762 for BERT base and GPT2 as in Wolf et al., 2020). More details on experiments and model settings can be found in Appendix E.

heading: 5.1 BASIC PROPERTIES
text: We examine the estimation of the partition function for three unit cases, namely dense, intermediate, and long-tailed distributions. Instead of simulating these unit cases, to make our experiments more realistic, we extract CRFs on-the-fly from different LVM training stages. We also study the effects of memory budget by setting K to 20, 200, and 400 (corresponding to 1, 10, and 20 percent of the full memory). We use TopK summation (Sun et al., 2019) as our main baseline. This method can be viewed as setting K1 = K and K2 = 0 in our framework, i.e., it does not use the random sample. For training LVMs, We consider 100 and 2,000 latent states. With 100 states we are able to perform the summation exhaustively which is the same as Fu et al. (2020). Full summation with 2,000 states is intractable, so we only compare with TopK summation and use K = 100.
Estimating the Partition Function As shown in Figure 2, TopK summation always underestimates the partition. The gap is quite large in the dense case (large entropy), which happens at the initial stages of training when the model is not confident enough. The long-tailed case represents later training epochs when the model has converged and is more concentrated. Our method effectively corrects the bias, and works well in all unit cases with significantly less memory.
Training Latent Variable Models We compare different LVMs in Table 1. Following common practice, we report negative log likelihood (NLL) and perplexity (PPL). We perform an extensive search over multiple hyperparameters (e.g., β, learning rate, word dropout) across multiple random seeds (3–6) and report the average performance of the best configuration for each method. Our model performs best in both 100 and 2,000 state settings. The advantage is modest (as there are no architecture changes, only different training methods) but consistent. RDP trades off exploitation (i.e., increasing K1) and exploration (i.e., increasing K2) while TopK summation always focuses on the local solutions by passing gradients through top states. Intuitively, we have the chance of discovering better latent states (i.e., larger likelihood) by randomly searching the unexplored space.

heading: 5.2 DISCOVERING LATENT NETWORKS FROM PRETRAINED EMBEDDINGS
text: We now discuss how latent structures induced with RDP reveal linguistic properties of contextualized representations. We focus on BERT Devlin et al. (2019) and set the number of latent states to

heading: 1291 Literature lines 1848 | read 701 | writes 502 | line 376 | book 319 | books 244 | write 203 | written 177 | text 171
text: 2,000. As BERT’s vocabulary size is 32K, one state would approximately handle 15 words in the uniform case, functioning as a type of “meta” word. After convergence, we use qψ to sample z for each x in the training set (recall we use the 20News dataset). These z can be viewed as samples from the aggregated posterior ∑ x qψ(z|x)p?(x) where p?(x) denotes the empirical data distribution. To get a descriptive summary of BERT’s latent topology, we compute the following statistics on z samples: state frequency (Fig. 3, A3); number words corresponding to each state (Fig. 3, A2); number of states corresponding to each word (Fig. 3, A1); and state bigrams (Fig. 4). We further differentiate stopwords (e.g., in, of, am, is) from content words.
State-word Relations Figure 3 gives a first impression of how latent states spread over the representation space. Overall, we observe that the joint space is Zipfian, and this property characterizes
the distribution of words (A1), states (A3), and word occurrence within each state (C). We also see that the top 500 states account for most word occurrence (A2) while the remaining states model tail phenomena (A3). We conjecture this number is related to the intrinsic dimension of the data manifold (see Aghajanyan et al. 2021). The induced states encode multiple linguistic properties (Fig. 3, C). Some states are similar to a lexicon entry encoding specific words and their morphological variants; other states exhibit clustering based on morphological features (-s, -er, -ly suffix). We believe this is closely related to the fact that BERT learns embeddings over subwords. Note that the past tense cluster contains words exhibiting both regular (-ed suffix) and irregular morphology (e.g., lost and built). Finally, we also see that some states are largely semantic, similar to a conventional topic model (e.g., Computer and Medicine clusters). See Appendix E.6 for more state-word examples.
State-State Relations As shown in Fig. 4, we observe a clear geometric difference between top and tail states. Most linguistic constructions seem to be captured by the top 500 states (A1). The connections of top states are visualized in (A2–A4). From a statistical perspective, the similarity of (A3) and (A4) clearly shows how the empirical prior (encoded by the transition matrix Φ) matches
Model iB4↑ B4↑ sB4↓ CGMH (Miao et al., 2018) 7.84 11.45 - UPSA (Liu et al., 2020) 9.26 14.16 - GUMBEL-CRF (Fu et al., 2020) 10.20 15.75 -
the aggregated posterior (encoded in the bigram sample from qψ), which is an important desideratum of generative modeling (Mathieu et al., 2019). Note that the number of edges linked to each node, again, follows a Zipfian distribution as top nodes have most of the connections. From a linguistic perspective, we see how the combination of states leads to meaningful syntactic and semantic constructions. Again, BERT encodes various syntactic configurations such as to infinitives, passive voice, and even manages to distinguish adverbials (e.g., in fact) from prepositional phrases (e.g., in Bosnia). In general, the latent network seems to have some grasp of syntax, semantic roles, and collocations. In the following section, we examine whether this inherent knowledge can be harvested for generation. See Appendix E.7 for more state transition examples.

heading: 5.3 PARAPHRASING THROUGH NETWORK TRAVERSAL
text: We now study how the latent network can be usefully employed to generate paraphrases without access to parallel training instances. Given a sentence, we generate its paraphrase by conditioning on the input which we represent as a bag-of-words Fu et al. (2020) and by sampling from latent states. This amounts to traversing the latent network then fill in the traversal path to assemble a sentence, as visualized in Fig. 4 C. We instantiate our approach with a latent network learned from GPT2 representations (Radford et al., 2019) and refer to our model collectively as GPTNET.
We compare against three previous unsupervised models (first block in Table 2), including CGMH (Miao et al., 2019), a general-purpose MCMC method for controllable generation; UPSA (Liu et al., 2020), a strong paraphrasing model with simulated annealing, and GUMBEL-CRF (Fu et al., 2020), a template induction model based on a continuous relaxation of the CRF sampling algorithm. We present GPTNET variants with 50 and 2,000 states, and show results with RDP and topK, and the full summation for 50 states. Following previous work, we use iBLEU (Sun & Zhou, 2012) as our main metric, which trades off fidelity to the references (BLEU) and variation from the input (self-BLEU). Table 2 shows that RDP is superior to TopK and full summation in terms of iBLUE. GPTNet models do not outperform GUMBEL-CRF or UPSA. This is expected as these methods are highly tailored to the task and more flexible (e.g., they do not fix the encoder), while we restrict the modeling within the GPT2 representation space (to infer its structure). So, our results should be viewed as a sanity check demonstrating the latent network is indeed meaningful for generation (see Appendix E.8 for more generation examples).

heading: 6 CONCLUSION
text: In this paper, we have developed a general method for scaling the inference of structured latent variable models with randomized dynamic programming. It is a useful tool for the visualization and inspection of the intrinsic structure of contextualized representations. Experiments with BERT reveal the topological structure of its latent space: state-word connections encapsulate syntactic and semantic roles while state-state connections correspond to phrase constructions. Moreover, traversal over a sequence of states represents underlying sentence structure.
Ethics Statement As this paper inspects the internal structure of pretrained language models, it is likely that it will reveal frequent linguistic patterns encoded in the language model. Specifically, the frequent words, phrases, and sentences associated with different gender, ethnic groups, nationality, interest groups, social status, and all other factors, are likely to be revealed by our model. When calling the generative part of our model for paraphrasing, these differences are likely to exist in the generated sentences (depending on the dataset). These facts should be considered when using this model.
Reproducibility Statement A step-by-step implementation guide for our randomized forward algorithm is provided in Appendix section C. The comparison of RDP versus other possible solutions for scaling the structured models is provided in Appendix section B. A detailed description of the model architecture is provided in the Appendix section E.1. A detailed description of data processing is provided in the Appendix section E.2. A detailed description of training strategy, hyperparameter search strategy, and model selection, is provided in Appendix section E.3. A detailed description of visualization procedure is provided in Appendix section E.5. We will release code after the anonymity period.

heading: Appendix, table of contents
text: • Section A. Theoretical analysis of the Sampled Forward algorithm. Bias and variance. • Section B. Challenges of full DP inference under Automatic Differentiation. • Section C. Implementation details of Sampled Forward. • Section D. Extensions of randomized dynamic programming. • Section E. Experimental details. • Section F. Additional experimentsl results.

heading: A THEORETICAL ANALYSIS OF SAMPLED FORWARD ALGORITHM

heading: A.1 BIAS ANALYSIS
text: In this section, we discuss the unbiasedness and variance of the Randomized Forward algorithm. We first show that Randomized Forward gives an unbiased estimator of the partition function. Theorem A.1 (Unbiasedness). For all t ∈ [1, 2, ..., T ], the sampled sum α̂t (Eq. 6) is an unbiased estimator of the forward variable αt (Eq. 3). The final sampled sum Ẑ (Eq. 7) is an unbiased estimator of the partition function Z (Eq. 3).
Proof. By the Second Principle of Mathematical Induction. Assume initialization α1(i) = φ(x1, i). Firstly at t = 2, for all i, we have:
Eq1 [α̂2(i)] = K1∑ j=1 α1(σ1,j)Φ(σ1,j , i)φ(x2, i) + 1 K2 K2∑ j=1 Eq1 [ Z̃1 q̃1(δ1,j)
α1(δ1,j)Φ(δ1,j)φ(x2, i)︸ ︷︷ ︸ =A
]
(14)
where the second term can be expanded as a masked summation with the index rearranged from σ2,K1+1 to σ2,N :
A = K2∑ j=1 Eq1 [ Z̃1 q̃1(δ1,j) α1(δ1,j)Φ(δ1,j)φ(x2, i) ] (15)
= K2∑ k=1 N∑ j=K1+1 Eq1 [ 1 q1(δ1,k) 1(δ1,k = j)α1(σ1,j)Φ(σ1,j)φ(x2, i) ] (16)
= K2∑ k=1 N∑ j=K1+1 Eq1 [ 1 q1(δ1,k) 1(δ1,k = j) ] ︸ ︷︷ ︸
=1
α1(σ1,j)Φ(σ1,j)φ(x2, i) (17)
= K2 N∑ j=K1+1 α1(σ1,j)Φ(σ1,j)φ(x2, i) (18)
Notice how we re-index the sum. Now put it back to Eq. 14 to get:
Eq1 [α̂2(i)] = K1∑ j=1 α1(σ1,j)Φ(σ1,j , i)φ(x2, i) + 1 K2 ·K2 N∑ j=K1+1 α1(σ1,j)Φ(σ1,j)φ(x2, i) (19)
= N∑ j=1 α1(σ1,j)Φ(σ,j , i)φ(x2, i) (20)
= N∑ j=1 α1(j)Φ(j, i)φ(x2, i) (21)
= α2(i) (22)
This verifies the induction foundation that Eq1 [α̂2(i)] = α2(i) for all i.
Now assume for time index t we have ∀i,Eq1:t−1 [α̂t(i)] = α2(i). Consider t+ 1, we have:
E1:qt [α̂t+1(i)] = K1∑ j=1 Eq1:t−1 [ α̂t(σt,j) ] Φ(σt,j , i)φ(xt+1, i) (23)
+ 1
K2 K2∑ j=1 Eq1:t−1 [ α̂t(δt,j) ] · Eqt [ Z̃t q̃t(δt,j) Φ(δt,j)φ(xt+1, i)] ]
(24)
= K1∑ j=1 αt(σt,j)Φ(σt,j , i)φ(xt+1, i) (25)
+ 1
K2 K2∑ j=1 αt(σt,j) · Eqt [ Z̃t q̃t(δt,j) Φ(δt,j)φ(xt+1, i)] ]
︸ ︷︷ ︸ =A′
(26)
Note how we decompose the expectation by using the independence: q1:t = q1:t−1 ·qt With a similar masked summation trick as A, we have:
A′ = K2 N∑ j=K1+1 αt(σt,j)Φ(σt,j , i)φ(xt+1, i) (27)
This gives us:
E1:qt [α̂t+1(i)] = K1∑ j=1 αt(σt,j)Φ(σt,j , i)φ(xt+1, i) + 1 K2 ·K2 N∑ j=K1+1 αt(σt,j)Φ(σt,j , i)φ(xt+1, i)
(28)
= N∑ j=1 αt(σt,j)Φ(σt,j , i)φ(xt+1, i) (29)
= N∑ j=1 αt(j)Φ(j, i)φ(xt+1, i) (30)
= αt+1(i) (31)
Thus showing α̂t is an unbiased estimator for αt at each step t. Setting t = T , the last step, gives us E[Ẑ] = Z (details similar to the above).
Corollary A.1.1. When we change the (sum, product) semiring to the (log-sum-exp, sum) semiring, the expectation of the estimator will become a lower bound of logαt and logZ.
Proof. Denote lt(i) = logαt(i) and Ω the set of sampled indices where |Ω| = K2. For simplicity, we omit the top K1 summation and only show the summation of the sample. Cases where t > 2 can be derived similarly as following:
l̂2(i) = log ∑ j∈Ω exp(l1(j) + log Φ(j, i) + log φ(xt, i))− logK2 (32)
E[l̂2(i)] = E [ log ∑ j∈Ω exp(l1(j) + log Φ(j, i) + log φ(xt, i)) ] − logK2 (33)
≤ logE [∑ j∈Ω exp(l1(j) + log Φ(j, i) + log φ(xt, i)) ] − logK2 (34)
= logK2 − logK2 + log ∑ j∈Ω exp(l1(j) + log Φ(j, i) + log φ(xt, i)) (35)
= l2(i) (36)
where Eq. 34 comes from Jensen’s inequality. Then by induction one can show at everystep, we have E[l̂t(i)] ≤ lt(i).
Although implementation in the log space makes the estimate biased, it reduces the variance exponentially in a rather trivial way. It also provides numerical stability. So, in practice we use it for training.
A.2 VARIANCE ANALYSIS
Now we analyze variance. We start with the estimator aδ/qδ , δ ∼ Categorical{qK1+1, ...qKN } in Eq. 2. Firstly, this is an unbiased estimator of the tail sum:
E[aδ/qδ] = N∑
i=K1+1
ai (37)
We have the folling variance arguments: Theorem A.2 (Variance of the tail estimator).
V[aδ/qδ] = N∑
i=K1+1
a2i /qi − ( N∑
i=K1+1
ai) 2 (38)
= S2K1( N∑ i=K1+1 2i /qi + 2 i) (39)
where
SK1 = N∑ i=K1+1 ai i = ai/SK1 − qi (40)
which says the variance is:
• quadratic to the tail sum SK1 , which will can be reduced by increasing K1 as the effect of Rao-Blackwellization.
• approximately quadratic to the gap i, as the differences between the proposal qi and the oracle ai/SK1 , which can be reduced by choosing a correlated proposal as the effect of importance sampling.
Optimal zero variance is achieved on
q∗i = ai/SK1 (41)
Proof. The variance is then:
V[aδ/qδ] = E[(aδ/qδ)2]− E[aδ/qδ]2 (42)
where the first term is the second moment:
(aδ/qδ) 2 = ( N∑ i=K1+1 ai1[δ = i]/qi) 2 (43)
= N∑ i=K1+1 a2i1[δ = i]/q 2 i + 2 N∑ i=K1+1 N∑ j=K1+1 aiaj1[δ = i]1[δ = j]
qiqj︸ ︷︷ ︸ =0
(44)
= N∑ i=K1+1 a2i1[δ = i]/q 2 i (45)
taking the expection of it we get:
E[(aδ/qδ)2] = E[ N∑
i=K1+1
a2i1[δ = i]/q 2 i ] (46)
= N∑ i=K1+1 a2i /qi (47)
Plug this back, variance is:
V[aδ/qδ] = N∑
i=K1+1
a2i /qi − ( N∑
i=K1+1
ai) 2 (48)
To get the optimal proposal, we solve the constrained optimization problem:
min qi N∑ i=K1+1 a2i /qi − ( N∑ i=K1+1 ai) 2 (49)
s.t. N∑ i=K1+1 qi = 1 (50)
By solving the corresponding Lagrangian equation system (omitted here), we get the optimal value achieved at:
q∗i = ai∑N
i=K1+1 ai
N∑ i=K1+1 a2i /q ∗ i − ( N∑ i=K1+1 ai) 2 = 0 (51)
This says zero variance is achieved by a proposal equal to the normalized summands.
Then define the gap between the proposal and the normalized summands as:
i = ai SK1 − qi (52)
Plug this to the variance expression we get:
V[aδ/qδ] = S2K1( N∑
i=K1+1
2i /qi + 2 i) (53)
Corollary A.2.1. When increasing the sample size to K2, the variance will reduce to
V[ 1
K2 N∑ j=1 aδj qδj ] = 1 K2 S2K1( N∑ i=K1+1 2i /qi + 2 i) (54)
Now we consider the variance of the Sampled Forward algorithm. An exact computation would give complicated results. For simplification, we give an asymptotic argument with regard to RaoBlackwellization and importance sampling: Theorem A.3 (Single Step Asymptotic Variance of Sampled Forward). At each step, the alpha varible estimator has the following asymptotic variance:
V[α̂t(i)] = O ( 1 K2 α2t,K1(i) · 2 t (i) )
(55)
where:
• αt+1,K1(i) = ∑N j=K1+1 α̃t(j, i) = ∑N j=K1+1 √ Eq1:t−1 [α̂2t (j)]Φ(j, i)φ(xt, i) is a tail sum
after the topK1 summands. This term will reduce if we increaseK1, as an instance of RaoBlackwellization.
• 2t (i) = ∑N j=K1+1
2t−1(j, i)/qt−1(j) and t−1(j, i) is the difference between the proposal qt−1(j) and the oracle proposal in Eq. 3. This term will reduce if the proposal is more correlated to the oracle, as an instance of Importance Sampling.
Proof. We start with a simple setting where K2 = 1. At step t + 1 we have the following variance recursion:
Eq1:t [α̂2t+1(i)] = N∑
j=K1+1
Φ2(j, i)φ2(xt+1, i)
qt(j) · Eq1:t−1 [α̂2t (j)] (56)
This is derived by plugging estimator 6 to the variance Eq. 38 we have just derived. Denote:
αt+1,K1(i) = N∑ j=K1+1 α̃t(j, i) = N∑ j=K1+1 √ Eq1:t−1 [α̂2t (j)]Φ(j, i)φ(xt+1, i) (57)
Then we have
Vq1:t [α̂t+1(i)] = α2t+1,K1(i) ( N∑ j=K1+1 2t (j, i) qt(i) + 2 t(j, i) ) (58)
where t(j, i) is the differences between the proposal and the normalized exact summands at step t state i:
t(j, i) = α̃t(j, i)∑N
j=K1=1 α̃t(j, i)
− qt(j) (59)
Dropping out the first order errors and increasing the number of sample to K2, we have the asymptotics:
V[α̂t(i)] = O ( 1 K2 α2t,K1(i) · 2 t (i) )
(60)
Theorem A.4 (Asymptotic Variance of Sampled Forward Partition Estimation). The alpha variable estimators has the following asymptotic variance recurrsion:
V[α̂t+1(i)] = O( 1
K2 · φ2t,K1 · 2 t,K1 · V[α̂t]) (61)
Compared with Eq. 55, this expression:
• Uses the product of the factors φt,K1 (a function of the sum-prod of the factor at step t) and the previous step variance V[α̂t] to substitute the αt,K1 in equation 55. Again, this term will decrease with a larger K1 (Rao-Blackwellization).
• 2t,K1(i) = ∑N j=K1+1
2t−1(j, i)/qt−1(j) and t−1(j, i) is the difference between the proposal qt−1(j) and the oracle proposal in Eq. 3. This term will reduce if the proposal is more correlated to the oracle, as an instance of Importance Sampling (same as Eq. 55).
Consequently, the partition function has the following asymptotic variance:
V[Ẑ] = O( T∏ t=1 1 K2 · φ2t,K1 · 2 t,K1) (62)
When implemented in the log space, the variance is trivially reduced exponentially:
V[log Ẑ] = O( T∑ t=1 log 1 K2 + 2 log φt,K1 + 2 log t,K1) (63)
Proof. Firstly for simplcity we assume K1 = 0 and K2 = 1. The the estimator variance is:
V[α̂t+1(i)] = Eq1:t [α̂2t+1(i)]− α2t+1(i) (64) = Eq1:t−1 [ N∑ j=1 α̂2t (j)Φ 2(j, i)φ2(xt, i) qt(j) · α 2 t (j) α2t (j) ] − α2t+1(i) (65)
Recall:
αt+1(i) = N∑ j=1 αt(j)φ(j, i)φ(xt, i) (66)
Let:
t(j, i) = αt(j)φ(j, i)φ(xt, i)
αt+1(i)︸ ︷︷ ︸ =q(zt=j|zt+1=i)
−qt(j) (67)
Then: V[α̂t+1(i)] = Eq1:t−1 [ N∑ j=1 α2t+1(i) ( 2t (j, i) qt(j) + 2 t(j, i) + qt(j) ) α̂2t (j) α2t (j) ] − α2t+1(i) (68)
= α2t+1(i) ( N∑ j=1 (Eq1:t−1 [α̂2t (j)] α2t (j) qt(j) ) + N∑ j=1 ( 2t (j, i) qt(j) + 2 t(j, i) )Eq1:t−1 [α̂2t (j)] α2t (j) ) (69)
− α2t+1(i) (70) Note that there exist Jt such that:
Eq1:t−1 [α̂2t (j)] α2t (j) <= Jt (71)
This is because of the bounded gap of the Jensen’s inequality. Also recall: N∑ j=1 qt(j) = 1 (72)
So we get: V[α̂t+1(i)] <= α2t+1(i) ( Jt − 1 + N∑ j=1 ( 2t (j, i) qt(j) + 2 t(j, i) ) · (V[α̂t(j)] α2t (j) − 1 ))
(73)
= α2t+1(i) ( Jt − 1− N∑ j=1 ( 2t (j, i) qt(j) + 2 t(j, i) ))
(74)
+ α2t+1(i) ( N∑ j=1 ( 2t (j, i) qt(j) + 2 t(j, i) ) · (V[α̂t(j)] α2t (j) )) (75)
empirically Jt is not the dominate source of variance (but could be in the worst case, depending on the tightness of Jensen’s inequality). We focus on the second term:
V[α̂t+1(i)] = O ( α2t+1(i) ( N∑ j=1 ( 2t (j, i) qt(j) + 2 t(j, i) ) · (V[α̂t(j)] α2t (j) )) (76)
= O (( N∑ j=1
αt(j) Φ(j, i)φ(xt, j)︸ ︷︷ ︸ O(φ2t )
)2 · ( N∑ j=1 ( 2t (j, i) qt(j) + 2 t(j, i) )
︸ ︷︷ ︸ O( 2t )
· 1 α2t (j) · V[α̂t(j)]︸ ︷︷ ︸ O(V[αt])
))
(77)
= O(φ2t · 2t · V[αt]) (78)
Note that a lot of higher-order sum-products are simplified here. Adding top K1 summation and increasing the sample size to K2 leads to variance reduction as:
V[α̂t+1(i)] = O( 1
K2 · φ2t,K1 · 2 t,K1 · V[αt]) (79)
Recursively expand this equation we get:
V[Ẑ] = O( T∏ t=1 1 K2 · φ2t,K1 · 2 t,K1) (80)
Chaning the implementation to the log space we reduce the variance exponentially:
V[log Ẑ] = O( T∑ t=1 log 1 K2 + 2 log φt,K1 + 2 log t,K1) (81)

heading: B CHALLENGES OF FULL DP UNDER AUTOMATIC DIFFERENTIATION
text: Now we analyze in detail why direct full summation causes memory overflow. Firstly, we note that in a bachified stochastic gradient descent setting, memory complexity for the forward algorithm is
O(BTN2) (82)
Where B is the batch size, T is the maximum sequence length after padding, and N is the number of states. Consider a typical setting where B = 100, T = 100, N = 1000, then the complexity is:
C = c× 100× 100× 1000× 1000 = c× 1010 (83) where c is a constant depending on specific hardware and libraries. At first sight this may seem reasonable with a large GPU memory (e.g., 16G). Also, one can come up with improvements, such as not storing the intermediate αt variables. If we only need one single forward pass, such engineering tricks are indeed effective. In our experiments, when setting B = 50, T = 15, and N = 2000, the actual memory consumption for the forward computation, together with other model components, is about 4G when implemented with PyTorch 1.8.0.
However, this is not the case when working under an automatic differentiation (AD) setting. If we want to compute the gradients of the partition function (e.g., to optimize the likelihood), we inevitably need to store all the intermediate steps to keep the full computation graph. Then, the adjacent graph, in principle, has the same complexity as the forward graph. But in practice, it will be more complicated due to the actual implementation of AD engines. Following the above example where B = 50, T = 15, N = 2000, we get a memory overflow on a 16G memory GPU when calling the PyTorch backward function. This situation also immediately invalidates many engineering tricks (e.g., we cannot drop the intermediate αt variables anymore), and substantially increases the difficulty of coming up with new ones, since we would also be working with the internal mechanism of automatic differentiation engines. Note that the internal mechanism of these engines is complicated and opaque to general practitioners (e.g., they may change the underlying computation graph for better speed), and many of these engines (like PyTorch) are not optimized for dynamic programming.
In fact, it would be unwise to overwrite the backward computation, even if the AD engine allowed us to do so, as it would significantly increase engineering difficulty, as we not only need to overwrite the first order gradients (like the gradients for the likelihood), but the second order gradients too (e.g., gradients for marginals or reparameterized samples). In fact, a brute force implementation would not only forego all the advantages of AD engines (like operator-level optimization), but would also require separate implementations for every graph (e.g., chains, trees, semi-Markovs, and IsingPotts), every inference algorithm (e.g., partition, marginal, sampling, reparameterization, entropy), and higher order gradients. In summary, this is an extremely difficult path that requires clever tricks to efficiently improve a large table of already complicated DP algorithms.
This is why we do not interfere with the AD framework and work on more general and efficient algorithms, rather than tailored implementation tricks. With the randomized DP, we do not re-implement
anything. We only choose the index by the pre-computed proposal, and introduce the re-indexed potentials to existing efficient libraries (like Torch-struct in Rush, 2020). This is least difficult from an implementation perspective and best exploits the power of efficient structured prediction libraries and AD engines.
C IMPLEMENTATION OF SAMPLED FORWARD ALGORITHM
Now we discuss how to implement the sample forward in detail. Our implementation aims to set all computation outside the actual DP, and reuse existing optimized DP libraries. Specifically, we:
1. normalize the log potentials; 2. construct the proposal according to the local potential and a global prior; 3. obtain the top K1 index and sample K2 from the rest, retrieve the potentials according to
these indices; 4. correct the bias before the DP, then insert the corrected potentials to an existing DP imple-
mentation.
Now we explain each step in detail.
Potential Normalization is a technique useful for numerical stability (see Fu et al., 2021). Specifically, in practice, the log-sum-exp function will saturate (it returns the max input value) when the maximum difference between inputs is larger than 10. So, we normalize all log potentials to be within range [1,m],m < 10. Given the emission potential φ and transition potential Φ from Eq. 8, we normalize as:
φ̃(xt, i) = m ∗ log φ(xt, i)−mini log φ(xt, i)
maxi log φ(xt, i)−mini log φ(xt, i) (84)
Φ̃(i, j) = m ∗ log Φ(i, j)−mini,j log Φ(i, j) maxi,j log Φ(i, j)−mini log Φ(i, j)
(85)
Then the actual CRF is constructed based on the normalized potentials. Proposal Construction Our proposal is constructed based on the normalized local emission and a
prior distribution. Specifically:
qt(i) = 1 2 · ( exp φ̃(xt, i)∑N
i=1 exp φ̃(xt, i) + exp ||si||1∑N i=1 exp ||si||1 ) (86)
Index Retrieval For each step t, we retrieve the top K1 index and get a K2 sized sample from the rest, as is discussed in the main paper:
[σt,1, ..., σt,K1 , ..., σt,N ] = arg sorti{qt(i)}Ni=1 (87) [δt,1, ..., δt,K2 ] ∼ Categorical{qt(σt,K1 + 1), ..., qt(σt,N )} (88)
Note that for simplicity, we use sampling with replacement. For sampling without replacement, see the procedure in Kool et al. 2020
Conventional Forward Before calling an existing implementation of Forward, we need to correct the bias of sampled terms. We do this by multiplying the probability of the sample to its emission potential (which becomes addition in log space). Note that this will be equivalent to multiplying them inside the DP:
φ̃′(σt,i) = φ̃(σt,i) (89)
φ̃′(δt,i) = φ̃(δt, i) + log qt(δt,i) (90)
We treat any duplicate indices as if they are different, and view the resulting potentials as if they are a new CRF. We also retrieve the transition potentials at each step according to the chosen index. So the new CRF has different transitions across steps. Finally, we run an existing efficient implementation of the Forward algorithm on the new CRF to get the estimated α̂ and Ẑ.

heading: D EXTENSION OF RDP
text: In this section, we extend the RDP to more graph structures and more inference operations. We fisrt discuss a randomized inside algorithm for the partition function estimation for tree-structured hypergraphs, which includes dependency trees and PCFGs. Then we discuss a randomized entropy DP estimation on chain-structured graphs. This estimation will call the randomized forward as its subroutine. Finally we discuss howto generalize RDP to the general sum-product algorithm for any graph structure.

heading: D.1 EXTENSIONS TO TREE-STRUCTURED HYPERGRAPHS: RANDOMIZED INSIDE
text: This section discusses how to extend RDP to Tree-structured hypergraphs, including Dependency TreeCRFs and PCFGs. We focus on randomizing the Inside algorithm for partition function estimation. Specifically, the core recursion of the inside is:
β[i, j, k] = sijk j−1∑ l=i ∑ k1∈Ω,k2∈Ω β[i, l, k1]β[l, j, k2] (91)
where β[i, j, k] is the summation of all subtrees spanning from location i to location j with label k and sjik is the local score and Ω is the full state space. Suppose Ω is large so we want to sample its subset Ωij to reduce the summation computation. Suppose the proposal is qijk where ∑ k∈Ω qijk = 1. Then a randomized inside recursion with K2 sample is (we omit the top K1 summation for simplicity):
β[i, j, k] = sijk j−1∑ l=i ∑ δ1∈Ωil,δ2∈Ωlj
1
K2qilδ1 β[i, l, δ1] ·
1
K2qljδ2 β[l, j, δ2] (92)
The analysis about bias and variance is similar to the previous analysis (Sec. A) on the linear-chain case.

heading: D.2 EXTENSIONS TO RANDOMIZED ENTROPY DP ESTIMATION
text: This section describes a randomized DP for estimating the entropy of a chain-structured model (HMMs and linear-chain CRFs). Recall that the core recursion of conventional entropy DP is:
p(zt = i|zt+1 = j) = Φ(i, j)φ(xt+1, j)αt(i)
αt+1(j) (93)
Ht+1(j) = N∑ i=1 p(zt = i|zt+1 = j)[Ht(i)− log p(zt = i|zt+1 = j)] (94)
where Ht(i) is the intermediate conditional entropy end at state i step t. In our sampled DP, we first call the sampled forward to estimate the alpha variables. Then we re-use the sampled indices for the entropy DP graph, and the core recursion becomes:
p̂(zt = i|zt+1 = j) = Φ(i, j)φ(xt+1, j)α̂t(i)
α̂t+1(j) (95)
Ht+1(j) = K2∑ δt=1
1
K2 · qt(δt) p̂(zt = δt|zt+1 = j)[Ht(i)− log p̂(zt = δt|zt+1 = j)]
(96)
where qt is the proposal at step t. Note how we re-use the estimated alpha variables α̂ for our entropy DP and correct the bias of each step.

heading: D.3 EXTENSIONS TO GENERAL SUM-PRODUCT
text: This section discusses the extension of sampled DP to general graph structures and message-passing algorithms, following the Bethe variational principle (Wainwright & Jordan, 2008). Recall the general message-passing algorithm computes pseudo marginals by recursively updating the message at
each edge:
Mts(xs)← κ ∑ x′t { φst(xs, x ′ t)φt(x ′ t) ∏ u∈N(t)/s Mut(x ′ t) } (97)
where Mts(xs) denotes the message from node t to node s evaluated at Xs = xs, φst is the edge potential, φt is the node potential, N(t)/s denotes the set of neighbor nodes of t except s, and κ the normalization constant. At convergence, we have the pseudo marginals
µs(xs) = κφs(xs) ∏
t∈N(s)
M∗ts(xs) (98)
and the Bethe approximation of the log partition function is given by the evaluation of the Bethe variational problem:
logZBethe = φ ᵀµ+ ∑ s Hs(µs)− ∑ s,t Ist(µst) (99)
where H denotes marginal entropy and I denotes marginal mutual information.
We consider the application of randomization to the computation of pseudo marginals and the Bethe log partition. Note that the Bethe approximation will be exact if the underlying graph is a tree. We consider a local proposal q̃ combined with a global prior τ :
q(xs) = 1
2 (q̃(xs) + τ(xs)) (100)
where the prior τ depends on our knowledge of specific problem structures. As one can always retreat to uniform proposals, we can explore more advanced choices. For q̃, one may consider: (a) local normalization as (̃q)(xs) = φ(xs)/ ∑ s φ(xs) and (b) pre-computed mean-field approximations from algorithms like SVI (Hoffman et al., 2013). For τ , one can further use a frequency-based empirical prior estimated from the data. To determine which nodes to perform DP on, one may consider the following principles:
• Special nodes of interest depending on their actual meaning. For example, one may be interested in some xs = 0 (e.g., if a user is under a bank fraud threat).
• Nodes with large local weight φ(xs). • Nodes with large global weight τ(xs). • Loop elimination. If two nodes have similarly small local and global weight, we could
drop those which eliminate loops in the resulting graph. Also note that we would prefer removing small-weighted nodes for loop elimination.
With the above principles, one may construct three subsets to perform the sum-product:
• Ω1 including nodes of special interest, where we perform exact computation. • Ω2 from the top items of the proposal, where we also perform exact computation. • Ω3 by sampling from the remaining items of the proposal (optionally with loop-reduction).
For this set of nodes, one needs to correct the estimation by dividing the proposal probabilty: φ(xs)← φ(xs)/q(xs).
After these steps, we treat nodes in Ω1 ⋃ Ω2 ⋃
Ω3 as if they are a new model, then feed them to an existing sum-product implementation.

heading: E EXPERIMENT DETAILS

heading: E.1 MODEL ARCHITECTURE DETAILS
text: For training LVMs, partition function estimation, and paraphrasing, we use the Huggingface checkpoint of GPT2 model1 since these experimental variables are more about autoregressive language
1https://huggingface.co/transformers/model_doc/gpt2.html
modeling and generation. For analyzing latent network topologies, we use Huggingface checkpoint of BERT base model2 since it has been the main focus of most previous analytical work. The decoder LSTM is a one-layer LSTM with hidden state size 762, the same as the size of the contextualized embeddings. It shares the embeddings of the inferred states with the encoder, and uses its own input word embeddings, whose size is also 762. This architecture suffices for training LVMs and inferring latent networks. For paraphrase generation, we change the decoder to be conditional by: (a) using the average word embeddings of content words of the input sentence as its initial hidden state (rather than zero hidden states); (b) letting it attend to the embeddings of content words of the input sentence, and copy from them. This effectively makes this decoder conditional on the BOW of the content words of the input sentence. Then decoding a paraphrase becomes a neural version of slot filling: a sequence of latent states by traversing the network becomes a template of the sentence which we fill with words.

heading: E.2 DATASET PREPROCESSING
text: For the MSCOCO dataset, we use the GPT2 tokenizer to process sentences. As the official website 3 does not release a test split (only a training and a development split), we use the official development split as our test split, and re-split the official training split to be a new training and a development split. For training LVMs, we only use 1/10 of the training set for quicker experiments (the full dataset takes over more than 20 hours to converge). For paraphrasing, we use the full dataset.
For the 20News dataset, we use the data from the sklearn website4. We follow its official split and process the sentences with the BERT tokenizer.

heading: E.3 HYPERPARAMETERS AND TRAINING STRATEGIES
text: To get meaningful convergence of the latent space without posterior collapse, the following techniques are important: (a). set β in the correct range. A large β force the posterior to collapse to a uniform prior, while a small β encourages the posterior to collapse to a Dirac distribution. (b). use word dropout in initial training epochs, otherwise the decoder may ignore the latent code. (c). use potential normalization, otherwise the logsumexp function used in the forward algorithm may saturate and only return the max of the input, this would consequently lead the posterior to collapse to few “activated” states and other states will never be inferred.
We further note that a full DP based entropy calculation will cause memory overflow with automatic differentiation. So we approximate it with local emission entropy, i.e., the sum of the entropy of the normalized local emission factors. Compared with the full entropy, this approximation does not directly influence the transition matrix, but encourages more activated states and mitigates the posterior collapse problem. To get a full entropy regularization, one can further regularize the transition matrix towards an all-one matrix, or extend our randomized DP to entropy computation. As the current local entropy regularization is already effective for inducing a meaningful posterior, we leave the full entropy computation to future work.
Our reported numbers are averaged over “good enough” runs. Specifically, for all hyperparameter configurations, we first run three random seeds and get the mean and standard deviation of the performance metrics. If the standard deviation is small enough, we report mean performance metrics (NLL, PPL, and iBLEU) and the standard deviation. If the standard deviation is large, we run extra five seeds. Then we drop runs with bad performance (usually 2-3), and compute the mean and standard deviation of the performance metrics again. In total, we experiment more than 100 runs over more than 7 hyperparameters: (a). learning rate 10−3, 10−4, 5 × 10−4; (b). optimizer: SGD, Adam, AdamW; (c). dropout: 0.2, 0.3, 0.4; (d). β variable: 10−4, 5 × 10−4, 10−3, 10−2; (e). Gumbel CRF v.s. Gumbel CRF straight-through (f). K1 v.s. K2; (g). word dropout schedule; and their combinations. Note that different models may achieve best performance with different hyperparameter combinations. We make sure that all models are searched over the same set of
2https://huggingface.co/transformers/model_doc/bert.html 3https://cocodataset.org/#home 4https://scikit-learn.org/stable/modules/generated/sklearn.datasets.
fetch_20newsgroups.html
hyperparameter combinations, and report average performance metrics over multiple seeds of the best hyperparameter configuration for each model.

heading: E.4 TIME COMPLEXITY ANALYSIS
text: Table 3 shows the actual efficiency comparison of our method. Time is measured by seconds per batch. For experiments with 500 and 2,000 states, we set K = 100. *When the number of states is small, our method underperforms due to the overhead of constructing the proposal. However, its advantage becomes clear as the number of states becomes large.
E.5 VISUALIZATION PROCEDURE
This section describes how we produce Fig. 3(B1-2) and Fig. 4(A1-4). We use the sklearn implementation 5 of tsne (Van der Maaten & Hinton, 2008). For Fig. 3(B1-2), the word embeddings are obtained by sampling 8,000 contextualized embeddings from the full word occurrences in the 20News training set. Then we put the sampled word embeddings and the 2,000 states into the tsne function. The perplexity is set to be 30. An important operation we use, for better seperating the states from each other, it to manually set the distances between states to be large, otherwise the states would be concentrated in a sub-region, rather than spread over words. Fig. 4 is produced similarly, except we do not use word embeddings as background, and change the perplexity to be 5. For Fig. 4(A3), we connect the states if their transition potential is larger than a threshold. For Fig. 4(A4), we connect the states if their bigram frequency is larger than a threshold. All our decisions of hyperparameters are for the purpose of clear visualization which includes reducing overlapping, overcrowding, and other issues. We further note that no single visualization method can reveal the full structure of high-dimensional data, and any projection to 2-D plain inevitably induces information loss. We leave the investigation of better visualization methods to future work.

heading: E.6 STATE-WORD PAIR EXAMPLES
text: See Tables 8 and 9 for more sample.

heading: E.7 STATE-STATE TRANSITION EXAMPLES
text: See Table 10 for transitions involving stopwords and Table 11 for transitions without stopwords. Also note their differences as transitions involving stopwords are more about syntactic constructions and transitions without stopwords are more about specific meaning.

heading: E.8 STATE TRAVERSAL EXAMPLES
text: See Table 12 for more sample.

heading: F ADDITIONAL EXPERIMENT RESULTS

heading: F.1 RDP V.S. TOPK MSE COMPARISON
text: Table 4 shows the mean square error (MSE) comparison between RDP and TopK on dense, intermediate and long-tail distributions. Again, MSE for RDP is significantly smaller with less memory consumption.

heading: F.2 MSE COMPARISON OVER PROPOSALS
text: Table 5 shows the mean square error (MSE) comparison between different proposals. Our proposed local and global proposal outperforms a baseline uniform proposal on all distributions (dense, intermediate and long-tail).

heading: F.3 CONTROLLING NUMBER OF STATES
text: Figure 5 shows state frequency with different N. The long-tail behavior becomes clearer only when N is large enough (larger than 500 in our case).

heading: F.4 CONTROLLING K1 V.S. K2 RATIO
text: Figure 6 show state frequency with different K1/K2 ratios at different N. We highlight that when K2 = 0, a pure topK summation approach would lead to posterior collapse where there exist inactive states that do not have any density. We also notice that an increasing K2 consistently increases the frequency of tail states. This observation can be explained by the exploitation-exploration tradeoff, where increasing K1 would encourage the exploitation of states already confident enough during training (consequently leading to high-frequency head states after convergence) while increasing K2 would encourage exploring states that are not confident enough during training, leading to a larger tail frequency after convergence.

heading: F.5 COMPARISON TO RANDOMLY INITIALIZED BERT
text: Figure 7 shows the comparison of reconstruction NLL (− log p(xt|zt, ·)) between a randomly initialized BERT and a pretrained BERT. States induced from pretrained BERT are more meaningful than random, making it easier to reconstruct words based on their corresponding states.
5https://scikit-learn.org/stable/modules/generated/sklearn.manifold. TSNE.html

heading: F.6 RANDOMIZED INSIDE ALGORITHM FOR TREE-STRUCTURED HYPERGRAPH
text: We simulate three unit cases of treecrfs (dense, intermediate and long-tail) by controlling the entropy. We set the number of latent states to 2000. We use a uniform proposal for comparison, and set K1 (top K sum size) = K2 = 50. We use RDP with 20, 200, 400 states, which correspond to 1%, 10%, and 20% of the full number of states. We use topK summation as our baseline with 500 and 1000 states (25% and 50% of the full number of states). Figure 8 shows performance of RDP to tree structured hypergraph. Our method still outperforms topK summation (grey lines) with less memory (as our estimates are visually closer to the red line, i.e., the true partition). Table 6 shows the MSE of our method compared to topK summation, and the results are consistent with Figure 8.

heading: F.7 RANDOMIZED ENTROPY DP ALGORITHM
text: Figure 9 shows the application of RDP to entropy estimation of linear-chain CRFs. Our method consistently outperforms topK summation (grey lines) with less memory (as our estimates are visually closer to the red line, i.e., the true partition). Table 7 shows the MSE of our method compared to topK summation, and the results are consistent with Figure 9.

heading: A. Dense distribution, entropy = 34.66 B. Intermediate, entropy = 18.08 C. Long-tail distribution, entropy = 6.49

Review summaries:
In this paper, the authors address the issue in the computational and memory efficiency in Dynamic Programming when it is used in the Forward Algorithm for Linear Chain CRF. The author purpose  a Sample Forward Algorithm by randomizing the summation. The authors proves that such randomization is an unbiased estimator for the true sums, and, using Amortized Variational inference for optimization, shows that it empirically reveals meaningful latent topology with much less required resource. 

Contextual embeddings of words are mapped to a sort of cluster of word embeddings with each one corresponding to one possible state in a linear CRF model. The idea is to build a network between the cluster of words (states of CRF) for analysis the manifold of language. To reduce quadratic cost of CRF, a sampling based approach is proposed to compute the partition function.   Empirical analysis is not very insightful, and the contributions are incremental. 

This paper proposes a new way to greatly reduce the memory requirements of inference on discrete latent variable models (LVMs). It does this by estimating the key inference steps in dynamic programming through enumerating a small subset of nodes and using importance sampling for the rest of the nodes. This technique is generally applicable to many dynamic programming algorithms and enables scaling LVMs to much more states by reducing memory requirements. Empirically, this work shows that the proposed approach gets better partition estimates than existing methods using top-k approximations. This work also demonstrates an application of the proposed approach on discovering meaningful latent codes and latent trajectories from pretrained language models.

This paper proposes a latent network analysis, using a proposed a Randomized Dynamic Programming (RDP) approach, to study the topology of the representation space of contextual embedding models. By constructing an unbiased estimator for the partition function with subsampled DP paths, the authors are able to reduce the memory complexity of forward-backward computation on Linear CRFs by 2 orders of magnitude. Using this RDP formulation, inference with structured latent variable models is scalable to 1000s of latent states. Latent states are modeled in a linear chain CRF with a long-tailed prior on transitions and learned using amortized variational inference similar to Fu et al. 2020, Li and Rush 2020. State-word connections in the representation space provide linguistic information that encapsulates syntactic and semantic roles and state-state connections correspond to the construction of phrases. The authors find anchoring states in the representation space that provide enough information to perform unsupervised paraphrase generation.