Inference-Masked Loss for Deep Structured Output Learning

1 Introduction

Structured learning considers learning to predict output variables that are interdependent and need to obey some structural constraints. Structured learning approaches involve an inference phase after computing local predictions along with their probabilities/scores which can be generally formulated by a Maximum A Posteriori (MAP) inference problem. On one hand, structured learning are powerful in its flexibility to incorporate domain knowledge and requirement of fewer data. On the other hand, deep neural networks have achieved significant results by using large number of parameters and examples. Exploiting structured output in deep neural networks can help to combine the power of learning representations and leveraging the use of output structure during training. It will enable deep neural networks in imposing explicit structural constraints on their predictions by using domain knowledge on top of the data.

∗Contact Author

Current neural network architectures are able to represent structures such as sequences [Cho et al., 2014] and graphs [Wu et al., 2019]. However, the structure is mostly considered in the input side based on processing real-valued tensors and there is a lack of systematic way to predict discrete structured outputs. Neural networks are data-hungry. If we rely on the model to capture the global structure without explicitly modeling them, the need to the large data is even more critical. While providing explicit structural information to the model can reduce the need to large parameters and examples of deep neural networks, introducing the nondifferentiable inference to gradient-based training is still a critical challenge. Deep neural networks are trained by backpropagation algorithms that calculate the gradient of the parameters. In most cases, the inference procedures are nondifferentiable and the gradients can not be calculated based on the model output and the loss function.

In the previous studies on combining inference with learning of neural networks, we can observe at least two different types of techniques. One tries to learn local models to generate local predictions and rely on the inference for a global result, for example, structured perceptron [Collins, 2002]. However, the studies are mostly limited to linear models. More studies focus on incorporating the constraints from output structure and domain knowledge in the architecture of neural networks [Li and Srikumar, 2019] or learning algorithms [Chen et al., 2015; Nandwani et al., 2019; Liang et al., 2019] so that the local neural networks learn to generate outputs respecting all the constraints.

In this work, we present a new loss function to extend deep neural network to use inference and exploit the structure of the output. Our proposed Inference-Masked Loss (IML) allows neural networks to keep the local predictions, even if they are false, as much as the inference can correct them. IML belongs to the first class of techniques mentioned above because instead of forcing the local output to respect the constraints, it trains the deep neural networks to rely on the inference. Compared to structured perceptron [Collins, 2002], IML integrates the inference result with the loss in a fully differentiable way and makes it compatible with arbitrary underlying deep neural networks. By relying on inference, the network needs less number of examples compared to training without information about the output structure. There are other recent works involving deep neural networks learning with inference. For example, Nandwani et al. formulates the learning with inference as a primal-dual problem, solves it by alternating optimization, and addresses the problem by training the neural networks to respect the constraints [Nandwani et al., 2019]. However, to produce the structured output and respect the constraints, it requires the neural networks to use global information even in producing a part of the output. In contrast, IML allows the neural networks to make local predictions without the constraints, and rely on the inference to correct the false predictions. It helps to decompose the dependencies of the output variables and allows the (sub-) neural networks to focus on training for parts of the output locally. The contribution of this paper is as follows:

• We propose the inference-masked loss (IML) that utilizes the global inference not only for prediction but also to train the parameters for structured output learning with deep neural networks.
• We evaluate the proposed method with two different tasks on three datasets under various settings to show the effectiveness of the proposed approach, especially with low training data.

2 RelatedWork

In structured learning, the inference is usually described by a Maximum A Posteriori (MAP) problem and formulated by a log-linear model, which takes the form of a combination of local feature functions or probabilistic factors [Sutton and McCallum, 2006].

The structured perceptron was proposed to extend the perceptron with structured output [Collins, 2002]. The algorithm updates the model parameters based on the difference between inference assignments and the ground-truth. A recent practice [Weiss et al., 2015] has shown that the algorithm works with deep features. The structured perceptron was applied on top of a deep neural network that generates the features. However, the structured perceptron was trained separately from the deep neural network. Our proposed IML shares the same intuition of structured perceptron, that is, training based on the inference error. Both algorithms are tolerant to the local inconsistencies where there is no error in the inference results. However, the structured perceptron is limited to linear models. IML is a differentiable loss function to train deep neural networks. It can be incorporated with all sophisticated deep neural network architectures and deep learning regularizers.

The recent research regarding training neural networks concerning structural constraints are mainly imposing the constraints either by modifying the architectures, loss functions, or the formulation of the optimization.

Shen et al. worked on designing specific architectures and proposed a cumulative softmax to enforce ordering relation among a sequence of neurons. Li and Srikumar proposed a new approach to augment the neural networks with first-order logic by modifying the architecture by applying soft logic on the named neurons inside an existing architecture.

For reformulating the loss function, there is a collection of classical structured loss [Edunov et al., 2018] used in combination with deep features. "SoftmaxMargin" loss [Gimpel and Smith, 2010] augments negative log-likelihood with a cost that penalizes high cost outputs proportional to their cost, which is similar to the "Margin Softmax" loss [Yang et al., 2019; Ilharco et al., 2019]. Researchers also investigated the approaches to incorporate soft and hard constraints with the loss function. Hard constraints are applied by smoothing out the loss function in [Pathak et al., 2015] and by using Krylov subspace methods in [Marquez Neila et al., 2017]. Also, Muralidhar et al. proposed to use an added term to the loss function imposing soft constraints. In this approach, hard constraints can be achieved by using a very large cost coefficient.

Training deep neural networks with inference can be formulated by a min-max optimization problem. The constrained maximization, which can be converted to unconstrained form with the Lagrangian multiplier, plays the role of the inference. The minimization updates the weights in the neural networks. To address the problem, Nandwani et al. took the dual form of the problems and solved the two optimizations alternatively. Chen et al. proposed to use a single step of message passing to estimate the inference and to train the neural networks with the estimated inference. Constraint guided semi-supervised learning algorithm was proposed to generate pseudo-labeled data by the constraints to train the parameters in a model with a supervised learning algorithm [Chang et al., 2012]. Taking a different type of approach, Maes et al. formalized the problem of structured prediction as a Reinforcement Learning task and proposed to solve it by learning the policy.

Instead of training neural networks to respect the constraints, our proposed IML allows neural networks to keep local predictions, even if they are false, as long as the global inference can correct them. The model will count on the inference for a global output and tends to learn better decomposition of the task that helps to improve the generalizability of the model.

As mentioned before, MAP inference is usually an essential component of structured learning models. Inference can be performed with various approaches including probabilistic models, classical search algorithms, dynamic programming, or constrained optimization techniques. Some of recent practices integrate predictions of deep neural networks with conditional random fields (CRF) [Ma and Hovy, 2016] based on a chain structure with Viterbi algorithm. Besides, the beamsearch algorithm is commonly used to improve the search efficiency of the inference. However, the efficiency of these approaches heavily depends on the independence structure of the underlying probabilistic model in a certain application. On a different thread, using integer linear programming (ILP) approach, one can represent the dependencies with linear constraints that can be solved by efficient off-the-shelf solvers 1 . Roth and Yih propose to model CRF inference as ILP for efficient solving and imposing global constraints over the output space beyond sequential structures [Roth and Yih, 2005].

1 Solving ILP is known to be NP-Hard. However, the derived ILPs are usually very sparse and can be solved in a satisfactory time.

One essential but orthogonal aspect to inference problem is the representation of the structure. Researchers have investigated approaches to represent domain knowledge in the form of first-order expressions and using inference over locally trained classifiers [Rizzolo and Roth, 2007], which was extended to the approaches that support joint training and inference [Chang et al., 2012; Kordjamshidi et al., 2015]. Ontologies were also investigated to generate constraints to improve local outputs of deep neural networks by global inference [Guo et al., 2019]. We follow these approaches to represent the output structure and domain knowledge by ontologies and logical expressions, which are converted to linear constraints in the inference problem.

3 Inference-Masked Loss

3.1 Model Formulation

We assume training a deep neural network given a set of examples {(X, Y )}, that is, pairs of inputs X and outputs Y . Both inputs and outputs can have arbitrary structures. We assume the structure of the output variables can be expressed with a set of linear constraints among them, denoted by C (Y ) ≤ 0. The constraints can be represented with a set of inequalities without loss of generality [Nandwani et al., 2019]. These constraints originate from our knowledge about the domain and are expressed using logical expressions or in ontologies that are converted to linear constraints -this is detailed in Section 3.4. In this study, we focus on the case where all parts of the output structure are discrete values and can be associated with a collection of binary predicates q ∈ Q where Q is the set of all possible predicates. Each component Yq of the structured output Y is a binary random variable and indicates whether q is true or not.

We are interested in the joint probability distribution of structured output variables y given input X, P (y|X). The inference is finding the best assignment of y that yields the maximum joint probability. Our framework and loss function are agnostic to the inference method. In this study, we exploit ILP, which solves the problems with constraints that are not satisfying first order Markov property, efficiently [Roth and Yih, 2005]. The joint probability distribution is estimated by a normalized exponential function with logarithm estimated by the total scoring function g (y,X; θ) that logP (y|X) ∝ g (y,X; θ). Given the linear constraints, the inference can be represented as an ILP problem as follows.

To train such a model, first, we make local predictions for each component of y. In other words, for each predicate q ∈ Q, we make a local prediction fq (X; θ) with deep neural networks, where fq is a local network and θ is the weights and biases. Following the log-linear model formulation, we calculate a local scoring function gq = log fq (X; θ) and a corresponding negative term g¬q = log (1 - fq (X; θ)). We calculate the total scoring function for the output y by a linear model,

Notations:

↑: a probability close to 1; ↓: a probability close to 0.

0-: a voided penalty in IML when the local prediction is correct;

0+: a voided penalty in IML when the false prediction is corrected by the inference.

Table 1: Penalty term of NLL and IML regarding one predicate q.

The inference problem in (1) can be solved efficiently by offthe-shelf solvers 2 .

3.2 The Loss Function

The commonly-used negative log-likelihood (NLL) loss function is given by

In fact, the ground-truth Yq and the term (1 - Yq) serves as a selective mask to select penalizing the parameters of negative log-likelihood - log fq (X; θ) or - log (1 - fq (X; θ)).

Inspired by the idea mentioned above, we introduce the IML. For each predicate q, if the associated component in the global inference y∗ q is correct according to the ground-truth Yq, IML will nullify the corresponding terms by the mask. The resulting IML is as follows,

It can be observed that both masks (1-y∗ q )Yq and y∗ q (1-Yq) will be zero if y∗ q = Yq. When y∗ q ̸= Yq, the masks will become Yq and (1-Yq) as in the NLL that, consequently, selects either the - log fq (X; θ) or - log (1 - fq (X; θ)) to be penalized according to the ground-truth.

Table 1 shows the term that is added to the total loss regarding a predicate q based on NLL and IML, respectively. The first column denotes the True/False Positive/Negative notations regarding the local prediction. The ↑ in the third column indicates a probability close to 1 while a ↓ indicates that close to 0.

In the NLL, the term is selected by the ground-truth Yq. Moreover, the magnitude is determined by how correct (or incorrect) the local prediction fq (X; θ) is. As we introduced the inference to IML, the conditions become different. 0+ and 0- are the cases that the global inference gives the correct result (H∗ q = Yq), while the others are for the cases when the inference is wrong. (FN, y∗ q = 0) and (FP, y∗ q = 1) are the criminals that we want to penalize. The 0+ are the false local predictions that can be corrected by inference. Therefore, we do not need to penalize them. The (TP, y∗ q = 0) and (TN, y∗ q = 1) are innocent cases because they were correct locally and inference changed them to be wrong. However, we argue that if they get a higher local score, they should be able to help in correcting other false predictions. The 0- are not problematic because they are correct, and they support the correct inference. No update needs to be applied in this case.

2 We solve the ILP problems by Gurobi https://www.gurobi.com

3.3 Combination

Training solely with IML in (4) could lead to "lazy" update and is more likely to be trapped in a local minimum compared to traditional loss because IML would ignore many of the local errors. The effect is two-fold. On one hand, with IML, the model receives fewer updates, which leads to a slower convergence. On the other hand, using IML would miss the opportunity to update a weak local model, which may cause trouble at testing time. We combine NLL and IML to enforce accurate local prediction.

where λ is weighting the trade-off between NLL and IML, which is a hyper-parameter for the proposed loss. λ is also interpreted as the ratio of penalty to apply when the inference can resolve the errors on the local prediction and correct them in the global inference. λ times the corresponding NLL loss term will be penalized. Otherwise, if there are still errors according to the ground-truth, the penalty term is the same as NLL. This leads to having a better local prediction even if the inference can resolve the incorrectness. It is shown to be helpful in the Experiments Section and Analysis Section. However, λ is a new (and the only) hyper-parameter in this work. Experimental results with different λ will be reported and analyzed in Section 4.

3.4 Structure and Domain Knowledge

To represent the structure of the output and the domain knowledge, we follow the line of study on declarative knowledge representation for inference. We consider three types of generic relations between the parts: is-a indicates sub-typing of categorical predicates; disjoint-with indicates mutual exclusiveness of categorical predicates; and has-a indicates the composition of parts. These relations can be expressed by logical expressions. More specifically, is-a and has-a are mapped to implication qu ⇒ qv. disjoint-with is implemented by alternative denial ¬(qu ∧ qv). Then we derive the linear constraints from the logical expressions by the rules in [Rizzolo and Roth, 2007].

4 Experiments

We evaluate the proposed approach with several structured learning tasks: Two different entity relation extraction (ER) tasks and spatial role labeling (SpRL) task. We investigate the entity and relation recognition corpora (CoNLL04) [Roth and Yih, 2004] and ACE 2005 Corpus (ACE2005) [Li and Ji, 2014] for ER task. The two datasets contain different types of entities and relationships. For SpRL task, CLEF 2017 mSpRL dataset (SpRL2017) [Kordjamshidi et al., 2017a; Kordjamshidi et al., 2017b] is investigated.

CoNLL04. CoNLL04 [Roth and Yih, 2004] is a publicly available corpus 3 for ER. The task is to recognize four types of entities among tokens in a sentence and classify five types of relations between entities. This corpus contains 5, 516 sentences, 11, 182 entities, and 2, 048 relations. The applied hard constraints are between the types of relations and the types of their two entities [Roth and Yih, 2004].

ACE2005. The ACE dataset contains documents with annotations defined for several tasks, including Named Entity Recognition, Relation Extraction, and Event Detection and Recognition. The dataset contains seven types of entities and 45 sub-entity types. We use the same data split used in [Li and Ji, 2014]. The training set contains 10, 360 sentences each of which includes at least one entity. The test set contains 2, 637 sentences some of which may not contain any entities. The total number of entities within the sentences of the training set is 47, 406, while the testing set contains 10, 675 of them.

CLEF 2017mSpRL. The SpRL task is to identify and classify the spatial arguments of the spatial expression in a sentence [Kordjamshidi et al., 2017a]. To be specific, we identify spatial roles, including "Trajector", "Spatial indicator", and "Landmark", and detect their spatial triplet relation. We evaluated with CLEF 2017 mSpRL dataset [Kordjamshidi et al., 2017a], which has 600 sentences in the training set and 613 sentences in the testing set. The dataset is more challenging because of the complicated triplet relations and fewer examples compared to other tasks.

4.1 Experimental Setup and Results

In this section, we will show the experiment settings for each task and report the results. We implemented all the experiments using Pytorch 4 . In all the tables IML(λ) denotes the training with the combined loss and a specified λ, and the prediction without inference. IML(λ)+ denotes training with the combined loss and a specified λ and the prediction with inference. Similarly, NLL denotes the training with NLL and predictionwithout inferencewhileNLL+means trainingwith NLL and prediction with inference.

CoNLL04

We employ 768-dimensional pre-trained BERT [Devlin et al., 2019], followed by an fully-connected layer of 768dimensional hidden neurons as the token representation. For pairs based on which the relations are classified, we concatenate the representations of its two tokens and map it to a representation of the pairs by an additional fully-connected layer with 768-dimensional hidden neurons. Leaky ReLU, with a negative slope of 0.01, is used in the fully-connected layers. Independent Softmax of two classes based on the representation of the token or the pair are applied for local predictions. We handcrafted a simple ontology based on the constraints

3 https://cogcomp.seas.upenn.edu/page/resource view/43

4 Our code is available at https://github.com/HLR/ Inference-Masked-Loss

B1, B2: Baseline [Kordjamshidi et al., 2015]

Table 3: F1 score compared with baseline, on CoNLL04.

in [Roth and Yih, 2004], adding disjoint constraints among all entity types. As shown in Table 2, we consider two types of constraints for CoNLL04. Entity disjoint constraint means each token can have at most one entity type. Relation composition constraint indicates that the arguments of a relation must have entity types consistent with the relation type as defined in [Roth and Yih, 2004]. The constraints are used for global inference in both training and testing. IML(0.6) (i.e. with λ = 0.6) is used and optimized by Adam optimizer [Kingma and Ba, 2014] in batches of 8 examples. We train for 100 epochs with the learning rate of 1e - 4, which is decayed 10 times every ten epochs. We use weight decay 1e - 5 and a dropout rate of 0.35 to avoid over-fitting. Focal loss γ = 2 and label smoothing 0.01 are used for imbalanced class labels. We conducted five-fold cross-validation over all samples for the final evaluation.

The baseline [Kordjamshidi et al., 2015] used sophisticated linguistic features and structured features with a linear model and constraints [Roth and Yih, 2004]. Different training and testing settings are reported: local models (LO), Local learning and global prediction (L+I), joint training (IBT). Table 3 shows our F1-score compared with the baseline using L+I and IBT setting. Our model shows significant improvement, however, mainly because of BERT, which is capable of encoding each token with its context. We also achieve significant additional improvement by making inference after local prediction for most of the cases. The only case that inference decreases the performance is for the "Kill" relation, where local prediction performance is very high. It is connected to two "People", whose local prediction performance is rather weak, which may cause trouble for inference of "Kill". However, "Kill" gives "People" a boost. This is a trade-off between strong and weak local models.

ACE2005

We have designed a model as a baseline to experiment with both NNL and IML loss functions. The model is using BERT [Devlin et al., 2019]+FLAIR [Akbik et al.,

NER: [Strakov´a et al., 2019]

Table 5: F1-score of the ACE2005 entity recognition with hierarchical constraint on reduced training set

2018]+GLoVe [Pennington et al., 2014] for the representation of tokens. The first layer of the model is a Bi-directional LSTM layer converting the 5, 236-dimensional input of each token to a 2 × 240 representation vector. It is followed by a shared fully connected layer to a 480-dimensional representation and then classified by logistic regression for each class of entity and sub-entities. We exploit two different types of constraints for ACE05 as shown in Table 2. First, the disjoint property of different entity types. Second, the hierarchical constraints between entity types and sub-types in the dataset. The results of this experiment are listed in Table 4. Training with NLL, the model exceeds the published results [Strakov´a et al., 2019] after inference. With IML(0.5), we gain better results even without inference and get additional improvement after inference.

The IML(0.5) (with λ = 0.5) is optimized by Adam optimizer. We train for 100 epochs, with the learning rate of 0.04. We use focal loss γ = 2 and label smoothing 0.01 to deal with imbalanced class labels.

We have also compared the results of training on a smaller dataset (1000 sentences) to show whether IML improves the predictions in the low data regime. Table 5 shows the result of this experiment. In this experiment, we increased our representation space from 480 to 1000. While overall constraints helped in improving the results in the previous experiment we noticed the hierarchical constraints did not help significantly when working on the whole dataset. However, both types of constraints were helpful to achieve some improvements when training on the smaller number of examples in Table 5.

CLEF 2017 mSpRL

In this task, we firstly encode the input phrases with different linguistic features generated by SpaCy 5 , such as POS-tag, lemma, dependency path. For spatial relation extraction, we concatenate the phrase encoding of the three spatial roles. Then, we use Recurrent Neural Networks and Logistic Regression to predict spatial role labels of phrases and extract spatial triplet relations among them.

We consider only one type of constraint according to the Table 2. Spatial triplets must include the three spatial roles, which are "Trajector", "Spatial Indicator", and "Landmark". We jointly predict the spatial roles and triplets. We trained IML (0.6) (λ = 0.6) by Adam for 20 epochs with a learning rate of 0.005, weight decay of 0.001 and dropout rate of 0.5 to avoid over-fitting.

5 https://spacy.io/

B: Baseline [Manzoor and Kordjamshidi, 2018]

Figure 1: Different λ.

Table 6 shows IML improves the performance of spatial role labeling. Specifically, the inference improves the prediction results in most cases, and IML strengthens the utilization of inference compared to NLL. We use totally different parsing tools from the published result [Manzoor and Kordjamshidi, 2018], which leads to different phrases and relation candidates. Therefore, the input features are different and our results of triplet extractions which needs a lot of feature engineering are not competitive. However, we improve the spatial role extractions and IML+ consistently provides the best results for the triplet extraction compared to NLL and other variations in the scope of our own feature representations.

5 Analysis

To investigate the impact of IML and to find the optimal configuration, we conducted a series of variant experiments on EMR with the CoNLL04. For the sake of experimental time efficiency, we only evaluated using the first data split of the five-folds. All the experimental setting remains the same except those specified for each experiment.

Variation on λ. In IML(λ), λ is the parameter that controls the penalty that we apply on a false local prediction even if it can be corrected by inference. IML(0) is identical to IML while IML(1) is equivalent to NLL. The performance of the model with different λ values is shown in Figure 1. The dashed line with boxes shows the performance of local predictions, and the solid line with triangles shows the performance of global inference. Inference improves local predictions most of the time. λ = 1 (NLL) achieved a strong baseline performance. However, it is not the best. λ = 0.6 achieves the best inference result in this experiment. The performance with λ = 0 (IML) is too low to be shown properly in the chart. However, the corresponding inference performance is still competitive. λ is the only hyper-parameter proposed in this work.

Data Subsampling. We evaluate the training approach by using different variants of IML (λ) with less training data by subsampling the training set (4,413 sentences). The result is shown in Figure 2. The dashed lines with boxes are the performance of local prediction, and the solid lines with triangles are the performance of the global inference. Blue, green, and red curves show the result of different λs. In terms of inference performance, λ = 0 is rather weak compared with λ = 1 and λ = 0.6 when training with full data. For λ = 1, 0.6, the performance keeps decreasing when the data scale is reduced, while λ = 0 performs quite stable. It becomes competitive when 1/2 of the data is used and among the best is when only 1/8 of the data is available. When only 1/16 of the data is used, λ = 0 surpasses the other two settings by 9.6%. The prediction performance with λ = 0 is clearly very stable all over the time. With IML (λ = 0), the model can be trained to use local neural networks to learn minimal facts from the data and count on the constraints to get a robust global result. This is because IML decouples the structured learning task and leads to the need for fewer data.

6 Conclusion

We investigated the combination of deep neural networks with structured output by introducing the inference-masked loss. The proposed structured output learning model imposes the structure of data and domain knowledge in the form of logical constraints that describe the correlations between output variables. Inference-masked loss takes the inference into account based on the domain knowledge compiled from logical expressions. It allows local deep neural networks to make false local predictions that can be corrected by the inference. The loss helps to decompose the learning task and let the neural networks focus on local representations and make local predictions. The inference collects the local predictions based on the structure of output and domain knowledge. One advantage of IML approach is being agnostic to the inference model and treating it as a black box. This helps to plug in any inference mechanism in the loop of deep learning iterations. The proposed approach improves the generalizability of the model and robustness of training, leading to state-of-the-art results on entity relation extraction and spatial role labeling tasks, with CoNLL04, ACE2005, and CLEF 2017 mSpRL datasets, respectively. In particular, it shows improvements when there is only a small set of annotated data available.

Acknowledgments

This work is (partially) supported by the Office of Naval Research grant N00014-19-1-2308.