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?
[N/A]
A Technical Proofs
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 ∥·∥ .
Proof.
Recall the primal problem
[Equation]
where B ( P emp ) :
= { Q : Q X
=
P emp X ∧ ∥
E P emp ϕ ( · ) − E Q ϕ ( · )
∥ ≤ ε } with ε ≥ 0 .
Note the feature function ϕ ( · ) is fixed and given.
Since P ˆ Y | X ∈ ∆ and Q X , ˇ Y ∈ ∆ ∩ B ( P emp ) where ∆ is the probability simplex with dimension omitted, the constraint sets are convex.
The objective function is convex in P and concave in Q because it is affine in both.
Therefore strong duality holds:
[Equation]
Let C := { u : ∥ u − E P emp ϕ ( · ) ∥ ≤ ε } .
Rewrite the problem with this constraint:
[Equation]
where I C ( · ) is the indicator function with I C ( x ) =
if x ∈ C and + ∞ otherwise.
The simplex constraints are omitted.
The dual problem by relaxing the equality constraint is
[Equation]
where θ is the vector of Lagrange multipliers.
Given X = x , optimization of Q
ˇ Y | x and P ˆ Y | x can be done independently.
Again by strong duality, we can rearrange the terms:
[Equation]
The associated dual norm ∥·∥ ∗ of the norm ∥·∥ is defined as
[Equation]
based on which we are able to simplify the optimization over u as
[Equation]
Plugging it back to the dual problem, we have
[Equation]
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]
ln(4 /ρ ) (cid:33) , 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 P max Q E Q
ˇ Y | x P ˆ Y | x ℓ ( ˆ Y , ˇ Y )
+ θ ⊺ ϕ
( x , ˇ Y ) .
Proof.
Define the adversarial surrogate risk of θ with respect to ˜ P as
[Equation]
Let θ ∗ true ∈ arg min θ R S P true ( θ ) and θ ∗ emp ∈ arg min θ R S P emp ( θ ) be the optimal parameters learned with
P true X , Y and P emp X , Y respectively.
Given x , define the decoded prediction by θ as
[Equation]
Let the original risk of loss ℓ under some distribution Q be
[Equation]
According to Proposition 1, for any fixed P , we have similarly
[Equation]
We start by looking at the worst-case risk of θ ∗ true and θ ∗ emp .
[Equation]
where the last inequality holds because θ ∗ emp is not necessarily a minimizer
[Equation]
Similarly for θ ∗ true ,
[Equation]
On the other hand,
[Equation]
where the first equality holds according to the definition of θ ∗ true .
The above two inequalities imply the equality:
[Equation]
Therefore,
[Equation]
The main idea is thus to use uniform convergence bounds.
Firstly, by substituting Q = P true , note that
[Equation]
We can get an upper bound of the norm of any optimal solution θ ∗ true or θ ∗ emp as follows:
[Equation]
Let ψ ( X , Y ) :
= θ ⊺ ϕ
( X , Y ) and ψ x :=
( ψ ( x , y ))
y ∈Y .
Define
[Equation]
Let q x ∈ ∆ be the probability vector of Q
ˇ Y | x and e y be the standard basis vector with y -th entry equal to 1 .
We have that for any ( x , y ) ,
[Equation]
where δ ( x , y ) is the Dirac point measure.
g ( · , ˜ P ) is therefore 2 -Lipschitz with respect to the ℓ 1 norm.
As per the assumption, ∥ ϕ ( · , · ) ∥ ≤ B .
This further implies that
[Equation]
We then follow the proof of Theorem 3 in Farnia and Tse [2016].
According to Theorem 26.12 in
Shalev-Shwartz and Ben-David [2014], by uniform convergence, for any ρ ∈ (0 , 2] , with a probability at least 1 − ρ 2 ,
[Equation]
According to the definition of θ ∗ true , the following inequality holds:
[Equation]
Since θ ∗ true do not depend on samples, according to the Hoeffding’s inequality, with a probability 1 − ρ/ 2 ,
[Equation]
Applying the union bound to the above three inequations, with a probability 1 − ρ , we have
[Equation]
As stated by Inequation (5), we conclude with the following excess risk bound:
[Equation]
Corollary 3.
When ε = 0 , ℓ adv is Fisher consistent with respect to ℓ .
Namely,
[Equation]
where θ ∗ true is learned with ℓ adv and P true as in Theorem 2.
Proof.
Our formulation differs from Nowak-Vila et al.
[2020] in the fact that we allow probabilistic prediction to be ground truth.
By defining y ∗ ( µ ) as the gold standard probabilistic prediction and Y as the set of all possible probabilistic predictions in Proposition C.2 in Nowak-Vila et al.
[2020], we have
[Equation]
Therefore,
[Equation]
Proposition 4.
Let G be a multi-graph.
A marb ≜
A arb .
Proof.
We follow the proof of Friesen [2019] for simple graphs.
Recall the definition of A marb :
[Equation]
On one hand, given a legal r -arborescence with characteristic vector z r , Eq.
(6) and Eq.
(7) hold by the definition of arborescences.
The equality also holds for a convex combination of the characteristic vectors of r -arborescences.
On the other hand, given z ∈ A marb .
Consider Edmond’s definition of r -arborescence polytope based on rank constraints:
[Equation]
We have Eq.
(6) directly implies Eq.
According to Eq.
[Equation]
Therefore, (cid:88)
[Equation]
which is exactly Eq.
Proposition 5.
Let G be a multi-graph.
A mdep ≜ A dep .
Proof.
Recall the definition of A mdep :
[Equation]
Input: Lagrange multipliers θ ; feature function ϕ ( · , · ) ; initial set of trees { y initial }
Output: A sparse Nash equilibrium ( ˆ T , ˇ T , P , Q )
Initialize ˆ T ← ˇ T ← { y initial } repeat ( P , ˆ v Nash ) ← SolveZeroSumGame ˆ T ( ℓ, θ ⊺ ϕ , ˆ T , ˇ T ) (ˇ y BR , ˇ v BR ) ← FindBestResponse ( ℓ, θ ⊺ ϕ , P , ˆ T ) if ˆ v Nash ̸
= ˇ v BR then ˇ
T ← ˇ
T ∪ { ˇ y BR } end if ( Q , ˇ v Nash ) ← SolveZeroSumGame ˇ T ( ℓ, θ ⊺ ϕ , ˆ T , ˇ T ) (ˆ y BR , ˆ v BR ) ← FindBestResponse ( ℓ, θ ⊺ ϕ , Q , ˇ T ) if ˇ v Nash ̸
= ˆ v BR then ˆ
T ← ˆ
T ∪ { ˆ y BR } end if until ˆ v Nash = ˇ v BR =
ˇ v Nash
= ˆ v BR return ( ˆ T , ˇ T , P , Q )
On one hand, given a legal dependency tree z r ∈
A dep , it satisfies Eq.
(6) and Eq.
(7) by Proposition 4.
It also satisfies Eq.
(10) by the definition of A dep .
On the other hand, given z r ∈ A mdep , firstly, z r must be in A arb by Proposition 4, which implies that we can write it as a convex combination of k r -arborescences vectors: z r ≜ α 1 t 1 + α 2 t 2 + · · · + α k t k .
All of them are legal r -arborescences, so (cid:80) a ∈ δ + ( r ) t ia ≥ 1 for all i ∈ [ k ] .
Now if (cid:80) a ∈ δ + ( r ) t ia >for some i , we would have a contradiction, (cid:80) a ∈ δ + ( r ) z ra > 1 .
B Algorithm Details
The pseudo-code of the constraint generation algorithm proposed in Section 3.2 is illustrated in
Algorithm 1.
C More on Experiments
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], the Dutch Lassy Small Treebank and the Turkish
Treebank in Universal Dependencies (UD v2.3)
[Nivre et al., 2016].
We follow conventions in Chen and Manning [2014], Dyer et al.
[2015] to prepare our data.
We make standard train/validation/test splits.
We use Stanford Dependencies (SD v3.3.0)
[De Marneffe and Manning, 2008] to convert dependencies in PTB and CTB.
The predicted POS tags with Stanford POS tagger [Toutanova et al., 2003] are adopted for PTB whereas gold POS tags are adopted for CTB and UD.
Punctuation is excluded during evaluation 6 .
The pretrained models are trained with the suggested hyperparameters in SuPar.
The pretrained models achieve 97 .
91% and 94 .
78% UAS on PTB, CTB and UD Dutch respectively, where
RoBERTa
[Liu et al., 2019], ELECTRA [Cui et al., 2020] and XLM-RoBERTa
[Conneau et al., 2019] are adopted as encoders.
No BERT embeddings are adopted for the UD Turkish dataset.
For our ADMM algorithm, we adopt the adaptive scheme of varying penalty parameters ( τ incr = τ decr = 1 . 1 , µ = 1 ) in Boyd et al.
[2011] and the stopping criterion ( ϵ tol = 10 − 2 ) for consensus
ADMM in Xu et al.
In FW, the learning rate is set to 2 t +2 .
The smoothness weight µ and ambiguity radius λ = 2 ε are tuned using a logarithmic scale on [10 − 7 , 1] .
The batch size for the game-theoretic algorithm is 10 .
The batch size for Stochastic is 200 .
The error tolerance in
Game is set to 10 − 2 .
In stochastic gradient training, we use Adam with lr = 10 − 2 , β 1 = 0 .
9 , β 2 = 0 .
999 , ϵ = 10 − 8 .
In our experiments, for efficiency, we again adopt the FW algorithm for the outer maximization in Marginal .
Complete main experimental results including all the metrics are shown in Table 2.
D Extension Details
For the dependency tree polytope, recall that the dual problem of projection onto U ′
r
:=
{ x : x ∈
U r ∧
[Equation]
is
[Equation]
where γ ( i,j,l ) := α j + 1 ( i = r ) β .
Following Zhang et al.
[2010] similarly, we sort 2 w ( i,j,l ) for each j and compute the optimal α ∗ j with β = 0 .
Let the sorted w ’s be ( w ( j ) 1 , . . .
, w ( j ) n ) for each j .
We blend create a set { w ( j ) x − α ∗ j } for all j and x .
Let the sorted sequence be −∞ = t 1 < t 2 < · · · < t n t = ∞ .
The derivative with respect to β is piecewise-linear in each interval [ t k , t k +1 ] .
Since the objective is concave in β , we can iterate over all the intervals or find the optimal β ∗ with binary search.
For higher-order tree local polytopes, the central problem is the projection onto
[Equation]
The only variables of interest are x a and x s , given x s , the optimal x a is simply x ∗ a = max( w a , x s ) .
We can sort ( w a , w s ) a ∈ s and enumerate the range x s takes over this set.
E Wong’s Arborescence Polytope
We introduce another extended formulation of the arborescence polytope based on a multi-commodity flow representation [Wong, 1980, Martins, 2012, Friesen, 2019] as follows, which may be of independent interest:
[Equation]
Thus we have the arborescence polytope:
[Equation]
) satisfy equations (11) − (13) } .
According to Martins [2012], Friesen [2019], A mc ≜
A arb instead of an outer polytope of A arb .
We are interested in the following quadratic programming problem with linear inequality constraints:
[Equation]
We can reformulate the problem as
[Equation]
[Table]
The scaled augmented Lagrangian function is
[Equation]
where
[Equation]
The ADMM algorithm updates the parameters as follows:
[Equation]
Projection onto X is decomposable over each j ∈ V .
And for each j , the optimal value of the group can be computed in
O ( n ) in almost closed form via Section 5.5.1 in Zhang et al.
[2010] or other simplex projection algorithms in O ( n log n ) .
Projection onto U is a minimum quadratic capacity expansion cost problem for fixed multi-commodity flows:
[Equation]
A partially relaxed problem is
[Equation] [Equation]
Given β , the sub-problem for u is
[Equation]
with an analytical solution
[Equation]
Given β , the sub-problem for f is
[Equation]
which is a minimum-cost multi-commodity flow problem.
With u ∗ and f ∗ , we can optimize β with sub-gradient ascent.
Alternatively, another partially relaxed problem is
[Equation]
where h ( a ) and t ( a ) are the head and tail of arc a respectively.
Given β , the inner minimization problem is decomposed over a :
[Equation]
which is a convex continuous knapsack problem for each a .
The above optimization requires sub-gradient methods, which are usually slower than FW ( O ( 1 ϵ 2 ) ).