L anguage F low: Advancing Diffusion Language Generation with Probabilistic Flows (2024)

Abstract: Recent works have demonstrated success in controlling sentence attributes (e.g., sentiment) and structure (e.g., syntactic structure) based on the diffusion language model. A key component that drives theimpressive performance for generating high-quality samples from noise is iteratively denoise for thousands of steps. While beneficial, the complexity of starting from the noise and the learning steps has limited its implementation to many NLP real-world applications. This paper proposes Language Rectified Flow (LF).Our method is based on the reformulation of the standard probabilistic flow models.Language rectified flow learns (neural) ordinary differentialequation models to transport between the source distribution and the target distribution, henceproviding a unified and effective solution to generative modeling and domain transfer.From the source distribution, our language rectified flow yields fast simulation and effectively decreases the inference time. Experiments on three challenging fine-grained control tasks and multiple high-quality text editing show that our method consistently outperforms its baselines. Extensive experiments and ablation studies demonstrate that our method can be general, effective, and beneficial for many NLP tasks.

## introduction traditional pretrained large-scale language models (lm) can generate high-quality text for specific real-world applications @xcite @xcite @xcite . however, updating the lm parameters or finding proper prompts for each control task can be expensive and unscalable given the combinatorially many possible compositions and the lack of supervised data. recent research thus has started to explore plugand-play solutions. with a pretrained language model (lm), the plug-and-play approaches @xcite @xcite @xcite are served as the light-weight constrained guidance to the targeted text sequence generations @xcite . the approaches, however, typically rely on search or optimization in the complex text sequence space. the discrete nature of text makes the search/optimization extremely difficult. though some recent work introduces continuous approximations to the discrete tokens @xcite , the high dimensionality and complexity of the sequence space still render it inefficient to find accurate high quality text. the most recent approach diffusion lm @xcite induces continuous latent representations, which enables efficient gradient-based methods for the controllable generation. however, despite the advantage of diffusion lm for separating the distribution map learning from a noise distribution to a meaningful shape distribution, it always requires thousands of steps. the transport trajectory learns from a noise distribution to a meaningful shape the distribution. this is a major efficiency bottleneck during inference since a standard diffusion model requires thousands of generation steps and a proper sde solver to produce high-reconstruction and diverse text. in addition, the existing denoising diffusion techniques requires substantial hyper-parameter search in involved design space. we propose language rectified flow, a surprisingly simple approach to the transport mapping problem, which unifiedly solves both generative modeling and domain transfer. the language rectified flow is an ordinary differential equation (ode) model that transports distribution source distribution π 0 to target distribution π 1 by following straight line paths as much as possible. the straight paths are theoretically preferred because it is the shortest path between two end points, and computationally preferred because it can be precisely simulated without time discretization. hence, flows with straight paths bridge the gap between few-step and continuous-time models. specifically, we formulate an ode transport flow as the initial text generator with a simpler trajectory compared with the diffusion model formulated in sde. meanwhile, we optimize the transport flow cost for the initial flow model to significantly straighten the learning trajectory while maintaining the model's performance. further, a vae is utilized to connect the text sequence space and the flow latent space. this can be used to produce the initial values in the language flow for our controllable text generation. to demonstrate control of language rectified flow, we consider three control targets ranging from span-anchored controls (e.g., length control and infilling) to complex structures (e.g., parse trees). our lf can generate high-quality text, performing favorably relative to the diffusion-based lm with a 27x faster sampling on the controllable text generation task. in addition to these individual control tasks, we conduct experiments on three challenging settings, including sequential editing of text. results show that composing operators within our method manage to generate or edit high-quality text, substantially improving over respective baselines in terms of quality and efficiency. furthermore, we provide extensive ablation studies on different design choices for the proposed method, including the evidence with different generation steps, influence of the constrained optimization, flow in the latent space. our analysis shows that lf contributes the performance improvement, helping the sampling efficiency and generation. with little modification, ours can be easily applied to other nlp tasks for better controllable text generation. our contributions are summarized as follows: ## method we now introduce our method, language rectified flow (lf), the transport flow for the controlled language generation. generating text with transport flow can be viewed as transporting from source distributions to target distributions by following a learned trajectory. specifically, we suggest a general recipe for the language rectified flow: 1) construct the continuous latent space, 2) learn a neural velocity flow network, construct an ode process in the latent space with the shortest transport path, and utilize a neural network to fit this process, 3) propose the lexicographic (lexico) optimization strategy for the joint learning. in this section, we present in detail our proposed method, lf (algorithm 1). ## experimental settings we evaluate our method on control tasks and text editing tasks. ## experiments we evaluate the performance of our language rectified flow and learning framework in this section. we bold the best result within each column block. the results of our method are obtained with three independent runs to determine the variance. see appendix a for full results with error bars. ## analysis what is the influence of the constrained optimization vs. manually tuning coefficient? here we verify whether the constrained optimization in lf is better than the manually tuning strategies. with the designed trade-off, lf targets optimizing between the flow and representation construction. instead of automatically searching the trade-off between the the flow optimization and representation construction, we manually set a constant λ in eqn (7) as 0.1, 1.0, 2.0. ❶ ablations on the components of language rectified flow. our language rectified flow focuses on latent diffusion. the latent space is constructed by vae. thereforce, the purpose of the latent flow is to learn the transport from the gaussian distribution to the gaussian distribution. under this setting, unet @xcite and transformer @xcite are two eligible considerations. following the diffusion language model hyperparameter and code base @xcite , we obtain the below result ( studies on the role of latent lf. we conduct the ablation study to exam the role of latent lf in the latent space. with the designed flow transport, lf targets the fine-grained text generation. we compare lf with and w/o latent lf settings. w/o latent lf here represents mapping the sentence in one domain to the latent space and directly mapping this latent code into a sentence in another domain without transporting by lf. as shown in in conclusion, there seem be a trade-off informative and the length. it can be handled well by ours, but for the baselines, they either lose informative or the accurate length. ## conclusion our work demonstrates the benefits of introducing the language rectified flow. the proposed flow can learn the neural ordinary differential equation model to transport between the source distribution and the target distribution, providing the unified and effective solution to generative modeling and domain transfer. our lf produces the fine-grained and high quality controllable text with fast simulation. the proposed strategy shows noticeable gains in performance across controllable text generation (parts-of speech, length and infill) and text editing (yelp and amazon). we further conduct the detailed study with the lf in different settings, e.g., comparing between constrained optimization vs. manually tuning coefficient, providing more evidence for faster simulation with different generation steps, and verifying the impact of different components. to summarize, the flow method is effective and general, with the potential to be incorporated into existing models for various nlp tasks.