Moment Distributionally Robust Tree Structured Prediction
Abstract
Structured prediction of tree-shaped objects is heavily studied under the name of syntactic dependency parsing.
Current practice based on maximum likelihood or margin is either agnostic to or inconsistent with the evaluation loss.
Risk minimization alleviates the discrepancy between training and test objectives but typically induces a non-convex problem.
These approaches adopt explicit regularization to combat overfitting without probabilistic interpretation.
We propose a momentbased distributionally robust optimization approach for tree structured prediction, where the worst-case expected loss over a set of distributions within bounded moment divergence from the empirical distribution is minimized.
We develop efficient algorithms for arborescences and other variants of trees.
We derive Fisher consistency, convergence rates and generalization bounds for our proposed method.
We evaluate its empirical effectiveness on dependency parsing benchmarks.
Introduction
Structured prediction is an important learning setting for joint prediction of interdependent variables.
The output space typically consists of an exponential number of structured objects whose inherent relations can be exploited to develop efficient learning algorithms and capture key properties of data [Ciliberto et al., 2019].
Trees are widely used structures that offer expressiveness and simplicity.
We distinguish between two different tree structured prediction tasks in the literature.
The first task is a structure learning problem in graphical models [Bradley and Guestrin, 2010], aimed at constructing trees underlying a predictive model from training data.
The optimal tree is found easily with greedy algorithms for generative models [Chow and Liu, 1968], while it is NP-hard for the discriminative max-margin setting [Meshi et al., 2013].
The second task requires prediction itself to be a tree-shaped object (e g., an incidence vector).
Dependency parsing is a crucial application of this problem that has inspired a flurry of work in natural language processing.
The first-order spanning tree prediction assuming factorization over arcs can be done in O ( n 2 ) [Stanojevi´c and Cohen, 2021], whereas exact inference is NP-hard for certain (non-projective) higher-order trees (e.g., considering siblings)
[McDonald and Satta, 2007].
We study the latter in this work.
A common evaluation criterion in dependency parsing is the attachment score, namely, the score we would like to maximize on test data.
It is cost-sensitive to allow partially correct prediction.
Ideally, the training objective should be aligned with the test objective.
An early attempt to directly mimic test conditions leads to a non-convex piece-wise constant objective [Och, 2003].
Risk minimization in appropriate parametric form has a non-convex smooth objective, solvable with gradient descent, but still losing global convergence and generalization guarantees.
Maximum likelihood approaches formulate a convex smooth problem minimizing a logistic loss, consistent with conditional probability estimates but oblivious to test losses.
Maximum margin methods have convex objectives able to implicitly incorporate custom losses by scaling margins, but are known to be inconsistent with test losses generally [Nowak-Vila et al., 2021].
Unfortunately, none of these approaches yield a Bayes optimal estimator for test losses with global convergence and finite-sample generalization guarantees.
Consistent structured prediction methods include Ciliberto et al.
[2016], Blondel [2019], Nowak-Vila et al.
[2020], the latter two of which are based on Fenchel-Young losses [Blondel et al., 2020].
However, none of them have addressed the tree structured prediction problem explicitly.
For instance,
Blondel [2019] calls for Euclidean or Kullback-Leibler projection oracles, which do not exist in an efficient sense from what we know for arborescence (directed tree) polytopes.
In addition, the FrankWolfe type algorithm adopted by Nowak-Vila et al.
[2020] requires a max-min oracle and converges in a rate of O ( 1 ϵ ) .
Furthermore, all of the above methods belong to empirical risk minimization (ERM) that requires explicit regularization to combat overfitting, which can be quite vulnerable in high-dimensional settings (e.g., scarce data).
To address the above issues, we propose an estimator from first principles in distributionally robust optimization (DRO).
It minimizes the worst-case risk over an ambiguity set of distributions within bounded moment divergence from the empirical distribution.
We seek probabilistic prediction by assuming non-deterministic groundtruth labels, which, together with the ambiguity set, models uncertainty about the unknown true distribution.
We interpret the primal problem as a dual-normregularized surrogate loss minimization problem.
Note that prior art applying moment-based DRO to tree-structured graphical models [Fathony et al., 2018b] and bipartite matching
[Fathony et al., 2018a] adopts a special case of our ambiguity set in which the empirical feature moments are matched exactly and regularization has to be imposed manually.
This moment-based DRO also allows us to derive generalization bounds regarding true worst-case risks.
When the ambiguity radius is zero, the DRO estimator is shown to be consistent.
We develop two practical algorithms, one based on game theory and the other based on marginal probabilities of tree parts.
We further propose efficient
Euclidean projection oracles onto the arborescence polytope with linearly convergent guarantees.
We conduct experiments on three common dependency parsing datasets, suggesting that our method is particularly effective with little training data.
Contributions.
Our contributions are summarized as follows.
(1) We propose a distributionally robust tree structured prediction method and show its equivalence to regularized surrogate minimization.
We derive its generalization bounds and consistency.
(3) We propose efficient algorithms based on projection oracles for arborescence polytopes.
(4) We perform empirical study on real-world datasets.
Paper structure.
We begin with problem setup and existing work in Section 2.
We present our method with theoretical analysis in Section 3.
Section 4 proposes efficient projection oracles.
Sectiondiscusses extensions beyond first-order directed trees.
Experimental results of comparing our method with a competitive baseline are given in Section 6.
We conclude the paper in Section 7.
Background and Related Works 2.1 Tree Structured Prediction
Consider a weighted directed multi-graph G =
( V , E ) where each arc ( i,
j, l ) ∈ E from node i to j has a label l .
By designating a root node r ∈ V , we say that A ⊆
E is an r -arborescence of G if ( V , A ) is a directed spanning tree rooted at r .
For any v ∈ V , denote by δ − ( v ) :
= { ( i, j
, l ) ∈ E : j = v } the set of its incoming arcs, and δ + ( v ) :
= { ( i, j
, l ) ∈ E : i = v } the set of its outgoing arcs.
Let X be the input space and Y ≜ (cid:83) x ∈X Y ( x ) be the output space where Y ( x ) represents the set of rarborescences of a graph G ( x ) formed by x .
Dependence on x is suppressed when context is clear.
Let
R ⊆
2 E be a set of parts with E ⊆ R .
Each part s ∈ R is a subset of arcs.
It is convenient to represent y ∈ Y as a binary vector with y s = 1 iff part s appears in y .
Let w θ ( x , y ) ≜ (cid:80) s ∈R w θ ( x , y s ) be a score function decomposing over parts, parameterized by θ .
Let { ( x ( i ) , y ( i ) ) } mi =1 be a set of m training examples drawn i.i.d.
from a distribution P ∈ P ( X × Y ) , where each y ( i ) is an r -arborescence.
The goal of tree structured prediction is to learn a function h :
X → Y from training data.
Assume that the evaluation criterion is a loss function ℓ :
Y × Y → R ≥ 0 .
We introduce existing methods in the setting of (graph-based, non-projective, syntactic) dependency parsing where x is a sequence of tokens and G ( x ) encodes dependencies among tokens.
2.2 Maximum Likelihood
A probabilistic modeling approach based on exponential family distributions maximizes the conditional log-likelihood of the training data:
[Equation]
where Z ( x ) ≜ (cid:80) y ∈Y ( x ) exp ( w θ ( x , y ))
This problem is convex for log-linear models, but intractable for general R [Koller and Friedman, 2009].
The first-order arc-factored model ( R = E ) is equivalent to a loop-free factor graph, rendering it tractable via the matrix-tree theorem [Kirchhoff, 1847, William, 1984, Koo et al., 2007, McDonald and Satta, 2007, Smith and Smith, 2007].
Neural parsers either leverage the same theorem to compute the partition function [Ma and Hovy, 2017] or consider the parent node distribution independently for each node by local normalization [Dozat and Manning, 2017, Zhang et al., 2017].
Higher-order models require approximate algorithms such as loopy belief propagation [Murphy et al., 1999] and Markov chain Monte Carlo [Brooks, 1998].
This approach does not incorporate task-specific losses.
In fact, with maximum a posteriori (MAP) decoding, it is not consistent with any specific loss in general [Nowak-Vila et al., 2019].
2.3 Maximum Margin
An alternative approach based on maximum margin Markov networks [Taskar et al., 2003] or structured support vector machines [Tsochantaridis et al., 2005] optimizes a hinge-type surrogate:
[Equation]
which inspires a rich line of work based on MAP inference with manual features [Taskar et al., 2004,
McDonald et al., 2005, McDonald and Pereira, 2006, Martins et al., 2009, 2010, 2015, Zhang et al., 2014] or deep learning [Kiperwasser and Goldberg, 2016, Wang and Chang, 2016].
Approximate
MAP inference is required for models beyond first-order.
A smooth variant called softmax-margin [Gimpel and Smith, 2010] incorporates the task-specific loss ℓ but still implicitly minimizes it.
Margin-based objectives are known to be consistent only under very restrictive conditions [Liu, 2007,
Nowak-Vila et al., 2021] (i.e., data with majority label, loss being a distance).
2.4 Minimum Risk
Empirical risk minimization suggests directly optimizing the expected target loss on training data:
[Equation]
which is commonly non-convex due to normalization of p θ .
There are a few parsers optimizing this objective via back-propagation [Stoyanov and Eisner, 2012], k -best lists [Smith and Eisner, 2006], semirings [Li and Eisner, 2009, Zmigrod et al., 2021] and other differentiable approximations [Gormley et al., 2015, Mensch and Blondel, 2018].
Local optima found by these algorithms do not satisfy the premise of Fisher consistency and make it difficult to quantify generalization errors.
2.5 Distributionally Robust Optimization
Distributionally robust optimization has attracted emerging interests in improving machine learning models due to its connections to robustness, regularization and generalization.
It proposes to minimize a risk with respect to the worst-case distribution chosen by an adversary in some uncertainty set:
[Equation]
where B is an ambiguity set that can be defined by discrepancies [Shafieezadeh-Abadeh et al., 2019, Duchi and Namkoong, 2019], moments [Delage and Ye, 2010, Farnia and Tse, 2016], shapes [Popescu, 2005, Hanasusanto et al., 2015] and kernels [Shang et al., 2017, Staib and Jegelka, 2019].
A thorough review can be found in Rahimian and Mehrotra [2019].
We focus on moment-matching discriminative approaches while a similar generative method is proposed in Ganapathi et al.
Method
We introduce the formulation, followed by practical algorithms for learning and inference.
Afterwards, we present the theoretical guarantees.
We defer all proofs to Appendix A.
3.1 Formulation
We assume that the evaluation criterion is the Hamming loss ℓ ( y , y ′ ) :
= (cid:80) i 1 ( y i ̸
= y ′ i ) with( · ) being the 0 - 1 indicator function, but the results in this paper generalize to losses with affine decomposition [Ramaswamy et al., 2013] easily.
Let P true be the true distribution and P emp be the empirical distribution.
Our approach builds upon a probabilistic predictor that non-parametrically minimizes the expected loss with regard to the most adverse distribution in an uncertainty set where the distributions are ε away from the empirical distribution in terms of feature moment difference:
[Equation]
(1) where B ( P emp ) :
= { Q : Q X
=
P emp X ∧∥ E P emp ϕ ( · ) − E Q ϕ ( · )
∥ ≤ ε } with ε ≥ 0 and ϕ : X ×Y → R d is a joint feature mapping decomposable over parts: ϕ
( x , y ) ≜ (cid:80) s ϕ
( x , y s ) .
In Farnia and Tse [2016], cross-moments are adopted: ϕ ( x , y ) :
= ϕ X ( x ) ⊗ ϕ Y ( y ) where ⊗ is the tensor product.
By Fenchel duality [Altun and Smola, 2006] and strong duality [Von Neumann and Morgenstern, 1947], we show that Eq.
(1) is analogous to dual-norm-regularized surrogate loss minimization:
Proposition 1.
The distributionally robust tree structured prediction problem based on moment divergence in Eq.
(1) can be rewritten as
[Equation]
where θ ∈ R d is the vector of Lagrangian multipliers and ∥·∥ ∗ is the dual norm of ∥·∥ .
3.2 Constraint Generation Solution
From a game-theoretic rationale [Topsøe, 1979, Grünwald and Dawid, 2004], Eq.
(1) is considered as an adversary-constrained zero-sum game.
A prediction player chooses a set of stochastic strategies (conditional distributions over arborescences) in order to minimize the expected payoff whereas an adversarial player chooses constrained strategies to maximize it.
The payoff for a pair of pure strategies is the incurred loss, ℓ (ˆ y , ˇ y ) .
The constrained game is transformed to a set of unconstrained ones in Eq.
(2) whose payoffs are parameterized by θ : payoff (ˆ y ,
ˇ y ) ≜
ℓ (ˆ y , ˇ y )
+ θ ⊺ ϕ
( x , ˇ y ) .
Note that the games in Eq.
(1) are jointly constrained for all x ’s in the support of P emp X while the ones in
Eq.
(2) are conditionally independent given x .
The unconstrained game can be solved by a linear program [Von Neumann and Morgenstern, 1947].
However, there are O ( n n ) spanning trees in a complete graph, thus making explicit construction of the full payoff matrix impractical.
We adopt a constraint generation algorithm named double oracle [McMahan et al., 2003], shown in Appendix B.
It builds a payoff sub-matrix starting from small initial sets of strategies.
In each iteration, each player takes their turn based on the game payoff sub-matrix by finding the best response among all possible strategies to the opponent’s optimal mixture strategies.
The response is added to a player’s strategy set if it improves the value of the game, with the sub-matrix updated.
The algorithm terminates and converges to a Nash equilibrium of the original game when the strategy sets no longer grow.
The size of the final sub-matrix is usually small in practice but there are no known theoretical guarantees, thus no way to analyze the convergence behavior.
Finding the best response requires an oracle, equivalent to finding the minimum weight arborescence.
The objective in Eq.
(2) is a convex function of θ , so we can optimize it with sub-gradients based on solutions of the inner zero-sum games.
Although lacking convergence guarantees, this algorithm is flexible with custom losses and provides a game-theoretic perspective to a typical DRO problem.
3.3 Marginal Distribution Formulation
The r -arborescence polytope is defined as the convex hull of all vectors representing r -arborescences:
A arb ( x ) :=
Conv ( { y ∈ R |R| : y ∈ Y ( x ) } ) .
Note that each p ∈
A arb is a convex combination of all r -arborescences: p ≜ (cid:80) y Prob ( y ) y , where p s denotes the marginal probability of part s .
Here we adopt the squared ℓ 2 norm as the dual norm and an ambiguity radius of ε = λ/ 2 .
By substituting the marginal probability vectors and switching min-max optimization orders, we simplify Eq.
(2) into
[Equation]
where Φ ( i ) ∈ R |R|× d denotes the feature matrix of the i -th training data, µ ∈ R ≥ 0 is a smoothing parameter to induce strong convexity.
We push the maximization over q to the outermost level because of its large computational cost.
If µ = 0 , the solution to Eq.
(3) is also optimal to Eq.
(2) by strong duality
but the problem becomes non-smooth.
Therefore we expect θ ∗ obtained with a very small positive µ to be a good approximation of θ ∗ obtained with µ = 0 .
To optimize it, with fixed q , due to strong convexity, the unconstrained minimization over θ yields
[Equation]
− p ( i ) emp ) .
In contrast, the constrained minimization over p admits no closed-form solution but can be cast as Euclidean projection onto A arb instead, independently for each i ∈ [ m ] : p ∗ =
min p ∈A arb ∥ p − 1 µ q ( i ) ∥ 22 ≜ Proj
A arb ( 1 µ q ( i ) ) .
Given θ ∗ and p ∗ , the outermost maximization can be solved by a projected quasi-Newton algorithm
[Schmidt et al., 2009] that also requires the projection oracle Proj
A arb ( · ) , elaborated in Section 4.
3.4 Inference
We propose two algorithms to make inference with given θ ∗ .
Weight construction.
Construct the part weights as Φ θ ∗ ∈ R |R| and find the maximum weight arborescence:
y ∗ ∈ arg max y y ⊺
Φ θ ∗ by the Gabow-Tarjan (GT) algorithm
[Gabow et al., 1986,
Zmigrod et al., 2020] or approximate methods for higher-order trees.
Minimum Bayes risk decoding.
The optimal probabilistic prediction P ∗ or p ∗ can be obtained from Eq.
(2) or Eq.
The marginal probabilities enable minimum Bayes risk decoding: y ∗ ∈ arg min y E P
∗ ˆ Y | x ℓ ( y , ˆ Y ) ≜ arg max y (cid:80) s : y s =1 p ∗ s , a maximum weight arborescence problem.
3.5 Statistical Properties
Basic generalization bounds of DRO methods derived from measure concentration are not appropriate for an ambiguity set defined by low-order moments in this paper since it fails to converge [Shafieezadeh-Abadeh et al., 2019].
We take an alternate approach following Farnia and Tse [2016] to obtain excess out-of-sample risk bounds by assuming boundedness on features and losses.
Theorem 2.
Given m samples, a non-negative loss ℓ ( · , · ) such that | ℓ ( · , · ) | ≤ K , a feature function ϕ ( · , · ) such that ∥ ϕ ( · , · ) ∥ ≤ B , a positive ambiguity level ε > 0 , then, for any ρ ∈ (0 , 1] , with a probability at least− ρ , the following excess true worst-case risk bound holds:
[Equation]
where θ ∗ emp and θ ∗ true are the optimal parameters learned in Eq.
(2) under P emp and P true respectively.
The original risk of θ under Q is R L Q ( θ ) :
=
E Q X , Y , P θ
ˆ Y | X ℓ ( ˆ Y , Y ) with Bayes prediction P θY | x ∈ arg min
[Equation]
Theorem 2 presents a bound based on uniform convergence and Rademacher complexities [Bartlett and Mendelson, 2002], which improves the results in Asif et al.
[2015], who merely show that the worst-case risk upper bounds the risk under any distribution in the ambiguity set.
The dual problem in Eq.
(2) suggests an adversarial surrogate loss ℓ adv ( θ , ( x , y )) in a ERM form.
The special case of ε = 0 in our DRO estimator has a similar form to the max-min surrogate loss in
Nowak-Vila et al.
[2020] except that we assume probabilistic prediction.
A conclusion of its Fisher consistency can thus be drawn based on Fathony et al.
[2018a], Nowak-Vila et al.
Corollary 3.
When ε = 0 , ℓ adv is Fisher consistent with respect to ℓ .
Namely, P θ ∗ true ˆ Y | X is the probabilistic prediction made by the Bayes optimal decision rule, where θ ∗ true is defined in Theorem 2.
If ε > 0 , the decoded prediction for each x will not belong to the convex hull of true conditional distributions, thus not a minimizer of ℓ .
On the other hand, if ε is chosen as m − α for 0 < α < 1 / 2 , ℓ adv will be universally consistent according to the comparison inequality in Nowak-Vila et al.
Projection onto Arborescence Polytopes
The Euclidean projection onto an r -arborescence polytope is a quadratic programming problem 1 :
[Equation]
We focus on first-order models and discuss the extensions to other classes of trees in Section 5.
4.1 Frank-Wolfe Algorithm
The Frank-Wolfe (FW) method [Frank et al., 1956] is an iterative first-order algorithm that enforces constraints by optimizing a linear objective over the feasible set at each iteration t :
[Equation]
(4) which is a minimum weight arborescence problem with weights ∇ f ( x t ) in our case.
The solution is updated and stays feasible: x t +1 ← x t + γ t ( s t − x t ) , where γ t is a step size typically set to 2 t +2 .
FW style algorithms are known to have a convergence rate of O ( 1 ϵ )
[Jaggi, 2013].
4.2 Martin’s Polytope
A compact representation of A arb with a polynomial number of linear constraints is attractive to lead to efficient algorithms.
To the best of our knowledge, there is no existing projection method exploiting special structures of this polytope.
An extended formulation of the arborescence polytope [Friesen, 2019, Martin, 1991] follows a lift-and-project approach.
It relates each element to existence of k -arboresences of the underlying undirected graph for all k ∈ V .
We extend it to multi-graphs:
[Equation]
where z r ∈ R |E| is associated with the original arcs E , z k ∈ R |E ′
| for k ̸
= r is associated with a simple directed graph ( V , E ′ ) formed by removing directions and splitting each edge { i, j } into two directed ones, E ij :=
{ a ∈ E : ¯ a = { i, j }} is the set of arcs connecting i and j with ¯ a ≜ ( i, j, l ) :
= { i, j } denoting the underlying undirected edge.
We show exact correspondence between A marb and
A arb based on a similar argument for simple graphs [Friesen, 2019]:
Proposition 4.
Let G be a multi-graph.
A marb ≜
A arb .
To solve min x ∈A marb ∥ x − w ∥
[Equation]
we propose to adopt the alternating direction method of multipliers (ADMM) and rewrite it into the following separable form:
[Equation]
where I U ( · ) is the characteristic function with I U ( x )
=
if x ∈ U and ∞ otherwise.
Let λ ′
k be the dual variables and λ k :=
ρ k λ ′ k .
The scaled augmented Lagrangian function is
L
[Equation]
u ,
[Equation]
The ADMM algorithm updates the parameters as follows:
[Equation]
This decomposes the original projection problem into simpler projection problems.
Projection onto
U k for k = r decomposes over j ∈ V into |V| projections onto simplex, solvable as fast as O ( n ) in the worst case [Condat, 2016].
For k ̸ = r , computation of u t +1 k can be done in parallel.
The
Lagrange dual problem of Proj U k ( · ) can be written as
[Equation]
where w ij := (cid:80) a ∈E ij w a , n ij :=
|E ij | , α ij :=
min( α i , α j ) and α k := + ∞ .
Strong duality holds by linear constraint qualification.
Primal solutions are recovered by x ∗ a = w a − min( α ∗ ¯ a / 2 , w ¯ a /n ¯ a ) .
Convergence.
The dual objective of Proj U k ( · ) is strongly concave on { α ∈ R |V| : ∀ i ∃ j { i, j } ∈
E ∧ α i ≤ α j ∧ α i ≤ 2 w ij /n ij } , with a unique global maximizer.
This implies fast convergence in practice given good initialization.
The negative Lagrange dual function has restricted strong convexity with ν = min ij ( n ij / 2) , near the optimum, suggesting linear convergence
[Zhang and
Cheng, 2015].
Alternatively, exact solutions can be found by enumerating rankings (with duplicates) of α in O ( |V| |V| ) .
In this manner, the ADMM algorithm with a strongly convex objective has a linear convergence rate O (log 1 ϵ ) with either exact [Deng and Yin, 2016] or linearly convergent approximate solution [Hager and Zhang, 2020] of Proj U k
Using Nesterov’s accelerated gradient algorithm
[Nesterov, 2003] to optimize Eq.
(3) leads to iteration complexity
O ( C log 1 ϵ ) with constant
C dependent on Lipschitz constants of gradients and µ .
Extensions 5.1 Undirected Spanning Trees
An straight-forward way of extending to undirected spanning trees is to split { i, j } into two arcs ( i, j ) , ( j, i ) and make the feature mapping direction-invariant, i.e., ϕ ( x , y s ) =
ϕ ( x , y s ′ ) for s and s
′ having the same underlying undirected graph.
We post-process the prediction by removing directions.
Alternatively, we seek projection oracles for undirected graphs.
Projection via FW is done by using any minimum spanning tree algorithm in Eq.
For ADMM, the formulation in Martin [1991] is originally for undirected trees:
A mund := { x : ∃ z ≥ 0 (cid:80) a ∈ δ − ( j ) z ka = 1 ( j ̸ = k ) ∧
z kij + z kji
=
x
{ i,j } ∀ k, i, j ∈ V} .
ADMM is easily adapted to this case with (cid:80) a ∈E ij x a replaced by x { i,j } .
5.2 Dependency Trees
The spanning tree structure in dependency parsing is a special one where the outdegree of root is restricted to be one.
We can use the GT algorithm for inference with either the same training objective or an aligned objective where a dependency tree polytope is considered: A dep ( x ) :=
Conv ( { y ∈
Y ( x ) : | δ + ( r ) | = 1 } ) .
A straightforward extension of A marb to characterizing dependency trees is
A mdep := { z r : z r ∈
A marb ∧ (cid:80) a ∈ δ + ( r ) z ra = 1 } , equivalent to A dep by the following proposition:
Proposition 5.
Let G be a multi-graph.
A mdep ≜ A dep .
FW methods leverage the GT algorithm in Eq.
As for ADMM, the dual problem of projection onto U ′ r :=
{ x : x ∈ U r ∧ (cid:80) a ∈ δ + ( r ) x a = 1 } becomes
[Equation]
where γ ( i,j,l ) := α j + 1 ( i = r ) β .
This can be solved in O ( |E| log |E| )
[Zhang et al., 2010].
5.3 Higher-order Polytope
Compact higher-order polytope descriptions exist for undirected spanning trees but are still unknown for arborescences with even one monomial [Friesen, 2019].
FW requires a linear oracle that is
NP-hard to solve exactly in higher-order settings [McDonald and Pereira, 2006].
Instead, we can approximate it with a local polytope where the marginal probabilities of each part s is required to be locally consistent with that of each arc a .
For simplicity, we consider only features for the all-true assignments, i.e., all arcs exist in part s .
The resulting polytope can be written as
A mloc := { x : x E ∈
A marb ∧ ∀ s ∈ R , a ∈ s p s ≤ p a } , which suggests an ADMM algorithm with additional constraint sets for each part: U s :=
{ x ∈ R |R|≥ 0 : x s ≤ x a ∀ a ∈ s } , the projection onto which can be done in O ( | s | log | s | ) .
See Appendix D for details.
Experiments
We evaluate our proposed method on dependency parsing tasks and compare its ability to BiAF [Dozat and Manning, 2017], arguably the state-of-the-art neural dependency parser.
We implement our methods in Python and C 2 .
We leverage the implementations in SuPar
[Zhang et al., 2020] for the baseline.
All experiments are conducted on a computer with an Intel Core i7 CPU (2.7 GHz) and an NVIDIA Tesla P100 GPU (16 GB).
We adopt three public datasets, the English Penn Treebank (PTB v3.0)
[Marcus et al., 1993],
the
Penn Chinese Treebank (CTB v5.1) [Xue et al., 2002] and the Universal Dependencies (UD v2.3)
[Nivre et al., 2016].
See Appendix C for data-processing details.
Representation learning is not the focus of this paper.
We follow Levy et al.
[2020] and compare our method with the last biaffine classification layer in BiAF on top of pretrained features preceding this layer (backbone’s output).
The pretrained embeddings produced by complicated non-linear models make Fisher consistency’s assumption of optimizing over all measurable functions less violated.
To featurize the data, for each dataset, we train a BiAF network with the whole training set to obtain a pretrained model.
Note that this may create unfair advantages for the baseline because the last layer was optimized together with the backbone network in an end-to-end manner during pretraining.
Moreover, pretraining uses a standard ERM objective with the cross-entropy loss and local normalization over head nodes.
The pretrained features are thus more adequate for the ERM objective than for our DRO objective.
To make use of the features as inputs in our method, we take the outer product of the embedding vectors for two nodes as the arc feature vector.
Our method and the biaffine layer therefore share the same number of parameters ( 501 × 501 , including bias terms).
We focus on predicting the unlabeled dependency tree while relying on pretrained models for relation label prediction.
The evaluation criteria are the labeled/unlabeled attachment scores (LAS/UAS) and
[Figure] [Figure]
labeled/unlabeled complete matches (LCM/UCM).
The attachment score can be transformed to the
Hamming loss with linear mapping: AS ( y , y ′ ) ≜ |V| − 1 − ℓ ( y , y ′ ) / 2 .
Full batch learning is adopted for Marginal (Eq.
Mini-batch training is adopted for Game , the game-theoretic algorithm, and Stochastic , which solves the inner min-max problem in Eq.
(2) using
Eq.
(3) with fixed θ .
All models are trained with the training set only.
The optimal hyperparameters and parameters are chosen based on the validation set.
See Appendix C for detailed parameter values.
To showcase the ability of DRO methods tackling scarce data, in each run, we randomly draw m ∈ { 10 , 50 , 100 , 1000 } samples without replacement from the training set and keep the original validation and test sets.
All the models are trained on the same set of sampled data.
The process is repeated 5 times for each m .
The main UAS results on the PTB, CTB and UD Dutch Lassy
Small datasets are reported in Table 1 with complete results provided in Appendix C. Our methods consistently deliver higher UAS than BiAF especially with a small amount of data 4 .
With little training data, DRO approaches minimize the worst-case risk to avoid overfitting.
With more training data available, our method is still comparable to BiAF which is not significantly better than our methods by statistical tests.
This illustrates the advantages of replacing conditional log-likelihood with our
Fisher consistent surrogate loss without changing the number of model parameters.
Moreover, we study a low-resource setting with the UD Turkish dataset in which only the sampled data is used for pretraining without BERT embeddings.
The binary cross-entropy loss (single normalization) is adopted during pretraining in this setting to avoid pretrained features biased towards the multi-class cross-entropy loss (local normalization) adopted by BiAF .
We observe consistently competitive performance of our methods in the low-resource setting in Table 1 as well.
We report computational time of one gradient descent step in the second column of Table 1, averaged across 10 runs.
For fair comparisons, all the models are run with CPU only, with a batch size of 200 .
All the methods achieve their optimal validation set performance in 150 - 300 steps.
BiAF and
Marginal are the fastest because the most time-consuming step of computing dot products of features and parameters is only performed once whereas the other two methods perform it multiple times.
However, since Marginal is unable to leverage stochastic gradients, its execution time grows linearly in the full batch size.
Henceforth, there is a trade-off between Marginal and Stochastic / Game for computational efficiency.
The extra cost compared to BiAF with cross entropy is expected because distributional robustness against a set of adversarial distributions is guaranteed.
We compare ADMM and FW by performing for 100 times projection of random points in [ − 5 , 5]on a graph with 5 nodes and 3 parallel arcs between each ( i, j ) .
We subtract the integral part of the observed minimum values in each run for better illustration.
As shown in Figure 1, ADMM usually finds a better solution in the arborescence polytope than FW does within 1000 iterations 5 .
That being said, the per-iteration cost of ADMM is about 8 n times higher than that of FW due to consensus optimization of n subproblems.
In practice, the solution computed with FW usually leads to an approximately good sub-derivative to optimize the DRO objective.
We have verified that the solutions suggested by ADMM satisfy the polytope constraints for graphs of up to 10 nodes.
We conduct sensitivity analysis by varying µ and λ on UD Dutch with 100 training samples.
Figureimplies that moderate smoothing is beneficial to generalization.
The ambiguity radius should be judiciously chosen because a small λ causes overfitting while a large λ leads to conservative models.
Discussion and Conclusion
We proposed a distributionally robust and consistent tree structured prediction method.
We showed its equivalence to regularized surrogate loss minimization.
We put forward a provably convergent algorithm based on efficient projection oracles for arborescence polytopes.
Our proposed method enjoys Fisher consistency and robustness against noise in conditional distributions in terms of feature moments.
Theoretical and empirical results validate its effectiveness.
We assume that an expressive feature mapping is given such that a sufficiently good linear discriminant rule can be learned.
The class-sensitive form ϕ ( x , y ) is general but consumes more memory than the decomposable form ϕ X ( x ) ⊗ ϕ Y ( y ) .
The ADMM projection algorithm is efficient theoretically with high per-iteration costs in practice.
We expect this work to be a principled way of learning to predict tree-structured objects.
Future directions include a more efficient implementation and general structured prediction with DRO.
Potential negative societal impacts of our work include using its prediction without verification to guide human-centered design in policy-making.
Representation learning.
Our method can be easily adapted to a representation learning framework with automatic differentiation.
Although this may lead to a non-convex problem without the theoretical guarantees derived in this paper, it is highly desired in practice if feature mappings are optimized as well.
We discuss a possible approach as follows.
Modern neural networks for supervised learning typically have a linear layer in the end without activation.
Assume the penultimate layer outputs Φ ( x ) for input x , the last layer with parameters θ will typically output ψ ( x ) :
= Φ ( x ) θ ∈ R k , sometimes called logits, with k = n 2 labels for all arcs when parsing a sentence of n tokens.
Note that θ in our formulation naturally serves as the parameters of this linear layer.
Moreover, knowing ψ ( x ) is sufficient for us to solve the inner minimax problem in Eq.
(2) to get P ∗ ˆ Y | x and Q ∗ ˇ Y | x .
In this way, our DRO method can be considered a loss layer without learnable parameters, which backpropagates the sub-derivative of the objective with respect to ψ ( x ) :
[Equation]
where B is the batch size.
The sub-derivative of the regularization term with respect to θ should be added to the linear layer.
Now we are able to take advantage of automatic differentiation and focus on solving the inner adversarial problem given ψ ( x ) and y .
Since the computational bottleneck lies in computing ψ ( x ) and backward passes, the overhead of computing the adversarial loss may be dominated and not significant compared to the cross-entropy loss.
We leave investigations on its effective applications to future work.
Acknowledgments and Disclosure of Funding
This material is based upon work supported by the National Science Foundation under Grant Nos.
1652530, 1910146, and 1934915.
Checklist
For all authors... (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope?
[Yes] (b) Did you describe the limitations of your work?
[Yes] (c) Did you discuss any potential negative societal impacts of your work?
[Yes] (d) Have you read the ethics review guidelines and ensured that your paper conforms to them?
[Yes] 2.
If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results?
[Yes] (b) Did you include complete proofs of all theoretical results?
[Yes] 3.
If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
[Yes] (b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?
[Yes] (c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[Yes] (d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?
[Yes] 4.
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?
[Yes] (b) Did you mention the license of the assets?
[No] (c) Did you include any new assets either in the supplemental material or as a URL?
[No] (d) Did you discuss whether
and how consent was obtained from people whose data you’re using/curating?
[N/A] (e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content?
[N/A] 5.
If you used crowdsourcing or conducted research with human subjects... (a) Did you include the full text of instructions given to participants and screenshots, if applicable?
[N/A] (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?
[N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?
[-where I C ( · ) is the indicator function with I C ( x ) = 
-]
[N/A]