Discovering Latent Network Topology in Contextualized Representations with Randomized Dynamic Programming

1 INTRODUCTION . 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 can not 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 . 2 RANDOMIZED DYNAMIC PROGRAMMING . 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 ) .

Summary: 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.