Latent Template Induction with Gumbel-CRFs

Recent work in NLP has focused on model interpretability and controllability [63, 34, 24, 55, 16], aiming to add transparency to black-box neural networks and control model outputs with task-specific constraints. For tasks such as data-to-text generation [50, 63] or paraphrasing [37, 16], interpretability and controllability are especially important as users are interested in what linguistic properties – e.g., syntax [4], phrases [63], main entities [49] and lexical choices [16] – are controlled by the model and which part of the model controls the corresponding outputs.

Most existing work in this area relies on non-probabilistic approaches or on complex Reinforcement Learning (RL)-based hard structures. Non-probabilistic approaches include using attention weights as sources of interpretability [26, 61], or building specialized network architectures like entity modeling [49] or copy mechanism [22]. These approaches take advantages of differentiability and end-to-end training, but does not incorporate the expressiveness and flexibility of probabilistic approaches [45, 6, 30]. On the other hand, approaches using probabilistic graphical models usually involve non-differentiable sampling [29, 65, 34]. Although these structures exhibit better interpretability and controllability [34], it is challenging to train them in an end-to-end fashion.

In this work, we aim to combine the advantages of relaxed training and graphical models, focusing on conditional random field (CRF) models. Previous work in this area primarily utilizes the score function estimator (aka. REINFORCE) [62, 52, 29, 32] to obtain Monte Carlo (MC) gradient estimation for simplistic categorical models [44, 43]. However, given the combinatorial search space, these approaches suffer from high variance [20] and are notoriously difficult to train [29]. Furthermore, in a linear-chain CRF setting, score function estimators can only provide gradients for the whole sequence, while it would be ideal if we can derive fine-grained pathwise gradients [44] for each step of the sequence. In light of this, naturally one would turn to reparameterized estimators with pathwise gradients which are known to be more stable with lower variance [30, 44].

Our simple approach for reparameterizing CRF inference is to directly relax the sampling process itself. Gumbel-Softmax [27, 38] has become a popular method for relaxing categorical sampling. We propose to utilize this method to relax each step of CRF sampling utilizing the forward-filtering backward-sampling algorithm [45]. Just as with Gumbel-Softmax, this approach becomes exact as temperature goes to zero, and provides a soft relaxation in other cases. We call this approach Gumbel-CRF. As is discussed by previous work that a structured latent variable may have a better inductive bias for capturing the discrete nature of sentences [28, 29, 16], we apply Gumbel-CRF as the inference model in a structured variational autoencoder for learning latent templates that control the sentence structures. Templates are defined as a sequence of states where each state controls the content (e.g., properties of the entities being discussed) of the word to be generated.

Experiments explore the properties and applications of the Gumbel-CRF approach. As a reparameterized gradient estimator, compared with score function based estimators, Gumbel-CRF not only gives lower-variance and fine-grained gradients for each sampling step, which leads to a better text modeling performance, but also introduce practical advantages with significantly fewer parameters to tune and faster convergence ($\mathsection$ 6.1). As a structured inference network, like other hard models trained with REINFORCE, Gumbel-CRF also induces interpretable and controllable templates for generation. We demonstrate the interpretability and controllability on unsupervised paraphrase generation and data-to-text generation ($\mathsection$ 6.2). Our code is available at https://github.com/FranxYao/Gumbel-CRF.

Latent Variable Models and Controllable Text Generation. Broadly, our model follows the line of work on deep latent variable models [14, 30, 54, 23, 11] for text generation [28, 42, 16, 67]. At an intersection of graphical models and deep learning, these works aim to embed interpretability and controllability into neural networks with continuous [7, 67], discrete [27, 38], or structured latent variables [29]. One typical template model is the Hidden Semi-Markov Model (HSMM), proposed by Wiseman et al. [63]. They use a neural generative HSMM model for joint learning the latent and the sentence with exact inference. Li and Rush [34] further equip a Semi-Markov CRF with posterior regularization [18]. While they focus on regularizing the inference network, we focus on reparameterizing it. Other works about controllability include [55, 24, 33], but many of them stay at word-level [17] while we focus on structure-level controllability.

Monte Carlo Gradient Estimation and Continuous Relaxation of Discrete Structures. Within the range of MC gradient estimation [44], our Gumbel-CRF is closely related to reparameterization and continuous relaxation techniques for discrete structures  [64, 36, 31]. To get the MC gradient for discrete structures, many previous works use score-function estimators [52, 43, 29, 65]. This family of estimators is generally hard to train, especially for a structured model [29], while reparameterized estimators [30, 54] like Gumbel-Softmax [27, 38] give a more stable gradient estimation. In terms of continuous relaxation, the closest work is the differentiable dynamic programming proposed by Mensch and Blondel [40]. However, their approach takes an optimization perspective, and it is not straightforward to combine it with probabilistic models. Compared with their work, our Gumbel-CRF is a specific differentiable DP tailored for FFBS with Gumbel-Softmax. In terms of reparameterization, the closest work is the Perturb-and-MAP Markov Random Field (PM-MRF), proposed by Papandreou and Yuille [46]. However, when used for sampling from CRFs, PM-MRF is a biased sampler, while FFBS is unbiased. We will use a continuously relaxed PM-MRF as our baseline, and compare the gradient structures in detail in the Appendix.

In this section, we discuss how to relax a CRF with Gumbel-Softmax to allow for reparameterization. In particular, we are interested in optimizing $\phi$ for an expectation under a parameterized CRF distribution, e.g.,

We start by reviewing Gumbel-Max [27], a technique for sampling from a categorical distribution. Let $Y$ be a categorical random variable with domain $\{1,..,K\}$ parameterized by the probability vector $\mathbf{\pi}=[\pi_{1},..,\pi_{K}]$, denoted as $y\sim\text{Cat}(\mathbf{\pi})$. Let $\text{G}(0)$ denotes the standard Gumbel distribution, $g_{i}\sim\text{G}(0),i\in\{1,..,K\}$ are i.i.d. gumbel noise. Gumbel-Max sampling of $Y$ can be performed as: $y=\arg\max_{i}(\log\pi_{i}+g_{i})$. Then the Gumbel-Softmax reparameterization is a continuous relaxation of Gumbel-Max by replacing the hard argmax operation with the softmax,

where $\tilde{y}$ can be viewed as a relaxed one-hot vector of $y$. As $\tau\to 0$, we have $\tilde{y}\to y$.

Now we turn our focus to CRFs which generalize softmax to combinatorial structures. Given a sequence of inputs $x=[x_{1},..,x_{T}]$ a linear-chain CRF is parameterized by the factorized potential function $\Phi(z,x)$ and defines the probability of a state sequence $z=[z_{1},..,z_{T}],z_{t}\in\{1,2,...,K\}$ over $x$.

The conditional probability of $z$ is given by the Gibbs distribution with a factorized potential (equation 3). The partition function $Z$ can be calculated with the Forward algorithm [58] (equation 4) where $\alpha_{t}$ is the forward variable summarizing the potentials up to step $t$.

To sample from a linear-chain CRF, the standard approach is to use the forward-filtering backward-sampling (FFBS) algorithm (Algorithm 1). This algorithm takes $\alpha$ and $Z$ as inputs and samples $z$ with a backward pass utilizing the locally conditional independence. The hard operation comes from the backward sampling operation for $\hat{z}_{t}\sim p(z_{t}|\hat{z}_{t+1},x)$ at each step (line 6). This is the operation that we focus on.

We observe that $p(z_{t}|\hat{z}_{t+1},x)$ is a categorical distribution, which can be directly relaxed with Gumbel-Softmax. This leads to our derivation of the Gumbelized-FFBS algorithm(Algorithm 2). The backbone of Algorithm 2 is the same as the original FFBS except for two key modifications: (1) Gumbel-Max (line 8-9) recovers the unbiased sample $\hat{z}_{t}$ and the same sampling path as Algorithm 1; (2) the continuous relaxation of argmax with softmax (line 8) that gives the differentiable¹¹1 Note that argmax is also differentiable almost everywhere, however its gradient is almost 0 everywhere and not well-defined at the jumping point [48]. Our relaxed $\tilde{z}$ does not have these problems. (but biased) $\tilde{z}$.

We can apply this approach in any setting requiring structured sampling. For instance let $p_{\phi}(z|x)$ denote the sampled distribution and $f(z)$ be a downstream model. We achieve a reparameterized gradient estimator for the sampled model with $\tilde{z}$:

We further consider a straight-through (ST) version [27] of this estimator where we use the hard sample $\hat{z}$ in the forward pass, and back-propagate through each of the soft $\tilde{z}_{t}$.

We highlight one additional advantage of this reparameterized estimator (and its ST version), compared with the score-function estimator. Gumbel-CRF uses $\tilde{z}$ which recieve direct fine-grained gradients for each step from the $f$ itself. As illustrated in Figure 1 (here $f$ is the generative model), $\tilde{z}_{t}$ is a differentiable function of the potential and the noise: $\tilde{z}_{t}=\tilde{z}_{t}(\phi,g)$. So during back-propagation, we can take stepwise gradients $\nabla_{\phi}\tilde{z}_{t}(\phi,g)$ weighted by the uses of the state. On the other hand, with a score-function estimator, we only observe the reward for the whole sequence, so the gradient is at the sequence level $\nabla_{\phi}\log p_{\phi}(z|x)$. The lack of intermediate reward, i.e., stepwise gradients, is an essential challenge in reinforcement learning [56, 66]. While we could derive a model specific factorization for the score function estimator , this challenge is circumvented with the structure of Gumbel-CRF, thus significantly reducing modeling complexity in practice (detailed demonstrations in experiments).

An appealing use of reparameterizable CRF models is to learn variational autoencoders (VAEs) with a structured inference network. Past work has shown that these models (trained with RL) [29, 34] can learn to induce useful latent structure while producing accurate models. We introduce a VAE for learning latent templates for text generation. This model uses a fully autoregressive generative model with latent control states. To train these control states, it uses a CRF variational posterior as an inference model. Gumbel CRF is used to reduce the variance of this procedure.

Generative Model    We assume a simple generative process where each word $x_{t}$ of a sentence $x=[x_{1},..,x_{T}]$ is controlled by a sequence of latent states $z=[z_{1},..,z_{T}]$, i.e., template, similar to Li and Rush [34]:

Where Dec$(\cdot)$ denotes the decoder, FF$(\cdot)$ denotes a feed-forward network, $h_{t}$ denotes the decoder state, $[\cdot;\cdot]$ denotes vector concatenation. $e(\cdot)$ denotes the embedding function. Under this formulation, the generative model is autoregressive w.r.t. both $x$ and $z$. Intuitively, it generates the control states and words in turn, and the current word $x_{t}$ is primarily controlled by its corresponding $z_{t}$.

Inference Model    Since the exact inference of the posterior $p_{\theta}(z|x)$ is intractable, we approximate it with a variational posterior $q_{\phi}(z|x)$ and optimize the following form of the ELBO objective:

Where $\mathcal{H}$ denotes the entropy. The key use of Gumbel-CRF is for reparameterizing the inference model $q_{\phi}({z}|{x})$ to learn control-state templates. Following past work [25], we parameterize $q_{\phi}({z}|{x})$ as a linear-chain CRF whose potential is predicted by a neural encoder:

Where Enc$(\cdot)$ denotes the encoder and $h^{(\text{enc})}_{t}$ is the encoder state. With this formulation, the entropy term in equation 10 can be computed efficiently with dynamic programming, which is differentiable [39].

Training and Testing    The key challenge for training is to maximize the first term of the ELBO under the expectation of the inference model, i.e.

Here we use the Gumbel-CRF gradient estimator with relaxed samples $\tilde{z}$ from the Gumbelized FFBS (Algorithm 2) in both the forward and backward passes. For the ST version of Gumbel-CRF, we use the exact sample $\hat{z}$ in the forward pass and back-propogate gradients through the relaxed $\tilde{z}$. During testing, for evaluating paraphrasing and data-to-text, we use greedy decoding for both $z$ and $x$. For experiments controlling the structure of generated sentences, we sample a fixed MAP $z$ from the training set (i.e., the aggregated variational posterior), feed it to each decoder steps, and use it to control the generated $x$.

Extension to Conditional Settings    For conditional applications, such as paraphrasing and data-to-text, we make a conditional extension where the generative model is conditioned on a source data structure $s$, formulated as $p_{\theta}(x,z|s)$. Specifically, for paraphrase generation, $s=[s_{1},...,s_{N}]$ is the bag of words (a set, $N$ being the size of the BOW) of the source sentence, similar to Fu et al. [16]. We aim to generate a different sentence $x$ with the same meaning as the input sentence. In addition to being autoregressive on $x$ and $z$, the decoder also attend to [1] and copy [22] from $s$. For data-to-text, we denote the source data is formed as a table of key-value pairs: $s=[(k_{1},v_{1}),...,(k_{N},v_{N})]$, N being size of the table. We aim to generate a sentence $x$ that best describes the table. Again, we condition the generative model on $s$ by attending to and copying from it. Note our formulation would effectively become a neural version of slot-filling: for paraphrase generation we fill the BOW into the neural templates, and for data-to-text we fill the values into neural templates. We assume the inference model is independent from the source $s$ and keep it unchanged, i.e., $q_{\phi}(z|x,s)=q_{\phi}(z|x)$. The ELBO objective in this conditional setting is:

Our experiments are in two parts. First, we compare Gumbel-CRF to other common gradient estimators on the standard text modeling task. Then we integrate Gumbel-CRF to real-world models, specifically paraphrase generation and data-to-text generation.

Datasets    We focus on two datasets. For text modeling and data-to-text generation, we use the E2E dataset[50], a common dataset for learning structured templates for text [63, 34]. This dataset contains approximately 42K training, 4.6K validation and 4.6K testing sentences. The vocabulary size is 945. For paraphrase generation we follow the same setting as Fu et al. [16], and use the common MSCOCO dataset. This dataset has 94K training and 23K testing instances. The vocabulary size is 8K.

Metrics    For evaluating the gradient estimator performance, we follow the common practice and primarily compare the test negative log-likelihood (NLL) estimated with importance sampling. We also report relative metrics: ELBO, perplexity (PPL), and entropy of the inference network. Importantly, to make all estimates unbiased, all models are evaluated in a discrete setting with unbiased hard samples. For paraphrase task performance, we follow Liu et al. [37], Fu et al. [16] and use BLEU (bigram to 4-gram) [47] and ROUGE [35] (R1, R2 and RL) to measure the generation quality. We note that although being widely used, the two metrics do not penalize the similarity between the generated sentence and the input sentence (because we do not want the model to simply copy the input). So we adopt iBLUE [57], a specialized BLUE score that penalize the ngram overlap between the generated sentence and the input sentence, and use is as our primary metrics. The iBLUE score is defined as: iB$(i,o,r)=\alpha\text{B}(o,r)+(1-\alpha)\text{B}(o,i)$, where iB$(\cdot)$ denotes iBLUE score, B$(\cdot)$ denotes BLUE score, $i,o,r$ denote input, output, and reference sentences respectively. We follow Liu et al. [37] and set $\alpha=0.9$. For text generation performance, we follow Li and Rush [34] and use the E2E official evaluation script, which measures BLEU, NIST [5], ROUGE, CIDEr [60], and METEOR [3]. More experimental details are in the Appendix.

In this work, we propose a pathwise gradient estimator for sampling from CRFs which exhibits lower variance and more stable training than existing baselines. We apply this gradient estimator to the task of text modeling, where we use a structured inference network based on CRFs to learn latent templates. Just as REINFORCE, models trained with Gumbel-CRF can also learn meaningful latent templates that successfully encode lexical and structural information of sentences, thus inducing interpretability and controllability for text generation. Furthermore, the Gumbel-CRF gives significant practical benefits than REINFORCE, making it more applicable to real-world tasks.

Generally, this work is about controllable text generation. When applying this work to chatbots, one may get a better generation quality. This could potentially improve the accessibility [8, 53] for people who need a voice assistant to use an electronic device, e.g. people with visual, intellectual, and other disabilities [51, 2]. However, if not properly tuned, this model may generate improper sentences like fake information, putting the user at a disadvantage. Like many other text generation models, if trained with improper data (fake news [21], words of hatred [12]), a model could generate these sentences as well. In fact, one of the motivations for controllable generation is to avoid these situations [63, 12]. But still, researchers and engineers need to be more careful when facing these challenges.