Language modeling via stochastic processes

1 INTRODUCTION . Large language models ( LLM ) such as GPT-2 have been extremely successful in text generation ( Radford et al. , 2019 ; Brown et al. , 2020 ) . However , LLMs are known to generate incoherent long texts . One reason is that they are unable to plan ahead or represent long-range text dynamics ( Kiddon et al. , 2016 ; Fan et al. , 2019 ; Hua & Wang , 2020 ; Duboue & McKeown , 2001 ; Stent et al. , 2004 ; Tamkin et al. , 2020 ) . As a result , they oftentimes produce wandering content with poor discourse structure and low relevance ( Hua & Wang , 2020 ; Zhao et al. , 2017 ; Xu et al. , 2020 ) —the text reads as if the model has no anchored goal when generating . These problems with coherence are further exacerbated when forcing autoregressive models to generate longer texts as the model struggles to extrapolate beyond its expected text end point . These problems suggest that LLMs currently fail to properly capture how documents evolve from beginning to end and condition generation on goals . Doing so is critical for succeeding in goal-oriented tasks such as story , dialog or recipe generation . Prior work has explored the use of planning-based methods for generating globally coherent text ( Kiddon et al. , 2016 ; Fan et al. , 2019 ; Hua & Wang , 2020 ; Duboue & McKeown , 2001 ; Stent et al. , 2004 ) . However , these methods rely on manually defining text dynamics for specific domains . Other work has attempted to use sentence representations for modeling text , such as with variational autoencoders ( Bowman et al. , 2016 ) or contrastive learning ( Gao et al. , 2021 ; Devlin et al. , 2019 ) . Their shortcoming in text generation settings is that the latent representations are static : they capture semantic similarity between sentence neighbors , but don ’ t capture how sentence embeddings evolve over a document . Methods including van den Oord et al . ( 2019 ) have tried to remedy this by learning a model of local latent dynamics . However , it is difficult to use learned local dynamics for generating accurate goal-conditioned trajectories , especially long-horizon ones . We explore an alternative that explicitly assumes a simple , fixed dynamics model with goal-conditioned generation . In this work , we propose Time Control as a way to learn a latent space with known , goal-conditioned dynamics . We begin by assuming that meandering text generated without a goal can be represented as Brownian motion in latent space ; this motion enforces the embeddings of neighboring sentences to be similar to each other , whereas those of distant sentences to be dissimilar . Goal-directed behavior can be incorporated into this model by conditioning on a fixed start and end point . In this case , the Brownian motion becomes a Brownian bridge . The resulting latent trajectories are controllable and abide by simple , closed-form dynamics . In Time Control , we derive a novel contrastive objective for learning a latent space with Brownian bridge dynamics . We can then use this latent space to generate text that retains local coherence and has improved global coherence . To perform text generation , Time Control first plans a latent trajectory via the Brownian bridge process pinned at a start and end point . It then conditionally generates sentences using this latent plan . In our work , we decode latent plans by fine-tuning GPT2 to generate text conditioned on Time Control ’ s latent trajectory . Trajectories from Time Control act as abstract semantic positions in a document that guide generation of fine-tuned language models . In summary , our work ’ s contributions are the following : • We derive Time Control , a language model which explicitly models latent structure with Brownian bridge dynamics learned using a novel contrastive objective . • Across a range of text domains , we show that Time Control generates more or equally coherent text on tasks including text infilling and forced long text generation , compared to task-specific methods . • We validate that our latent representations capture text dynamics competitively by evaluating discourse coherence with human experiments . • We ablate our method to understand the importance of the contrastive objective , enforcing Brownian bridge dynamics , and explicitly modeling latent dynamics . 2 RELATED WORKS . Generating long , coherent text is conceptually difficult for autoregressive models because they lack the ability to model text structure and dynamics ( Lin et al. , 2021 ) . This means that they struggle to plan and look ahead which leads to generating globally incoherent text . Forcing autoregressive models to generate longer texts exacerbates this incoherence because the models struggle to extrapolate beyond their expected text end point . Prior work has tried to address the problem of generating globally coherent text with planning-based approaches ( Puduppully et al. , 2019 ; Moryossef et al. , 2019 ; Fan et al. , 2019 ; Kiddon et al. , 2016 ) . However , planning-based approaches rely on domain-specific heuristics for capturing text structure and dynamics . Our work uses a contrastive objective to learn latent dynamics in text without domain-specific heuristics . Contrastive objectives have been applied to several domains , including language ( Devlin et al. , 2019 ; Iter et al. , 2020 ; Liu & Liu , 2021 ) , vision ( Chen et al. , 2020 ) , and general time series data ( Hyvarinen & Morioka , 2016 ; Hyvarinen et al. , 2019 ) . In particular for language , contrastive objectives have been applied to the next-sentence prediction task for improving BERT embeddings ( Devlin et al. , 2019 ) and to the discourse coherence setting ( Nie et al. , 2019 ; Chen et al. , 2019 ) for evaluating how coherent pairs of sentences are . However , these methods have two shortcomings which we address with our work . One is that the resulting sentence embeddings are often static : they capture semantic similarity between sentence neighbors , but don ’ t capture how sentence embeddings evolve over a document . Two is that they are not used for generation and are limited to classification tasks like discourse coherence . Prior work has also tried fitting latent variable models , however these generally result in poor language generation ( Bowman et al. , 2016 ) or are domain-specific ( Weber et al. , 2020 ; Arora et al. , 2016 ) . Our work is closely related to Contrastive Predictive Coding ( CPC ) from van den Oord et al . ( 2019 ) which learns the implicit latent dynamics with an autoregressive model . CPC is different from our work in that we impose constraints with known dynamics onto our latent space and assume goalconditioned generation , whereas they learn an unconditioned dynamics model . By assuming that text dynamics follows a well-behaved stochastic process , we benefit significantly during generation . Generating from future time points with CPC results in extrapolation error due to compounding errors from its forward model . In contrast , our approach has an exact , closed-form forward model which allows us to extrapolate substantially further in time . Our method builds off of recent findings that contrastive objectives can be used to approximate local transition kernels of stochastic processes ( Liu et al. , 2021 ) . The main difference between Liu et al . ( 2021 ) and our work is that they focus on providing provable conditions for latent recovery , whereas we focus on building empirically effective methods that leverage similar insights for recovering latent representations from language . 3 TIME CONTROL . The intuition behind Time Control is to learn a latent space with smooth temporal dynamics for modeling and generating coherent text . We detail Time Control in three steps : The first step trains an encoder to map sentences to a Brownian bridge ( Revuz & Yor , 2013 ) latent space . The second step trains a decoder to reconstruct sentences from this latent space . The third step generates text by first sampling a latent plan , then generating using a decoder on this latent plan . 3.1 TRAINING AN ENCODER WITH BROWNIAN BRIDGE DYNAMICS . Our encoder is a nonlinear mapping from raw input space to latent space , fθ : X → Z . The objective for the encoder is to map high-dimensional sequential data into low-dimensional latents which follow a stochastic process of interest—in this paper , it is the Brownian bridge process . The density of a Brownian bridge process between an arbitrary start point z0 at t = 0 and end point zT at t = T is , p ( zt|z0 , zT ) = N ( ( 1− t T ) z0 + t T zT , t ( T − t ) T ) . ( 1 ) This density is intuitive to understand : It acts like a noisy linear interpolation between the start and end point of the trajectory , where zt should be more like z0 at the start and more like zT at the end of the trajectory . Uncertainty is highest in the middle region , and low near the end points . See Figure 1 . Consider a set of triplet observations , ( x1 , x2 , x3 ) . The goal of our work is to ensure that fθ ( x1 ) , fθ ( x2 ) , fθ ( x3 ) follow the Brownian bridge transition density in Equation 1 . We ensure this using a contrastive objective . Formally , given multiple sequences of data points , X = { x1 , ... , xN } , we draw batches consisting of randomly sampled positive triplets x0 , xt , xT where 0 < t < T : B = { ( x0 , xt , xT ) } .1 Our encoder is optimized by , LN = EX [ − log exp ( d ( x0 , xt , xT ; fθ ) ) ∑ t′∈B exp ( d ( x0 , xt′ , xT ; fθ ) ) ] , where ( 2 ) d ( x0 , xt , xT ; fθ ) = − 1 2σ2 ‖ fθ ( xt ) ︸ ︷︷ ︸ zt − ( 1− t T ) fθ ( x0 ) − t T fθ ( xT ) ︸ ︷︷ ︸ mean in Equation 1 ‖22 ( 3 ) σ2 is the variance in Equation 1 : t ( T−t ) T . Note that Equation 2 sums over negative middle contrasts , xt′ . This objective can be viewed as maximizing the extent to which true triplets from the data follow the Brownian bridge process while minimizing the extent to which an alternative mid-point sampled from another sequence does so . 2 1We use indices 0 , t , T to denote the start , middle and end point of a Brownian bridge , but these do not correspond to strictly sampling the first , middle and last sentence of a document . x0 , xt , xT can be any sentence in a document as long as x0 comes before xt and xt before xT in the document . 2Empirically , we found Brownian bridge dynamics easier to recover with triplets rather than pairs of contrasts . Appendix M.2 discusses some of the pair-wise contrast results . Figure 1 illustrates how the objective translates into the language setting for training the encoder . The objective samples triplet sentences from a document . Sentences drawn from the same document make up a smooth latent trajectory ; they should be close to each other and follow the conditional density in latent space . Sentences drawn from different documents should not make up a smooth trajectory and should less likely follow bridge dynamics . Connection to mutual information estimation and triplet classification We draw connections between our contrastive loss and the mutual information estimation setup from van den Oord et al . ( 2019 ) ; Poole et al . ( 2019 ) ( as |B| → ∞ ) as well as the classification setup from Liu et al . ( 2021 ) ( |B| = 1 ) . Following van den Oord et al . ( 2019 ) ; Poole et al . ( 2019 ) , this objective can be seen as a lower bound on the mutual information between the two end points and the middle point : I ( Xt , { X0 , XT } ) ≥ log ( N ) − LN . Hence , by minimizing the contrastive loss , we are maximizing the amount of information between the trajectory and the linear interpolation of its end points . Assuming |B| = 1 , we can draw a connection to the classification setup studied in Liu et al . ( 2021 ) where they train a classifier to distinguish in- vs. out-of-order pairs of inputs . They show that the Bayes optimal logits for pair-wise classification can be written as a function of the stochastic process transition kernel , which is equivalent to the loss on a single triplet i in our contrastive objective : li = − log exp ( d ( x0 , xt , xT ; fθ ) ) exp ( d ( x0 , xt , xT ; fθ ) ) +exp ( d ( x0 , xt′ , xT ; fθ ) ) . We show in Appendix A that the pairwise classification setup is equivalent to the triplet classification setup .

Summary: This paper proposes to model the evolution of sentences in a document via a stochastic process; specifically a Brownian Bridge process. The paper start off by assuming that the generated sequences by autoregressive models like GPT-2 follow Brownian motion in that they tend to get incoherent and "meander" in the semantic space. This paper aims to reduce this random behavior by pinning the endpoints of the trajectory and model the generation by Brownian bridge process instead. The key intuition behind this process is that given two endpoints z0 and zT, the evolution of z along time t is a Gaussian with mean that is some linear combination of z0 and zT. This paper models text by training an encoder for sentences x that produces the embedding z by training over triplets (x0, xt, xT) where 0<t<T that encourages zt to follow Brownian bridge dynamics and uses contrastive loss with a negatively sampled x't for training.  The approach is tested for local coherence, long range order sensitivity, and generation of long sequences and is compared against ablative and external baselines. The proposed approach does lead to learning of embeddings that are obtainable via linear combination and this leads to improved performance on sensitivity to sentence order in documents and document generation.