Notes on Latent Structure Models and SPIGOT

TL;DR

This paper analyzes the SPIGOT method for training neural networks with discrete latent variables, providing new interpretations, linking it to existing methods, and proposing alternative variants for future exploration.

Contribution

It offers a novel interpretation of SPIGOT's gradient, connects it to other training methods, and introduces alternative variants for discrete latent variable models.

Findings

Provides a new perspective on SPIGOT's gradient

Links SPIGOT to other neural network training methods

Suggests alternative variants for future research

Abstract

These notes aim to shed light on the recently proposed structured projected intermediate gradient optimization technique (SPIGOT, Peng et al., 2018). SPIGOT is a variant of the straight-through estimator (Bengio et al., 2013) which bypasses gradients of the argmax function by back-propagating a surrogate "gradient." We provide a new interpretation to the proposed gradient and put this technique into perspective, linking it to other methods for training neural networks with discrete latent variables. As a by-product, we suggest alternate variants of SPIGOT which will be further explored in future work.

Tables

	unstructured	structured
vertices	$e_k$	$z_k$
interior points	$p$	$\mu$
maximization	argmax	MAP
expectation	softmax	Marg
euclidean projection	sparsemax	SparseMAP

Equations

$$\arg \underset{\mu \in \mathrm{conv}(\mathcal{Z})}{\max }{s}^{\top }\mu$$

$$\arg \underset{\begin{array}{c}p\in {\Delta }^{|\mathcal{Z}|}\\ \mu ={\mathbb{E}}_p[z]\end{array}}{\max }{s}^{\top }\mu \underset{H(p)}{\underbrace{-\sum _{z\in \mathcal{Z}}{p}_z\log p_z}}$$

$$\arg \underset{\mu \in \mathrm{conv}(\mathcal{Z})}{\max }{s}^{\top }\mu -\frac{1}{2}\Vert \mu {\Vert }^2$$

$$\hat{z}(s)=\arg \underset{z\in \mathcal{Z}}{\max }{s}^{\top }z.$$

$$\hat{y}=g_{\theta }\left(x,\stackrel{\hat{z}(s)}{\overbrace{\arg \underset{z\in \mathcal{Z}}{\max }{f}_{\varphi }(x{)}^{\top }z}}\right).$$

$${\nabla }_{\varphi }L(\hat{y},y^{\star })=\frac{\partial f_{\varphi }(x)}{\partial \varphi }\underset{=0}{\underbrace{\frac{\partial \hat{z}(s)}{\partial s}}}{\nabla }_{z}L(\hat{y}(\hat{z}),y^{\star })=0,$$

$${\tilde{\nabla }}_{s}L(\hat{y}(\hat{z}),y^{\star })$$

$${\mu }^{\star }=\arg \underset{\mu \in \mathrm{conv}(\mathcal{Z})}{\min }L(\hat{y}(\mu ),y^{\star }).$$

$${\mu }^{(t+1)}={\Pi }_{\mathrm{conv}(\mathcal{Z})}\left[{\mu }^{(t)}-{\eta }_t{\nabla }_{p}L(\hat{y}({\mu }^{(t)}),y^{\star })\right],$$

$$\tilde{\mu }={\Pi }_{\mathrm{conv}(\mathcal{Z})}\left[\hat{z}-\eta {\nabla }_{p}L(\hat{y}(\hat{z}),y^{\star })\right].$$

$${\ell }_{\mathrm{perc}}(\hat{z}(s),\tilde{\mu })$$

$${\nabla }_s{\ell }_{\mathrm{perc}}(\hat{z}(s),\tilde{\mu })$$

$${\nabla }_s{\ell }_{\mathrm{perc}}(\hat{z}(s),\tilde{\mu })=\hat{z}-(\hat{z}-\eta \gamma (\hat{z}))=\eta \gamma (\hat{z}).$$

$${\nabla }_s{\ell }_{\mathrm{cross}}(p(s),\tilde{p})$$

$$p^{(t+1)}$$

$$\tilde{p}$$

$${\nabla }_s{\ell }_{\mathrm{cross}}(p(s),\tilde{p})$$

$${\nabla }_{s}L(\hat{y}(\rho (s)),y)$$

$${\nabla }_s{\mathbb{E}}_s[L(\hat{y}(z),y)]$$

Full text

1 Introduction

In these notes, we assume a general latent structure model involving input variables $x\in\mathcal{X}$, output variables $y\in\mathcal{Y}$, and latent discrete variables $z\in\mathcal{Z}$. We assume that $\mathcal{Z}\subseteq\{0,1\}^{K}$, where $K\leq|\mathcal{Z}|$ (typically, $K\ll|\mathcal{Z}|$): i.e., the latent discrete variable $z$ can be represented as a $K$-th dimensional binary vector. This often results from a decomposition of a structure into parts: for example, $z$ could be a dependency tree for a sentence of $L$ words, represented as a vector of size $K=O(L^{2})$, indexed by pairs of word indices $(i,j)$, with $z_{ij}=1$ if arc $i\rightarrow j$ belongs to the tree, and [math] otherwise.

Notation.

In the following, we denote the $(|\mathcal{Z}|-1)$-dimensional probability simplex by $\Delta^{|\mathcal{Z}|}=\{p\in\mathbb{R}^{|\mathcal{Z}|}\mid p\geq 0,\sum_{z\in\mathcal{Z}}p(z)=1\}.$ Given $p\in\Delta^{|\mathcal{Z}|}$, we denote the expectation of a function $h:\mathcal{Z}\rightarrow\mathbb{R}^{D}$ under the probability distribution $p$ by $\mathbb{E}_{z\sim p}[h(z)]=\sum_{z\in\mathcal{Z}}p_{z}h(z)$. We denote the convex hull of the (finite) set $\mathcal{Z}\subseteq\mathbb{R}^{K}$ by $\mathrm{conv}(\mathcal{Z})=\left\{\mathbb{E}_{z\sim p}[z]\mid p\in\Delta^{|\mathcal{Z}|}\right\}$. The set $\mathrm{conv}(\mathcal{Z})$ can be interpreted as the smallest convex set which contains $\mathcal{Z}$.

Background.

In the literature on structured prediction, the set $\operatorname{conv}(\mathcal{Z})$ is sometimes called the marginal polytope, since any point inside it can be interpreted as some marginal distribution over parts of the structure (arcs) under some distribution over structures. There are three relevant problems that may be formulated in a structured setting:

• •

Finding the highest scoring structure, a.k.a. maximum a-posteriori (MAP): identify

[TABLE]

• •

Marginal inference: finding the (unique) marginal distribution induced by the scores $s$, given by an entropy projection onto the marginal polytope:

[TABLE]

• •

SparseMAP: finding the (unique) marginal distribution induced by the scores $s$, given by an Euclidean projection onto the marginal polytope:

[TABLE]

Unstructured setting.

We may encode the simple case of an unstructured categorical case by setting $\mathcal{Z}=\{e_{1},\dots,e_{K}\}$ which leads to $\operatorname{conv}(\mathcal{Z})=\Delta^{K}.$ The optimization problems above then recover some well known transformations, as described in the table below.

2 Latent structure model

Throughout, we assume a neural network classifier, parametrized by $\phi$ and $\theta$, which consists of three parts:

• •

An encoder function $f_{\phi}$ which, given an input $x\in\mathcal{X}$, outputs a vector of “scores” $s\in\mathbb{R}^{K}$, as $s=f_{\phi}(x)$;

• •

An argmax node which, given these scores, outputs the highest-scoring structure:

[TABLE]

• •

A decoder function $g_{\theta}$ which, given $x\in\mathcal{X}$ and $z\in\mathcal{Z}$, makes a prediction $\hat{y}\in\mathcal{Y}$ as $\hat{y}=g_{\theta}(x,z)$. We will sometimes write $\hat{y}(z)$ to emphasize the dependency on $z$. For reasons that will be clear in the sequel, we assume that the decoder also accepts convex combinations of latent variables as input, i.e., it may also output predictions $g_{\theta}(x,\mu)$ where $\mu\in\mathrm{conv}(\mathcal{Z})$.

Thus, given input $x\in\mathcal{X}$, this network predicts:

[TABLE]

To train this network, we assume a loss function $L(\hat{y},y^{\star})$, where $y^{\star}$ denotes the true output. We want to minimize this loss over the training data, using the gradient backpropagation algorithm.

We assume $\nabla_{\theta}L(\hat{y},y^{\star})$ is easy to compute: it can be done by performing standard gradient backpropagation from the output layer until the output of the argmax node. The main challenge of this model is to sidestep the argmax node to propagate gradient information to the encoder parameters. Indeed, we have:

[TABLE]

so no gradient will flow to the encoder. Common approaches to circumvent this problem include:

• •

Replace the argmax node by a stochastic node where $z$ is a random variable parametrized by $s$ (e.g., using a Gibbs distribution). Then, compute the gradient of the expected loss $\mathbb{E}_{z\sim p_{s}}[L(\hat{y}(z),y^{\star})]$. This is the approach underlying REINFORCE, score function estimators, and minimum risk training (Williams, 1992; Smith and Eisner, 2006; Stoyanov et al., 2011). Niculae et al. (2018b) explore a sparse alternative to the Gibbs distribution.

• •

Keep the network deterministic, but do a continuous relaxation of the argmax node, for example replacing it with softmax or sparsemax (Martins and Astudillo, 2016). In the structured case, this gives rise to structured attention networks (Kim et al., 2017) and their sparse variant, SparseMAP (Niculae et al., 2018a). Mathematically, this corresponds to moving the expectation inside the loss, optimizing $L(\hat{y}(\mathbb{E}_{z\sim p}[z]),y^{\star})]$.

• •

Keep the argmax node and perform the usual forward computation, but backpropagate a surrogate gradient. This is the approach underlying straight-through estimators (Bengio et al., 2013) and SPIGOT (Peng et al., 2018). We will develop this approach in the remainder of these notes.

In what follows, we assume that:

• •

We have access to the gradient $\gamma(z):=\nabla_{z}L(\hat{y}(z),y^{\star})$;111This gradient would not exist if the decoder $g_{\theta}$ were defined only for the vertices $z\in\mathcal{Z}$ and not convex combinations thereof. This assumption is not needed in the minimum risk training approach discussed toward the end of this note.

• •

We want to replace the (zero) gradient $\nabla_{s}L(\hat{y}(z),y^{\star})$ by a surrogate $\tilde{\nabla}_{s}L(\hat{y}(z),y^{\star})$.

3 SPIGOT as the approximate computation of a pulled back loss

We now provide an interpretation of SPIGOT as the minimization of a “pulled back” loss with respect to the latent variable $z$. SPIGOT uses the following surrogate gradient:

[TABLE]

where we used the fact that the SparseMAP transformation (Niculae et al., 2018a) is equivalent to an Euclidean projection, i.e. $\mathrm{SparseMAP}(s)=\displaystyle\arg\max_{\mu\in\operatorname{conv}(\mathcal{Z})}s^{\top}\mu-\frac{1}{2}\|\mu\|^{2}=\arg\min_{\mu\in\operatorname{conv}(\mathcal{Z})}\|\mu-s\|$.

3.1 Intermediate loss on the latent variable

Let us start by stating an obvious fact, which will draw intuition for the rest: if we had supervision for the latent variable $z$ (e.g., if the true label $z^{\star}$ was revealed to us), we could simply define an intermediate loss $\ell(\hat{z},z^{\star})$ which can induce nonzero updates to the encoder parameters. In fact, if $K=|\mathcal{Z}|$ is small, we can enumerate all possible values of $z$ and define $z^{\star}$ as the one that minimizes the downstream loss, $z^{\star}=\arg\min_{z}L(\hat{y}(z),y^{\star})$, using the current network parameters $\theta$; this $z^{\star}$ would become our “groundtruth.”

While this seems somewhat sensible, we may expect some instability in the beginning of the training process, since the decoder parameters $\theta$ are likely to be very suboptimal at this stage. A more robust procedure is to allow for some label uncertainty: instead of picking a single label $z^{\star}\in\mathcal{Z}$, pick the convex combination $\mu\in\operatorname{conv}(\mathcal{Z})$ that minimizes $L(\hat{y}(\mu^{\star}),y^{\star})$. In fact, it is likely that the $\mu^{\star}$ that minimizes the downstream loss will not put all the probability mass on a single label, and we may benefit from that if the downstream loss is what we care about. With this in mind, we define:

[TABLE]

For most interesting predictive models $\hat{y}(\mu)$, this optimization problem is non-convex and lacks a closed form solution. One common strategy is the projected gradient algorithm, which iteratively performs the following updates:

[TABLE]

where $\eta_{t}$ is a step size and $\Pi_{\mathcal{S}}(u)=\arg\min_{u^{\prime}\in\mathcal{S}}\|u^{\prime}-u\|$ denotes the Euclidean projection of point $u$ onto the set $\mathcal{S}$. With a suitable choice of step sizes, the projected gradient algorithm is guaranteed to converge to a local optimum of Eq. 8. If we initialize $\mu^{(0)}=\hat{z}=\arg\max_{z\in\mathcal{Z}}s^{\top}z$ and run a single iteration of projected gradient, we obtain the following estimate $\tilde{\mu}$ of $\mu^{\star}$:

[TABLE]

We can now treat $\tilde{\mu}$ as if it were the “groundtruth” label distribution, turning the optimization of the encoder $f_{\phi}(x)$ as if it were a supervised learning problem. If we use a perceptron loss,

[TABLE]

we get the following gradient:

[TABLE]

which is precisely the SPIGOT gradient surrogate presented in Eq. 7. This leads to the following insight into how SPIGOT updates the encoder parameters:

SPIGOT minimizes the perceptron loss between $z$ and a pulled back target computed by one projected gradient step on $\displaystyle\min_{\mu\in\operatorname{conv}(\mathcal{Z})}L(\hat{y}(\mu),y)$ starting at $\hat{z}$.

This construction suggests some possible alternate strategies. The first results in a well known algorithm, while the rest result in novel variations.

Relaxing the $\operatorname{conv}(\mathcal{Z})$ constraint.

The constraints in Eq. 8 make the optimization problem more complicated. We relax them and define \mu^{\star}=\arg\min_{\mu\in{\color[rgb]{.75,0,.25}\mathbb{R}^{K}}}L(\hat{y}(\mu),y^{\star}). This problem still must be tackled iteratively, but the projection step can now be avoided. One iteration of gradient descent yields $\tilde{\mu}=\hat{z}-\eta\gamma{\hat{z}}$. The perceptron update then recovers straight-through,222Specifically, the “identity” variant of STE, in which the backward pass acts as if $\frac{\partial\hat{z}(s)}{\partial s}=\mathrm{Id}$ (Bengio et al., 2013). via a novel derivation:

[TABLE]

This leads to the following insight into straight-through and its relationship to SPIGOT:

STE minimizes the perceptron loss between the latent $z$ and a pulled back target computed by one gradient step on $\displaystyle\min_{\mu\in\mathbb{R}^{K}}L(\hat{y}(\mu),y)$ starting at $\hat{z}$.

Multiple projected gradient steps.

Instead of a single projected gradient step, we could have run multiple steps of the iteration in Eq. 9. We would expect this to yield an estimate $\tilde{\mu}$ closer to $\mu^{\star}$, at the cost of more computation.

Different initialization.

The projected gradient update in Eq. 10 uses $\mu^{(0)}=\hat{z}=\arg\max_{z\in\mathcal{Z}}s^{\top}z$ as the initial point. This is a sensible choice, if we believe the encoder prediction $\hat{z}$ is close enough to the “groundtruth” $\mu^{\star}$, and it is computationally convenient because $\hat{z}$ has already been computed in the forward propagation step and can be cached. However, other initializations are possible, for example $\mu^{(0)}=\mathrm{Marg}(s)$, or $\mu^{(0)}=0$.

Different intermediate loss function.

For simplicity, consider the unstructured case. Let $p(s):=\mathrm{softmax}(s)$. If we use the cross-entropy loss $\ell_{\mathrm{cross}}(p(s),\tilde{p})=-\sum_{k=1}^{K}\tilde{p}_{k}\log p_{k}(s)$ instead the perceptron loss, we get

[TABLE]

This generalizes easily to the CRF loss in the structured case.

Exponentiated gradient instead of projected gradient.

Also in the unstructured case, the exponentiated gradient algorithm (Kivinen and Warmuth, 1997) tackle the constrained optimization problem in Eq. 8 with the following multiplicative updates:

[TABLE]

where each point $p^{(t)}$ is strictly positive. This includes the initializer $p^{(0)}$, so we cannot have $p^{(0)}=\hat{z}$; for this reason we assume $p^{(0)}=p(s)=\mathrm{softmax}(s)$. A single iteration of exponentiated gradient with this initialization gives:

[TABLE]

With the cross-entropy loss, i.e. the Kullback-Leibler divergence $\operatorname{KL}(p(s)\mid\tilde{p})$, we obtain:

[TABLE]

i.e., the surrogate gradient is the difference of a softmax with a softmax with “perturbed” scores. This generalizes to an instance of mirror descent with Kullback-Leibler projections in the structured case.

4 Relation to other methods for latent structure models

4.1 Continuous relaxation of argmax

To simplify, let us consider the case where $z$ is a categorical variable. If we replace the argmax node by a continuous transformation $\rho:\mathbb{R}^{K}\rightarrow\Delta^{K}$ (e.g., a softmax with a temperature), the gradient $\nabla_{s}L(\hat{y}(\rho(s)),y^{\star})$ can be exactly computed by the chain rule:

[TABLE]

where $J_{\rho}(s)\in\mathbb{R}^{K\times K}$ is the Jacobian of transformation $\rho$ at point $s$.

4.2 Minimum risk training

In this case, the network has a stochastic node $z\sim\rho(s)=p_{z}(s)$, with $\rho$ as above. The gradient of the risk with respect to $s$ is:

[TABLE]

where $\ell\in\mathbb{R}^{K}$ is a vector where the $z$th entry contains the loss value $L(\hat{y}(z),y)$.

Another way of writing the gradient above, noting that $\nabla_{s}p_{z}(s)=p_{z}(s)\nabla_{s}\log p_{z}(s)$, is:

[TABLE]

It is interesting to compare this gradient with the SPIGOT surrogate gradient in Eq. 7. Also here a “pulled-back loss” (now $-\log p_{z}(s)$) is used in the gradient computation, this time as part of a weighted sum, where the weights are the reward $-L(\hat{y}(z),y)$ and the probability $p_{z}(s)$. For example, if the downstream loss minimizer is $z^{\star}$ and all $z\neq z^{\star}$ are equally bad (i.e., if they have the same loss), then we obtain

[TABLE]