L ANGUAGE MODELING VIA STOCHASTIC PROCESSES
A BSTRACT
Modern language models can generate high-quality short texts.
However, they often meander or are incoherent when generating longer texts. These issues arise from the next-token-only language modeling objective.
To address these issues, we introduce Time Control (TC), a language model that implicitly plans via a latent stochastic process.
TC does this by learning a representation which maps the dynamics of how text changes in a document to the dynamics of a stochastic process of interest.
Using this representation, the language model can generate text by first implicitly generating a document plan via a stochastic process, and then generating text that is consistent with this latent plan.
Compared to domain-specific methods and fine-tuning GPT2 across a variety of text domains, TC improves performance on text infilling and discourse coherence.
On long text generation settings, TC preserves the text structure both in terms of ordering (up to +40% better) and text length consistency (up to +17% better).
Human evaluators also prefer TC’s output 28 .
6% more than the baselines.
I NTRODUCTION
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.
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 and the resulting latent trajectories 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 GPT 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.
R ELATED 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., 2019b) 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 (Bowman et al., 2016), however these generally result in poor language generation (He et al., 2018) 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.
The key difference is CPC implicitly learns unconditioned latent dynamics, whereas we impose known goal-conditioned dynamics on our latent space.
Doing so, we can extrapolate successfully further in time.
Additionally, 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.
and our work is that they focus on provable conditions for latent recovery; we focus on empirically effective methods that leverage similar insights for recovering latent representations from language.
Finally, our use of stochastic processes draws similarities to diffusion models (Song et al., 2020; Sohl-Dickstein et al., 2015) which apply a chain of diffusion steps onto the data and learn to reverse the diffusion process.
However, our application
[Figure]
is conceptually different: diffusion processes characterize properties of our latent space and are not a fixed inference method in our work.
T IME C ONTROL
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 sections. The first section discusses training the encoder via contrastive learning to map sentences to a Brownian bridge (Revuz & Yor, 2013) latent space.
The second section discusses training a decoder to reconstruct sentences from this latent space.
The third section discusses generating text from Time Control.
3. T RAINING AN ENCODER WITH B ROWNIAN 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 z 0 at t = 0 and end point z T at t =
T is,
[Equation]
This density is intuitive to understand: It acts like a noisy linear interpolation between the start and end point of the trajectory, where z t should be more like z 0 at the start and more like z T at the end of the trajectory.
Uncertainty is highest in the middle region, and low near the end points (rf.
Figure 1).
Consider a set of triplet observations, ( x 1 , x 2 , x 3 ) .
The goal of our work is to ensure that f θ ( x 1 ) , f θ ( x 2 ) , f θ ( x 3 ) follow the Brownian bridge transition density in Equation 1.
We ensure this using a contrastive objective.
Formally, given multiple sequences of data points, X = { x 1 , ..., x N } , we draw batches consisting of randomly sampled positive triplets x 0 , x t , x T where 0 < t < T :
B = { ( x 0 , x t , x T ) } .
Our encoder is optimized by,
[Equation]
  , where (2)
[Equation]
σ 2 is the variance in Equation 1: t ( T − t ) T .
Note that Equation 2 sums over negative middle contrasts, x t ′ .
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.
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.
; Poole et al.
(2019) (as |B| → ∞ ) and the classification setup from Liu et al.
(2021) ( |B| = 1 ).
Following van den Oord et al.
; 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 ( X t , { X 0 , X T } ) ≥ log( N ) − L N .
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.
They train a classifier to distinguish in- vs. out-of-order input pairs and show that the Bayes optimal logits for pair-wise classification can be written as a function of the stochastic process transition kernel. This is equivalent to the our loss on a single triplet i : l i = − log [Equation]
Liu et al.
(2021) consider pairs whereas our work considers triplets; we show in Appendix A the pairwise and triplet setups are equivalent.
3. T RAINING A DECODER WITH LATENT PLANS
Here we discuss how to train a language model to decode latent plans for generation.
We first map all the sentences in the training dataset to our learned latent space using the pretrained encoder f θ .
This gives us a Brownian bridge trajectory of sentence-level latent codes ( z 0 , . .
z t , . . .
, z T ) for a document in the dataset.
Then, rather than learning a decoder from scratch, we fine-tune GPT (Radford et al., 2019) to generate text conditioned on past context and the latent plan.
We fine-tune in the following manner.
Let x 1 . . .
x W be a document with W tokens and T sentences used to train the decoder.
Using the encoder f
θ , we can obtain embeddings z 1 . . .
z T for each sentence.
The decoder is a standard auto-regressive language model that is modified in the following way: at time t , the decoder must predict x t using all tokens in the past x <t , as well as the sentence embedding z s t , where the index s t ∈ [ T ] is a map which takes each token to its corresponding sentence.
This is a form of a reconstruction objective, as the identity of x t is encoded in z s t .
3. G ENERATING TEXT WITH LATENT PLANS AT INFERENCE TIME
Figure 2 illustrates how the trained decoder generates text at inference time.
Given two end points z 0 , z T , we sample a trajectory from a latent Brownian bridge, and then generate from the decoder conditioned on this bridge.
In many situations, we may not know the endpoints of the Brownian bridge explicitly.
In this case, we encode a set of sentences corresponding to start and end points (eg.
the first and last sentences of our training set), and fit a Gaussian to these points to form a density estimate.
Generating in this case involves first sampling from the Gaussian, and then generating as before from the bridge.
More details on training and generation can be found in Appendix B.
E XPERIMENTS
We now evaluate the ability of Time Control to capture text dynamics.
Specifically, we aim to answer the following research questions (RQ): zt [USER]
Hi, I’d like to order tickets for tomorrow.
[EOS] [ASSISTANT] Sure, what movie are these for?
[EOS]
T IME C ONTROL . . .
[USER] I’m thinking to watch Iron Man 3.
[EOS] zz. . . T IME C ONTROL T IME C ONTROL z T Generate latent plan pinned at and . zz T
Conditionally generate each sentence using the latent plan.
[Figure]
RQ1: Can Time Control model local text dynamics?
Section 4.1 investigates this question using a sentence ordering prediction task: given two sentences from the same document, we evaluate whether different models can predict their original order.
RQ2: Can Time Control generate locally coherent text?
Section 4.2 investigates this question using the text-infilling task: given prefix and suffix, we evaluate how well different models can fill in between.
RQ3: Can Time Control model global text dynamics?
Section 4.3 investigates this question on text generation for Wikipedia city articles by examining the length of generated sections.
RQ4: Can Time Control generate long coherent documents?
Section 4.4 investigates this question on forced long text generation: we evaluate how well models preserve global text statistics (such as typical section orders and lengths) when forced to extrapolate during generation.
We run Time Control with different latent dimensions ( d = 8 , 16 , 32 ).
Our encoder architecture is a frozen, pretrained GPT2 model from Huggingface (Radford et al., 2019; Wolf et al., 2020) and a trainable MLP network.
We extract GPT2’s last layer hidden state that corresponds to the end-ofsentence (EOS) token and train the 4-layer MLP on top of the hidden state.
The MLP network has intermediate ReLU activations and is trained with stochastic gradient descent with a learning rate of 1e-4 and with momentum 0.9.
Ablations We perform three ablations on our encoder model.
Recall that Time Control (A) explicitly models latent structure with (B) Brownian bridge dynamics using a (C) contrastive loss.
(A) replaces explicit dynamics with Implicit Dynamics (ID) where future latents are directly predicted with an autoregressive model (van den Oord et al., 2019).
(B) replaces Brownian bridge dynamics with Brownian motion (BM): latents follow the transition density z t | z s ∼ N ( z s , t − s ) in Equation 3.
Note z t is centered at z s and is not conditioned on a goal end-point.
(C) replaces the contrastive loss with a Variational Auto-Encoder (VAE) and centers the priors over z 0 to 0 and z T to 1, as done in our setup.
We use GPT2 as the decoder.
Appendix D includes more detail on the ablations.
Datasets We use language datasets that elicit different kinds of structure, from section structure to discourse structure to narrative structure.
Time Control does not take in any information about the structure, treating each domain the same under its encoding objective.
More information and dataset examples are provided in Appendix E. Wikisection (Arnold et al., 2019) includes Wikipedia articles on cities split by sections.
We adapt this dataset such that each article contains four ordered sections (abstract, history, geography, demographics) marked with section id tokens: Each article is represented as, “[ABSTRACT] text [HISTORY] text [GEOGRAPHY] text [DEMOGRAPHICS] text”.
Wikihow (WH) (Koupaee & Wang, 2018) contains how-to articles organized by a title, method, and steps.
We mark each with its own section id tokens: Each article is represented as “[TITLE] text [METHOD] text [STEP] 1 text [STEP] 2 text ...”
Recipe NLG (Bie´n et al., 2020) contains recipes, each with a title, ingredients and set of directions.
A recipe is constructed as “[TITLE] text [Table] [INGREDIENTS] text [DIRECTIONS] text”.
Taskmaster-
(TM-2) (Byrne et al., 2019) contains conversations on finding restaurants between an assistant and a user.
The assistant’s turn is marked with an “[ASSISTANT]” tag, and the user’s turn is marked with a “[USER]” tag.
TicketTalk
(Byrne et al., 2021) contains conversations on booking movie tickets between an assistant and a user.
The assistant’s and user’s turns are similarly marked as in TM-2.
ROC Stories (Mostafazadeh et al., 2016) is a short 5-sentence stories dataset.
No additional tokens are added in this dataset.
4. M ODELING LOCAL TEXT DYNAMICS
We evaluate how well Time Control models local text dynamics ( RQ1 ) on the discourse coherence setting (Jurafsky, 2000).
Discourse coherence is often measured by how well representations capture discourse structure by testing for whether a linear classifier can detect in-order and vs. out-of-order sentence pairs (Chen et al., 2019a).
We compare Time Control’s encoder against GPT2’s last layer’s hidden state corresponding to the EOS token
(Radford et al., 2019), BERT (Devlin et al., 2019), ALBERT (Lan et al., 2019), Sentence BERT (Reimers et al., 2019), and SimCSE (Gao et al., 2021).
The latter 4 methods are designed as sentence embedding models.
We also compare to our ablations.
The setup is the following:
The encoder takes in two sentences x t , x t + k to produce their latents z t , z t + k ∈ R d .
At random, the latents are fed either in- or out-of-order.
A linear classifier is trained on 100 epochs with stochastic gradient descent with a learning rate of 1e-4 and with momentum 0.9.
We varied the sentence distance k ∈ { 1 , 5 , 10 } and found that on some domains and for k = 1 , all methods scored near random accuracy; we have omitted those results in the main paper.
Otherwise, the results are summarized in Table 1 where we report the mean and standard error accuracy on a held-out discourse dataset of 3000 examples on 3 runs.
Our method is able to outperform or compete with sentence embedding specific methods, like ALBERT and SimCSE.
Additionally, though GPT is used as a base encoder for Time Control, we observe that Time Control greatly improves upon GPT2 with significant gains on TM-2, TicketTalk, and Wikisection (around 7 − 25% ).
Neither VAE nor BM perform better than random accuracy on any of domains.
This suggests that the variational lower bound does not recover the latent embeddings as well as contrastive objectives and the choice of stochastic processes matter for learning informative structural embeddings.
ID does best on the conversation datasets, which indicates that implicitly learning text dynamics yields latent representations that can be used for inferring discourse though not as effectively as Time Control.
[Table]
[Table]
[Table] [Table]
Noticeably, discourse is more easily inferred on task-oriented conversations like TM-2 and TicketTalk.
We hypothesize that the ordering of responses matters more in conversations than enumerative, factual settings like Wikisection.
Nonetheless, our model performs above chance on both types of domains.
This answers RQ1 in the positive: Time Control can model local text dynamics, like in conversations and articles.
4. G ENERATING LOCALLY COHERENT TEXT
We evaluate how well Time Control generates locally coherent text ( RQ2 ) on the text-infilling setting.
Text-infilling requires a model to take an incomplete text, with missing sentences, and complete it.
An example input could be, “Patty was excited about having her friends over.
[blank] Patty had a great time with her friends.”
The challenge in performing text infilling is to generate a sentence that is locally coherent with the left and right neighboring sentences.
We follow the task setup in Donahue et al.
(2020) where they use the ROCStories dataset.
Each story in the dataset contains 5 sentences, one of which is randomly masked.
Time Control fills in the masked sentence by running the bridge process pinned at the latent for the prefix, z 0 , and the latent for the suffix, z T .
We compare our method against ILM (Donahue et al., 2020), a state-of-the-art method for text-infilling as well as Donahue et al.
(2020)’s
LM model, an autoregressive baseline that only sees the prefix.
The VAE ablation takes the average of the prefix and suffix embeddings for generation.
When using the average, BM is equivalent to Time Control we omit the BM results.
ID’s autoregressive model predicts the next latent given the prefix for generation.
We evaluate the text coherence with the BLEU score (Papineni et al., 2002), ROUGE (Lin, 2004),
BLEURT (Sellam et al., 2020) and BERTScore (Zhang et al., 2019) between the generated and ground truth infill sentence.
Due to space constraints, the last three metrics are in the Appendix,
Table 17.
We also report human evaluations on how coherent the generated sentence is as a fill-in sentence.
Participants were asked to rank the generated fill-in sentence from ILM, LM, and Time
Control on a scale of 1-5 (not reasonable to very reasonable).
Appendix H includes more details on the human experiment setup.
The BLEU scores are summarized in Table 2.
Time Control generates fill-in sentences that much more closely overlap with the ground truth than ILM and LM.
The ablated methods perform worse to varying degrees.
On average, VAE performs worse than LM and ID.
This suggests that the interpolating embeddings learned via the variational objective yields embeddings which hurt the autoregressive model’s ability to generate locally coherent text.
[Table]
[Table] [Table]
[Table]
Human evaluations in Table 6 indicate that Time Control performs competitively with ILM.
4 Upon inspection of the model output, we find that the ablated models struggle to meaningfully decode from the interpolated latent embeddings (Appendix G.1).
Oftentimes, the model would start a sentence coherently, but end the sentence by repeating a word over and over again; examples can be found in Appendix G. This suggests the importance of learning a latent space that supports interpolation by construction.
All together, the results provide positive evidence for RQ2 : Time Control can generate locally coherent text due to its well-defined latent dynamics.
4. M ODELING GLOBAL TEXT DYNAMICS
We evaluate how well Time Control models global text dynamics ( RQ3 ) by assessing whether the methods mimic document structure on Wikisection. We check whether the generated section lengths match with the average lengths in the dataset.
We focus on Wikisection because it is the only dataset with long, well-defined sections (rf. Appendix E on dataset statistics).
Each document contains an abstract, history, geography, and demographics section on a city.
Time Control plans a latent trajectory by running the bridge process between the start and end latent, z 0 ∼ p ( z 0 ) and
z T ∼ p ( z T ) .
We compare Time Control to fine-tuned GPT2.
We also include two oracle methods that are fine-tuned GPT2 models with additional section-embedding supervision.
One is “sec-dense GPT2” (SD) where each token’s embedding contains the current section identity; section i ’s embedding is added onto the token’s positional token embedding.
The other is “secsparse GPT2” (SS) where the token embedding contains an indicator that is 1 if the token is the start of a new section, and 0 otherwise.
For VAE and ID, we calculate the density estimate p ( z 0 ) and p ( z T ) and run a linear interpolation between the start and end latents.
For BM, we calculate the density estimate p ( z 0 ) and run Brownian motion.
Table 3 reports in percentage how much the generated section lengths deviate from the average section lengths.
Time Control best matches the section lengths out of all the methods.
One surprising observation was that fine-tuned GPT2 mismatches the section lengths significantly by almost 20% .
Upon further investigation, we noticed GPT2 overshoots short sections and undershoots long sections; this causes the section length deviations.
Adding domain-specific supervision like with SD and SS slightly alleviates this issue.
Surprisingly, we find that our method beats the oracle on certain settings, eg.
d = 16 .
ID performs extremely poorly in contrast to VAE and BM; this highlights the challenge of interpolating on learned dynamics models which is exacerbated in the long text genera- [Table] [Table]
tion settings. The results affirm the importance of Time Control for modeling global text dynamics, such as in matching document structure, thus answering RQ3 positively
[Table]
4. G ENERATING GLOBALLY COHERENT TEXT
We evaluate how well Time Control generates globally coherent text ( RQ4 ) where the EOS token is omitted.
We refer to this as the forced long text generation setup because the model must extrapolate beyond its natural end point in generation.
Appendix ??
For reference, 1000 tokens is about 50% longer than the average Wikisection document (the longest text domain).
This setting is intrinsically difficult for auto-regressive models which do not have the ability to “look ahead” during generation (Lin et al., 2021).
We evaluate model extrapolation on three metrics.
First is ordering: how well do the models maintain the structure of the document (eg. turn-ordering in TicketTalk)?
Second is length consistency.
Length consistency captures common failure modes such as a model which stops modeling a conversation between two agents, and instead outputs a long monologue from one agent.
Third is human evaluations, which measures the long-tail text quality.
We examine the long-tail behavior as the complete 1000-token generation is long.
The results are summarized in Table 4 for length consistency, Table 5 for ordering and Table 7 for human evaluations.
Overall, our method maintains text flow the best according to these metrics.
A common failure mode of GPT2 is that it produces nonsensical text where it would naturally end generation; this is particularly noticeable on TM-2 and TicketTalk.
Figure 8 shows an example of our model’s behavior versus GPT2’s behavior on generating long texts in TicketTalk.
The example illustrates how Time Control continues the dialog whereas GPT2 utters nonsensical text.
These results are additionally confirmed by our human evaluation experiments: human evaluators rank Time
Control’s extrapolation ability better on all latent dimension settings than that of GPT2 (Table 7).
ID performs poorly in forced long text generation, similar to the normal text generation setting: It significantly misorders sections and overshoots section lengths.
Although VAE and BM do poorly on length consistency, they perform slightly better than ID on ordering; in Appendix J, we investigate why the VAE baseline performs much better than TC on ordering.
This suggests the importance of good latent dynamics for long text generation.
These results answer RQ4 positively: Time Control generates globally coherent text thanks to its ability to plan ahead and correctly generate future latent variables following its latent dynamics.
C ONCLUSION
We propose Time Control, a language model that implicitly plans via a latent stochastic process.
Specifically, Time Control looks to Brownian bridge processes as a desirable latent space.
Time Control learns to map coherent text to smooth Brownian bridge trajectories.
Empirically, we demonstrate this leads to several benefits such as generating more locally and globally coherent text.
Although our work focuses on benefits of learning stochastic process latents for language, it can be extended to other domains with sequential data like videos or audio, or extended to handle arbitrary bridge processes without known fixed start and end points.
R EPRODUCIBILITY S TATEMENT
In the supplemental we include a zip file containing our code and processed datasets.
We’ve also included in the appendix how we processed the datasets.
Our results on showing equivalence between the triplet classification setting and pairwise classification setting are also included in the appendix. A CKNOWLEDGMENTS
REW is supported by the National Science Foundation Graduate Research Fellowship.
The authors would give special thanks to CoColab members Mike Wu, Gabriel Poesia, Ali Malik and Alex
Tamkin for their support and incredibly helpful discussions.
The authors would like to thank the rest of CoCoLab and Rishi Bommasani for reviewing the paper draft. The authors also thank Chris Donahue for helpful pointers on using ILM. A S HOWING THE EQUIVALENCE OF THE PAIRWISE AND TRIPLET CLASSIFICATION TASK
In this section we focus on the classification setup where the goal is to train a classifier to detect invs.
out-of-order groups of inputs.
Specifically, we want to show that the classification setup from
Liu et al.
(2021) where the inputs come in pairs is equivalent to the setup where the inputs come in triplets.
The triplet inputs are what we assume for Time Control.
Notation We denote latent space as Z ; we typically do not observe this space directly.
We denote the observation space as X ; this is the space we observe directly.
We define the Brownian bridge for t ∈ [0 , T ] as
[Equation]
where W ( t ) is a standard Wiener process, N (0 , t ) .
We state the assumptions:
• We assume latents { z t } t ≥ 0 ∈ R
d are drawn from the Brownian bridge process defined by the stochastic differential equation,
[Equation]
The intervals at which latents are sampled are every ∆ t step: z ′ =
z t +∆ t .
• We denote the transition probabilities as
[Equation]
We denote the proposal distribution over possible intermediate triplets q ( z ′ 0 , z ′ t , z ′ T ) .
Liu et al.
(2021) characterize properties of an optimal classifier h ∗ ( z, z ′ ) which observes pairs of latents ( z, z ′ ) and it outputs in [0 , 1] , a probability indicating whether the pair comes in order (ie.
z ′ =
z t + ∆ t · dz t ) or not in order (ie.
a randomly sampled latent).
They train h ∗ using an L2 loss,
L ( h, { ( z, z ′ ) , y } ) .
Lemma 1 of their work states the following : The optimum of the contrastive learning objective arg min
[Equation] E (( z,z ′ ) ,y ) [ L ( h, { ( z, z ′ ) , y } )] satisfies
[Equation]
Manipulating this equality, we observe that the transition kernel has the following relation to the classifier which takes in pairs of observations,
[Equation]
Our setting however assumes that the algorithm receives triplets of data points, z t 1 ,
z
t 2 , z t 3 .
We want to show below that minimizing L with triplet contrasts in the classification setting still approximates the transition kernel.
In particular, we’re interested in transitions of a Brownian bridge pinned at z 0 , z T : p ∆ t ( z 0 , z t , z T ) =
Pr ( z t | z 0 , z T ) .
Let’s say we have two positive triplet samples, ( z it 1 , z it 2 , z it 3 ) and ( z jt 1 , z jt 2 , z jt 3 ) , where t 1 < t 2 < t 3 .
Following Liu et al.
[Equation] minimizing L yields the following on each triplet:
[Equation]
Taking the difference in log probabilities between Equations 10-11 results in log p
[Equation]
Similar to the pair-wise classification setting, we’ve shown that minimizing L in the triplet classification setting results in approximating the transition kernel of the Brownian bridge process.
B T RAINING DETAILS
Here we describe the training procedure for the encoder and the fine-tuning procedure for decoding.
Brownian bridge encoder The encoder architecture is a frozen, pretrained GPT2 model from
Huggingface (Radford et al., 2019; Wolf et al., 2020) and a trainable MLP network.
We extract the
GPT2’s last layer hidden state that corresponds to the end-of-sentence (EOS) token and train the 4-layer MLP on top of the hidden state.
The MLP network has intermediate ReLU activations and is trained with stochastic gradient descent with a learning rate of 1e-4 and with momentum 0.9.
We train the encoder for 100 epochs on each of the datasets.
The text fed into GPT2 are fed in on a sentence level.
This means that the input x t refers to the t ’th sentence of a document.
The sentences are separated from each other in the main text as “ . ” which is added to the tokenizer as a separate token for indexing convenience.
Fine-tuning GPT2 with latent embeddings
After training the encoder, we run it on the training dataset to collect an accompanying latent trajectory for each text.
The encoder is run on the dataset at a sentence level: we separate the text by sentences and pass the sentences through the encoder.
The sentence latent embeddings are aligned with the tokens of that sentence and offset by one token before the start of that sentence token.
Let’s illustrate by an example.
We denote [SOS] as the start of the document token, [s1] as sentence 1 tokens and [s2] as sentence 2 tokens.
[ . ] is the period token which we’ve added into the tokenizer.
z i denote the latent variable corresponding to the i ’th sentence.
Let’s say that the sequence fed into GPT2 is “[SOS] [s1] [s1] [s1] [ . ] [s2] [s2] [s2]”.
Then the corresponding latent trajectory is “ z 1 , z 1 , z 1 , z 1 , z 2 , z 2 , z 2 , z 2 ”.
The latent variables are added onto the positional embeddings.
We then fine-tune GPT2 as normal.
C G ENERATION WITH EMBEDDINGS C. N ORMAL TEXT GENERATION
We first sample a start and end latent, z 0 ∼ p ( z 0 ) , z T ∼ p ( z T ) where p ( z 0 ) , p ( z T ) are calculated as the density estimates over the training dataset.
We pin our trajectory to the start and end latent, and run the Brownian bridge using Equation ?? .
For normal long text generation, we set T to be the average number of sentences in each of the dataset.
For forced long text generation, we set T to be proportional to the number of sentences needed in order to generate 1000 tokens.
By the end, we have a trajectory z 0 , z 1 , ..., z T .
Generation starts with feeding the SOS token and the first latent z 0 .
Once GPT2 emits a [ . ] token and terminates sentence t , we transition to the next latent z t +1 .
This process continues until GPT2 is finished with generation.
If GPT2 generates more sentences than there are latents in the trajectory, the last latent z T is used until the end of generation.
C. F ORCED LONG TEXT GENERATION
Let S avg ( D ) denote the average number of sentences in a document and T avg ( D ) denote the average number of tokens in a document.
Rather than planning a trajectory of length S avg ( D ) (the average number of sentences in a document) which is what is done in Section 4.3 for normal text generation, we scale the trajectory length to c · S avg ( D ) .
c is determined by how many more tokens we need in order to fill up to GPT-2 maximum context length of 1024: c = 1024 − T avg ( D ) T avg ( D ) .
D A BLATIONS
The following methods ablate the encoder model in Time Control; in other words, the ablations dissect the assumptions we make in the first step of Time Control described in Section 3.1.
Recall that Time Control (A) explicitly models latent structure with (B) Brownian bridge dynamics using a (C) contrastive loss.
(A) replaces explicit dynamics with Implicit Dynamics (ID) where future latents are directly predicted with an autoregressive model (van den Oord et al., 2019).
(B) replaces
Brownian bridge dynamics with Brownian motion (BM) which doesn’t rely on pinning trajectories.
Latents follow the transition density z t | z s ∼ N ( z s , t − s ) .
(C) replaces the contrastive loss with the
Variational Auto-Encoder (VAE).
Below we detail how these ablations are implemented.
D. I MPLICIT D YNAMICS (ID)
The Implicit Dynamics ablation is where we compare our explicit dynamics objective to van den
Oord et al.
(2019) which suggests an implicit latent dynamics objective.
This ablation thus changes the learned latent dynamics.
In van den Oord et al.
(2019), they train two models. One is a non-linear encoder g enc ( x t ) = z t which takes observation x t (eg. a sentence) and maps it to a latent representation z t . Two is an autoregressive context model g ar ( z ≤ t ) =
c t which summarizes a sequence of latent variables z ≤ t into a context latent representation c t . Rather than directly predicting a future observation x t + k ( k steps from the current timestep t ), they model the density ratio that preserves the mutual information between x t + k and the context variable c t : [Equation] They model this with a log-bilinear model f k ( x t + k ) =
exp( z Tt + k W k c t ) which applies a linear transformation W k for modelling latents k-steps away from timestep t .
This way, they avoid directly learning the generative model p ( x t + k | c t ) .
They train both models jointly via a contrastive InfoNCE loss.
Given a set X = { x 1 , . . .
, x N } of
N random samples (eg. sentences from different documents) containing one positive sample from p ( x t + k | c t ) and N − 1 negative samples from a proposal distribution p ( x t + k ) , they optimize:
[Equation]
The encoder g enc for Implicit Dynamics has the same architecture as Time Control.
The context encoder g ar using a 2400-hidden-unit GRU, as done in van den Oord et al.
We then train both encoders using the InfoNCE loss, as done in prior work.
Since g enc and g ar are trained to align latents up to a linear rotation, we use g enc for extracting the sentence embeddings.
D. B ROWNIAN MOTION (BM)
The Brownian motion ablation is where we remove goal-conditioning in the dynamics.
The main change is in distance function in Equation 3.
BM instead optimizes the following equation:
[Equation]
where σ 2 is the variance in of the Wiener process, σ 2 = t ′ − t .
D. V ARIATIONAL A UTO -E NCODER (VAE)
The Variational Auto-Encoder ablation is where we compare our contrastive objective with a variational one.
This ablation changes the encoder objective from Section 3.1 from Equation 2 to the
ELBO objective.
Below we derive the ELBO objective.
Similar to the contrastive objective, we’re deriving the ELBO over the triplet dynamics in the Brownian bridge.
[Equation]
We assume the priors over z 0 is 0-centered and
z T is 1-centered, which is similar to our Brownian
Bridge setup.
The encoder q ϕ ( z | x ) is parameterized with the same architecture as our encoder.
The decoder p ( x | z ) is a fine-tuned GPT2 model [Equation]
E D ATASET INFORMATION
For each dataset, text examples were filtered out if they did not fit within GPT2’s context length of 1024 tokens.
We also added the token “ . ”
for each setting to mark the end of a sentence.
This was done for indexing purposes, eg.
when aligning the latent embeddings.
Wikisection (Arnold et al., 2019) includes Wikipedia articles on cities split by sections.
We adapt this dataset such that each article contains four ordered sections (abstract, history, geography, demographics) marked with section id tokens: Each article is represented as, “[ABSTRACT] text [HISTORY] text [GEOGRAPHY] text [DEMOGRAPHICS] text”.
These section id tokens are added to the tokenizer.
The training dataset contains 1420 articles.
The section lengths have the following breakdown measured in the number of BPE tokens (GPT2 tokenizer):
• Abstract: 75 . 8 ± 1 .•
History: 191 . 5 ± 3 .•
Geography: 83 . 9 ± 1 .• Demographics: 342 .
6 ± 4 .
The test dataset contains 431 articles.
The section lengths have a similar breakdown:
• Abstract: 73 .
5 ± 2 .•
History: 180 . 2 ± 6 .•
Geography: 85 .
2 ± 2 .•
Demographics: 332 . 5 ± 8 .
The ordering metric used in Table 5 is 1 if all four section ids occur exactly once and come in the order as they are listed above.
The length mismatch in % used in Table 3 is calculated with respect to the training set lengths.
Wikihow (Koupaee & Wang, 2018) contains how-to articles organized by a title, method, and steps.
Each article includes multiple methods for completing a multi-step procedural task such as “How to Register to Vote” .
We scraped all the available English articles covering a wide range of represented
The training dataset contains 1566 articles.
The section lengths are,
• Title: 10 . 4 ± 2 .•
Method: 8 . 7 ± 2 .•
Steps (total step length): 480 . 2 ± 231 .
The test dataset contains 243 articles.
The section lengths are,
• Title: 10 .
7 ± 2 .•
Method: 8 . 6 ± 2 .•
Steps (total step length): 480 .
1 ± 224 .
The ordering metric used in Table 5 is 1 if all the section ids appear in order, the TITLE and
METHOD section ids are not repeated, and the step numbers come in order.
It’s 0 otherwise.
The deviation in section length measured in Table 4 is calculated with respect to the training set lengths.
We check for whether the models are able to maintain the section lengths when it has to extrapolate.
The most common failure mode is that the model generates incoherent text and allocates this text to the last section of what it’s generated thus far, resulting in the deviation in section lengths.
Recipe NLG (Bie´n et al., 2020) contains recipes, each with a title, ingredients and set of instructions.
A recipe is constructed as “[TITLE] text [INGREDIENTS] text [DIRECTIONS] text”.
The training dataset contains 4000 recipes.
The section lengths are,
Title: 9 . 7 ± 3 .•
Ingredients: 23 .
8 ± 4 .•
Directions (total step length): 62 . 0 ± 14 .
The test dataset contains 1000 recipes.
The section lengths are,
Title: 4 ± 3 .•
Ingredients: 24 . 1 ± 4 .•
Directions (total step length): 63 .
3 ± 13 .
The ordering metric used in Table 5 is 1 if all the section ids appear exactly once and in order.
It’sotherwise.
Taskmaster-2 (TM-2) (Byrne et al., 2019) which contains conversations user’s
The training dataset contains 2000 conversations.
The section lengths are,
• User: 11 . 8 ± 5 .•
Assistant: 18 . 0 ± 3 .
The test dataset contains 1276 conversations.
The section lengths are,
• User: 11 .
7 ± 5 .•
Assistant: 19 .
5 ± 3 .
The ordering metric used in Table 5 does not apply here because the user and assistant don’t take turns in the dialog. [Table]
The deviation in section length measured in Table 4 is calculated with respect to the training set lengths.
We check for whether the models are able to maintain the utterance lengths between the user and assistant when it has to extrapolate.
The most common failure mode is that the model generates incoherent text and allocates this text to the last section of what it’s generated thus far, resulting in the deviation in section lengths.
TicketTalk (Byrne et al., 2021) which contains conversations on booking movie tickets between an assistant and a user.
The assistant’s and user’s turned are similarly marked as in TM-2.
The training dataset contains 2000 conversations.
The section lengths are,
• User: 11 . 8 ± 5 .•
Assistant: 18 . 0 ± 3 .
The test dataset contains 1276 conversations.
The section lengths are,
• User: 11 .
7 ± 5 .•
Assistant: 19 .
5 ± 3 .
The deviation in section length measured in Table 4 is calculated similarly to TM-2.
The ordering metric used in Table 5 applies because the user and assistant take turns in the dialog.
The ordering metric is 1 if the user and assistant take turns in the conversation.
It’s 0 otherwise.
ROC Stories (Mostafazadeh et al., 2016) is a short stories dataset.
Each story contains 5 sentences.
No additional tokens are added in this dataset.
The training dataset contains 2000 stories, and the test dataset contains 1000 stories.
F P ERPLEXITY AFTER FINE - TUNING
Table 9 reports the final perplexity scores after fine-tuning GPT2 on the different domains with the methods.
We fine-tune for 10 epochs and checkpoint the models every 1000 steps; we keep the model checkpoint that scores the lowest PPL on a held-out validation set.
G G ENERATION EXAMPLES
We provide example outputs from the different models on text infilling.
We denote the prefix sentences as “LEFT CONTEXT”, the suffix sentences as “RIGHT SENTENCES”, the generated output from the models as “GENERATED” and the ground truth infill sentence as “GT”.
G. T EXT INFILLING
LM See Table 10.
[Table]
ILM See Table 11.
[Table]
Time Control See Table 12.
[Table]
LEFT CONTEXT:
Tom was jealous of his brother.
His brother was older and stronger.
His brother went wherever he wanted.
GENERATED: Tom’s brother was very jealous of his brother.
GT:
One day his brother was grounded for staying out too late.
RIGHT CONTEXT:
This made Tom really happy.
LEFT CONTEXT:
Jackie was 11 and had to get braces.
She was worried about what her friends would think.
She tried to hide them when she first got them.
GENERATED:
But she was too embarrassed to tell them .
GT:
Eventually her friends saw them and she was embarrassed.
RIGHT CONTEXT:
Her friends noticed she was embarrassed and decided to comfort her.
LEFT CONTEXT:
Sally was going to surprise the office with a cake.
Sally felt that a cake would be a good way to make them smile.
She went to the supermarket to pick up the cake.
GENERATED: She bought the cake and was very happy with it.
GT:
At the office she gathered the employees around the conference table.
RIGHT CONTEXT:
She then brought out a cake and they all felt better.
LEFT CONTEXT:
Lars was playing XBOX.
His controller stopped working during a game.
Lars didn’t have a car so he had to walk all the way to the store.
The store was being remodeled when he got there so he went to another.
GENERATED: Video games were all over the floor.
GT:
Lars wasn’t able to find a controller.
RIGHT CONTEXT:
LEFT CONTEXT:
Emma had been working as a dishwasher.
Her hands cracked and bled from the hot soapy water.
Then her mom noticed and concocted a special salve for her.
Emma used the salve every night before bed.
GENERATED: Soon she was able to keep her hands clean and happy.
GT:
Her hands got better.
RIGHT CONTEXT:
LEFT CONTEXT:
Jerome dribbled the ball quickly.
Sam tried to grab it from him, but wasn’t fast enough.
He chased Jerome down the court.
GENERATED:
He scored a point with a long shot .
GT:
Jerome pushed Sam backwards and threw the ball.
RIGHT CONTEXT:
Jerome scored points for his team.
LEFT CONTEXT:
Class started in 10 minutes and I had a math assignment due.
My older friend volunteered to do it for me.
Unfortunately, my teacher found out about the copying.
GENERATED: She was very upset with me for doing the assignment.
GT:
She crumbled my paper and threw it in the trash.
RIGHT CONTEXT:
My teacher gave me an F for the assignment.
LEFT CONTEXT:
I went to the park to play frisbee with my dog.
GENERATED: We played all day .
GT:
I tossed the frisbee to my dog and he would catch it in his mouth.
RIGHT CONTEXT:
I accidentally threw the frisbee too far.
The frisbee landed into the pond.
But my dog went and got it.
LEFT CONTEXT:
I was tired of working at Walmart.
The hours were bad.
The store music was awful.
GENERATED: I was very bored .
GT:
I handed my two weeks in to the manager.
RIGHT CONTEXT:
I then found another job and was happy.
LEFT CONTEXT:
Sam bought a new SUV.
It was all wheel drive.
He figured he would take it off road.
He hit a few hard bumps and broke his suspension.
GENERATED: Unfortunately he had to pay a lot of money for it.
GT:
Sheepishly, he brought it to the dealership for repair.
RIGHT CONTEXT:
LEFT CONTEXT:
Missy got drunk and went to get a tattoo.
She decided to get a tattoo on her forehead.
The next day, Missy was horrified at what she had done.
GENERATED: Her tattoo was on her forehead!
GT:
Missy scraped up her money to pay for a tattoo removal procedure.
RIGHT CONTEXT:
After much wasted money, the tattoo was gone.
LEFT CONTEXT:
Jake was going on a road trip to see his family.
He Got in the car and drove.
GENERATED: Jake’s family was driving down the road .
GT:
The car’s tires exploded due to too much air.
RIGHT CONTEXT:
Jake hitchhiked for 30 miles.
When Jake got to his family he was happy his trip was over.
LEFT CONTEXT:
Bob decided to start a business.
GENERATED:
However, he did not know the market at all .
GT:
He opened up a grocery store and was doing very well.
RIGHT CONTEXT:
After a year, his profits dropped and he had to declare bankruptcy.
Bob was sad to see his business fail.
Bob worked hard and reopened his business.
LEFT CONTEXT:
I hadn’t seen my girlfriend in a while.
She got a new job so it’s hard to talk.
The job takes up all of her time.
GENERATED: I had to ask her to go out with me .
GT:
Finally she called me to hang out.
[Table]
[Table] [Table]
G. F ORCED LONG TEXT GENERATION
We include examples of GPT2 long text generation in Figures 3 and 4.
H H UMAN EXPERIMENTAL SETUP
We use Mechanical Turk to collect human evaluation for long text generation and text infilling models.
Examples from all the systems are selected randomly and each example is evaluated byannotators.
H. L ONG T EXT G ENERATION E VALUATION
The annotators are provided with a dialogue between a movie assistant and a user generated by GPT and Time Control, where the goal of the assistant is to help the user to book a movie ticket.
We show the last few generated interactions between the assistant and the user to the annotators, and ask them to assess whether the provided snippet could be a reasonable ending for such a conversation.
We use a 5-point Likert scale where 1 corresponds to a “Very Unreasonable Ending” and 5 corresponds to a “Very Reasonable Ending”.
H. T EXT I NFILLING E VALUATION
Given the preceding and following context, we ask annotators to evaluate whether the generated fill-in sentence is a reasonable completion.
We evaluate generated fill-in sentence by ILM, LM and
Time Control on a scale of 1-5.
I E XAMPLE TRAJECTORIES OF THE LEARNED LATENT SPACE
See Figures 7-10 for example latent trajectories over coherent vs. incoherent (randomly scrambled) held-out Wikisection documents.
These trajectories show the recovered latent structure by Time
Control and the three ablations from our work.
Figure 7 are latent trajectories from Time Control.
Figure 8 are latent trajectories from Implicit Dynamics.
Figure 9 are latent trajectories from
Brownian motion.
Figure 10 are latent trajectories from Variational Auto-Encoder.
J W HY DOES V ARIATIONAL A UTO -E NCODER PERFORM WELL ON THE R ECIPE DOMAIN ?
In Table 5, VAE does much better than TC.
Here we investigate potential reasons why.
We looked into comparing the latent structure recovered by the VAE baseline and TC; see Figure 11.
What we found is the Time Control is best at extracting time over the course of the document (ie.
its latents recover a Bridge process correlated with time), whereas the VAE doesn’t elicit strong temporal structure.
We hypothesize that temporal structure is not all you need to succeed in the
Recipe domain.
K A DDITIONAL EXPERIMENTS
On discourse coherence, we ran ALBERT on k = 1 since ALBERT is trained on this setting.
We found that it got slightly above random chance, unlike the most of the baselines we compare against.
The discourse accuracy is 50 .
1 ± 9 . 0 on Wikisection, 55 .
8 ± 4 . 0 on TicketTalk and 60 .
8 ± 2 . 8 on
TM-2.
L P REVIOUS EXPLORATORY THREADS L. S UCCESS AND FAILURE MODES IN FINE - TUNING GPT
We found that fine-tuned GPT2 was able to replicate certain aspects of a document corpus, such as section header ordering ( 92% accurate).
However, when going from document-level to sectionlevel statistics, we noticed that GPT-2 seemed to be either undershooting or overshooting the section lengths.
These results are reported in Table 3.
L. I NVESTIGATING THE IMPORTANCE OF DISTANCES BETWEEN SAMPLED SENTENCES
In earlier iterations of the work, we did explore learning with pairwise contrasts with fixed tdistances, e g. distances of length 1, 5, and 10 sentences.
We observed that the resulting latent trajectories elicited different fits to Brownian bridge dynamics, and the quality in fit varied in t across domains; we’ve include some examples in Figure 12 on Wikisection.
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
[Table]