Improving Joint Training of Inference Networks and Structured Prediction Energy Networks (2020)

Abstract: Deep energy-based models are powerful, but pose challenges for learning and inference (Belanger and McCallum, 2016). Tu and Gimpel (2018) developed an efficient framework for energy-based models by training “inference networks” to approximate structured inference instead of using gradient descent. However, their alternating optimization approach suffers from instabilities during training, requiring additional loss terms and careful hyperparameter tuning. In this paper, we contribute several strategies to stabilize and improve this joint training of energy functions and inference networks for structured prediction. We design a compound objective to jointly train both cost-augmented and test-time inference networks along with the energy function. We propose joint parameterizations for the inference networks that encourage them to capture complementary functionality during learning. We empirically validate our strategies on two sequence labeling tasks, showing easier paths to strong performance than prior work, as well as further improvements with global energy terms.

## introduction energy-based modeling @xcite associates a scalar compatibility measure to each configuration of input and output variables. @xcite formulated deep energy-based models for structured prediction, which they called structured prediction energy networks (spens). spens use arbitrary neural networks to define the scoring function over input/output pairs. however, this flexibility leads to challenges for learning and inference. the original work on spens used gradient descent for structured inference @xcite . @xcite found improvements in both speed and accuracy by replacing the use of gradient descent with a method that trains a neural network (called an "inference network") to do inference directly. their formulation, which jointly trains the inference network and energy function, is similar to training in generative adversarial networks @xcite , which is known to suffer from practical difficulties in training due to the use of alternating optimization @xcite . to stabilize training, @xcite experimented with several additional terms in the training objectives, finding performance to be dependent on their inclusion. moreover, when using the approach of @xcite , there is a mismatch between the training and test-time uses of the trained inference network. during training with hinge loss, the inference network is actually trained to do "costaugmented" inference. however, at test time, the goal is to simply minimize the energy without any cost term. @xcite fine-tuned the cost-augmented network to match the test-time criterion, but found only minimal change from this fine-tuning. this suggests that the cost-augmented network was mostly acting as a test-time inference network by convergence, which may be hindering the potential contributions of cost-augmented inference in max-margin structured learning @xcite . in this paper, we contribute a new training objective for spens that addresses the above concern and also contribute several techniques for stabilizing and improving learning. we empirically validate our strategies on two sequence labeling tasks from natural language processing (nlp), namely part-of-speech tagging and named entity recognition. we show easier paths to strong performance than prior work, as well as further improvements with global energy terms. we summarize our list of contributions as follows. while spens have been used for multiple nlp tasks, including multi-label classification @xcite , part-of-speech tagging @xcite , and semantic role labeling (belanger et al., 2017), they are not widely used in nlp. structured prediction is extremely common in nlp, but is typically approached using methods that are more limited than spens (such as conditional random fields) or models that suffer from a train/test mismatch (such as most auto-regressive models). spens offer a maximally expressive framework for structured prediction while avoiding the train/test mismatch and therefore offer great potential for nlp. however, the training and inference have deterred nlp researchers. while we have found benefit from training inference networks for machine translation in recent work @xcite , that work assumed a fixed, pretrained energy function. our hope is that the methods in this paper will enable spens to be applied to a larger set of applications, including generation tasks in the future. ## background we denote the input space by x . for an input @xmath0, we denote the structured output space by y(x). the entire space of structured outputs is denoted @xmath1. a spen @xcite defines an energy function e θ : x × @xmath2. at test time, for a given input x, prediction is done by choosing the output with lowest energy: however, solving equation (1) requires combinatorial algorithms because y is a structured, discrete space. this becomes intractable when e θ does not decompose into a sum over small "parts" of y. @xcite relaxed this problem by allowing the discrete vector y to be continuous; y r denotes the relaxed output space. they solved the relaxed problem by using gradient descent to iteratively minimize the energy with respect to y. the energy function parameters θ are trained using a structured hinge loss which requires repeated cost-augmented inference during training. using gradient descent for the repeated costaugmented inference steps is time-consuming and makes learning unstable @xcite . @xcite replaced gradient descent with a neural network trained to do efficient inference. this "inference network" a ψ : @xmath3 when training the energy function parameters θ, @xcite replaced the cost-augmented inference step in the structured hinge loss from @xcite with a costaugmented inference network f φ : here is the specific objective to jointly train θ (parameters of the energy function) and φ (parameters of the cost-augmented inference network): one challenge with the above optimization problem is that it requires training a separate inference network a ψ for test-time prediction after the energy function is trained. in this paper, we propose an alternative that trains the energy function and both inference networks jointly. in particular, we use a "compound" objective that combines two widely-used losses in structured prediction. we first present it without inference networks: as indicated, this loss can be viewed as the sum of the margin-rescaled hinge and perceptron losses for spens. two different inference problems are represented. the margin-rescaled hinge loss contains cost-augmented inference, shown as part of eq. ( 3 ). the perceptron loss contains the test-time inference problem, which is shown in eq. ( 1 ). tu and gimpel (2018 ) used a single inference network for solving both problems, so it was trained as a cost-augmented inference network during training and then fine-tuned as a test-time inference network afterward. we avoid this issue by training two inference networks, a ψ for test-time inference and f φ for cost-augmented inference: we treat this optimization problem as a minimax game and find a saddle point for the game similar to @xcite and @xcite . we use minibatch stochastic gradient descent and alternately optimize θ, φ, and ψ. the objective for the energy parameters θ in minibatch m is: when we remove 0-truncation (see sec. 4.1 for the motivation), the objective for the inference network parameters in minibatch m is: ## training stability and effectiveness we now discuss several methods that simplify and stabilize training spens with inference networks. when describing them, we will illustrate their impact by showing training trajectories for the twitter part-of-speech tagging task. ## energies for sequence labeling for our sequence labeling experiments in this paper, the input x is a length-t sequence of tokens, and the output y is a sequence of labels of length t . we use y t to denote the output label at position t, where y t is a vector of length l (the number of labels in the label set) and where y t,j is the jth entry of the vector y t . in the original output space y(x), y t,j is 1 for a single j and 0 for all others. in the relaxed output space y r (x), y t,j can be interpreted as the probability of the tth position being labeled with label j. we then use the following energy for sequence labeling @xcite : ## experimental setup we consider two sequence labeling tasks: twitter part-of-speech (pos) tagging @xcite and named entity recognition (ner @xcite . twitter part-of-speech (pos) tagging. we use the twitter pos data from @xcite and @xcite which contain 25 tags. we use 100-dimensional skip-gram @xcite embeddings from @xcite . like @xcite , we use a bilstm to compute the input feature vector for each position, using hidden size 100. we also use bilstms for the inference networks. the output of the inference network is a softmax function, so the inference network will produce a distribution over labels at each position. the ∆ is l1 distance. we train the inference network using stochastic gradient descent (sgd) with momentum and train the energy parameters using adam (kingma and ba, 2014). we also explore training the inference network using adam when not using the local ce loss. 1 in experiments with the local ce term, its weight is set to 1. ## results and analysis effect of zero truncation and local ce loss. ## related work there are several efforts aimed at stabilizing and improving learning in generative adversarial networks (gans) @xcite @xcite . progress in training gans has come largely ## conclusions and future work we contributed several strategies to stabilize and improve joint training of spens and inference networks. our use of joint parameterizations mitigates the need for inference network fine-tuning, leads to complementarity in the learned inference networks, and yields improved performance overall. these developments offer promise for spens to be more easily applied to a broad range of nlp tasks. future work will explore other structured prediction tasks, such as parsing and generation. we have taken initial steps in this direction, considering constituency parsing with the sequence-to-sequence model of @xcite . preliminary experiments are positive, foot_4 but significant challenges remain, specifically in defining appropriate inference network architectures to enable efficient learning. the seq2seq baseline achieves 82.80 f1 on the development set in our replication of @xcite . using a spen with our stacked parameterization, we obtain 83.22 f1.