Understanding the Mechanics of SPIGOT : Surrogate Gradients for Latent Structure Learning (2020)

Abstract: Latent structure models are a powerful tool for modeling language data: they can mitigate the error propagation and annotation bottleneck in pipeline systems, while simultaneously uncovering linguistic insights about the data. One challenge with end-to-end training of these models is the argmax operation, which has null gradient. In this paper, we focus on surrogate gradients, a popular strategy to deal with this problem. We explore latent structure learning through the angle of pulling back the downstream learning objective. In this paradigm, we discover a principled motivation for both the straight-through estimator (STE) as well as the recently-proposed SPIGOT – a variant of STE for structured models. Our perspective leads to new algorithms in the same family. We empirically compare the known and the novel pulled-back estimators against the popular alternatives, yielding new insight for practitioners and revealing intriguing failure cases.

## introduction natural language data is rich in structure, but most of the structure is not visible at the surface. machine learning models tackling high-level language tasks would benefit from uncovering underlying structures such as trees, sequence tags, or segmentations. traditionally, practitioners turn to pipeline approaches where an external, pretrained model is used to predict, e.g., syntactic structure. the benefit of this approach is that the predicted tree is readily available for inspection, but the downside is that the errors can easily propagate throughout the pipeline and require further attention @xcite @xcite . in contrast, deep neural architectures tend to eschew such preprocessing, and instead learn † work partially done while vn was at the instituto de telecomunicac ¸ões, lisbon. soft hidden representations, not easily amenable to visualization and analysis. the best of both worlds would be to model structure as a latent variable, combining the transparency of the pipeline approach with the endto-end unsupervised representation learning that makes deep models appealing. moreover, largecapacity model tend to rediscover structure from scratch @xcite , so structured latent variables may reduce the required capacity. learning with discrete, combinatorial latent variables is, however, challenging, due to the intersection of large cardinality and null gradient issues. for example, when learning a latent dependency tree, the latent parser must choose among an exponentially large set of possible trees; what's more, the parser may only learn from gradient information from the downstream task. if the highestscoring tree is selected using an argmax operation, the gradients will be zero, preventing learning. one strategy for dealing with the null gradient issue is to use a surrogate gradient, explicitly overriding the zero gradient from the chain rule, as if a different computation had been performed. the most commonly known example is the straight-through estimator (ste; @xcite , which pretends that the argmax node was instead an identity operator. such methods lead to a fundamental mismatch between the objective and the learning algorithm. the effect of this mismatch is still insufficiently understood, and the design of successful new variants is therefore challenging. for example, the recently-proposed spigot method @xcite found it beneficial to use a projection as part of the surrogate gradient. in this paper, we study surrogate gradient methods for deterministic learning with discrete structured latent variables. our contributions are: figure 1 : a model with a discrete latent variable z. given an input x, we assign a score s @xmath0, and pick the highest scoring one, ẑ, to predict ŷ = g θ (ẑ). for simplicity, here g θ does not access x directly. (a). since argmax has null gradients, the encoder parameters φ do not receive updates. (b). if ground truth supervision were available for the latent z, φ could be trained jointly with an auxiliary loss. (c). as such supervision is not available, we induce a best-guess label µ by pulling back the downstream loss. this strategy recovers the ste and spigot estimators. while the discrete methods do not outperform the relaxed alternatives using the same building blocks, we hope that our interpretation and insights would trigger future latent structure research. the code for the paper is available on https: //github.com/deep-spin/understanding-spigot. ## related work discrete latent variable learning is often tackled in stochastic computation graphs, by estimating the gradient of an expected loss. an established method is the score function estimator (sfe) @xcite @xcite . sfe is widely used in nlp, for tasks including minimum risk training in nmt @xcite and latent linguistic structure learning @xcite . in this paper, we focus on the alternative strategy of surrogate gradients, which allows learning in deterministic graphs with discrete, argmax-like nodes, rather than in stochastic graphs. examples are the straight-through estimator (ste) @xcite and the structured projection of intermediate gradients optimization technique (spigot; @xcite . recent work focuses on studying and explaining ste. @xcite obtained a convergence result in shallow networks for the unstructured case. @xcite show that ste can be interpreted as the simulation of the projected wasserstein gradient flow. ste has also been studied in binary neural networks @xcite and in other applications @xcite . other methods based on the surrogate gradients have been recently explored @xcite . a popular alternative is to relax an argmax into a continuous transform such as softmax or sparsemax (martins and astudillo, 2016), as seen for instance in soft attention mechanisms @xcite , or structured attention networks @xcite @xcite @xcite . in-between surrogate gradients and relaxation is gumbel softmax, which uses the gumbel-max reparametrization to sample from a categorical distribution, applying softmax either to relax the mapping or to induce surrogate gradients @xcite . gumbel-softmax has been successfully applied to latent linguistic structure as well @xcite . for sampling from a structured variable is required, the perturb-and-map technique @xcite has been successfully applied to sampling latent structures in nlp applications @xcite . ## structured prediction preliminaries we assume a general latent structure model involving input variables @xmath1, output variables @xmath2, and latent discrete variables @xmath3. we assume that z ⊆ {0, 1} k , where @xmath4, k |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 @xmath5, indexed by pairs of word indices (i, j), with z @xmath6, and 0 otherwise. this allows us to define the score of a structure as the sum of the scores of its parts. given a vector @xmath7, containing scores for all possible parts, we define score(z) := s z. notation. we denote by e k the one-hot vector with all zeros except in the k th coordinate. we denote the simplex by background. in the context of structured prediction, the set m := conv(z) is known as 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: unstructured setting. as a check, we consider the encoding of a categorical variable with k distinct choices, encoding each choice as a one-hot vector e k and setting @xmath8, . . . , e k }. in this case, conv(z) = k . the optimization problems above then recover some well known transformations, as described in ## latent structure models throughout, we assume a classifier parametrized by φ and θ, which consists of three parts: thus, given input @xmath9, this network predicts: to train this network, we minimize a loss function l(ŷ, y), where y denotes the target label; a common example is the negative log-likelihood loss. the gradient w.r.t. the decoder parameters, ∇ θ l(ŷ, y), is easy to compute using automatic differentiation on g θ . the main challenge is to propagate gradient information through the argmax node into the encoder parameters. indeed, so no gradient will flow to the encoder. we list below the three main categories of approaches that tackle this issue. introducing stochasticity. replace the argmax node by a stochastic node where z is modeled as a random variable z parametrized by s (e.g., using a gibbs distribution). then, instead of optimizing a deterministic loss l(ŷ(ẑ), y), optimize the expectation of the loss under the predicted distribution: the expectation ensures that the gradients are no longer null. this is sometimes referred to as minimum risk training @xcite , and typically optimized using the score function estimator (sfe; @xcite @xcite . relaxing the argmax. keep the network deterministic, but relax the argmax node into a continuous function, for example replacing it with softmax or sparsemax (martins and astudillo, 2016). in the structured case, this gives rise to structured attention networks @xcite and their sparsemap variant @xcite . this corresponds to moving the expectation inside the loss, optimizing l ŷ(e z∼p(z;s) [z] µ ), y . inventing a surrogate gradient. keep the argmax node and perform the usual forward computation, but backpropagate a different, non-null gradient in the backward pass. this is the approach underlying straight-through estimators (hinton, 2012; @xcite and spigot @xcite . this method introduces a mismatch between the measured objective and the optimization algorithm. in this work, we proposed a novel, principled justification for inducing surrogate gradients. in what follows, we assume that: ## spigot as the approximate optimization of a pulled back loss we next provide a novel interpretation of spigot as the minimization of a "pulled back" loss. spigot uses the surrogate gradient: ## experiments armed with a selection of surrogate gradient methods, we now proceed to an experimental comparison. for maximum control, we first study a synthetic unstructured experiment with known data generating process. this allows us to closely compare the various methods, and to identify basic failure cases. we then study the structured case of latent dependency trees for sentiment analysis and natural language inference in english. full training details are described in appendix a. ## conclusions in this work, we provide a novel motivation for straight-through estimator (ste) and spigot, based on pulling back the downstream loss. we derive promising new algorithms, and novel insight into existing ones. unstructured controlled experiments suggest that our new algorithms, which use the cross-entropy loss instead of the perceptron loss, can be more stable than spigot while accurately disentangling the latent variable. differentiable relaxation models (using softmax and sparsemax) are the easiest to optimize to high downstream accuracy, but they fail to correctly identify the latent clusters. on structured nlp experiments, relaxations (sparsemap and marginals) tend to overall perform better and be more stable than straightthrough variants in terms of classification accuracy. however, the lack of gold-truth latent structures makes it impossible to assess recovery performance. we hope that our insights, including some of our negative results, may encourage future research on learning with latent structures.