Policy gradient is an efficient technique for improving a policy in a reinforcement learning setting. However, vanilla online variants are on-policy only and not able to take advantage of off-policy data. In this paper we describe a new technique that combines policy gradient with off-policy Q-learning, drawing experience from a replay buffer. This is motivated by making a connection between the fixed points of the regularized policy gradient algorithm and the Q-values. This connection allows us to estimate the Q-values from the action preferences of the policy, to which we apply Q-learning updates. We refer to the new technique as ‘PGQL’, for policy gradient and Q-learning. We also establish an equivalency between action-value fitting techniques and actor-critic algorithms, showing that regularized policy gradient techniques can be interpreted as advantage function learning algorithms. We conclude with some numerical examples that demonstrate improved data efficiency and stability of PGQL. In particular, we tested PGQL on the full suite of Atari games and achieved performance exceeding that of both asynchronous advantage actor-critic (A3C) and Q-learning.
1 INTRODUCTION
In reinforcement learning an agent explores an environment and through the use of a reward signal learns to optimize its behavior to maximize the expected long-term return. Reinforcement learning has seen success in several areas including robotics (Lin, 1993; Levine et al., 2015), computer games (Mnih et al., 2013; 2015), online advertising (Pednault et al., 2002), board games (Tesauro, 1995; Silver et al., 2016), and many others. For an introduction to reinforcement learning we refer to the classic text by Sutton & Barto (1998). In this paper we consider model-free reinforcement learning, where the state-transition function is not known or learned. There are many different algorithms for model-free reinforcement learning, but most fall into one of two families: action-value fitting and policy gradient techniques.
Action-value techniques involve fitting a function, called the Q-values, that captures the expected return for taking a particular action at a particular state, and then following a particular policy thereafter. Two alternatives we discuss in this paper are SARSA (Rummery & Niranjan, 1994) and Q-learning (Watkins, 1989), although there are many others. SARSA is an on-policy algorithm whereby the action-value function is fit to the current policy, which is then refined by being mostly greedy with respect to those action-values. On the other hand, Q-learning attempts to find the Qvalues associated with the optimal policy directly and does not fit to the policy that was used to generate the data. Q-learning is an off-policy algorithm that can use data generated by another agent or from a replay buffer of old experience. Under certain conditions both SARSA and Q-learning can be shown to converge to the optimal Q-values, from which we can derive the optimal policy (Sutton, 1988; Bertsekas & Tsitsiklis, 1996).
In policy gradient techniques the policy is represented explicitly and we improve the policy by updating the parameters in the direction of the gradient of the performance (Sutton et al., 1999; Silver et al., 2014; Kakade, 2001). Online policy gradient typically requires an estimate of the action-value function of the current policy. For this reason they are often referred to as actor-critic methods, where the actor refers to the policy and the critic to the estimate of the action-value function (Konda & Tsitsiklis, 2003). Vanilla actor-critic methods are on-policy only, although some attempts have been made to extend them to off-policy data (Degris et al., 2012; Levine & Koltun, 2013).
In this paper we derive a link between the Q-values induced by a policy and the policy itself when the policy is the fixed point of a regularized policy gradient algorithm (where the gradient vanishes). This connection allows us to derive an estimate of the Q-values from the current policy, which we can refine using off-policy data and Q-learning. We show in the tabular setting that when the regularization penalty is small (the usual case) the resulting policy is close to the policy that would be found without the addition of the Q-learning update. Separately, we show that regularized actor-critic methods can be interpreted as action-value fitting methods, where the Q-values have been parameterized in a particular way. We conclude with some numerical examples that provide empirical evidence of improved data efficiency and stability of PGQL.
1.1 PRIOR WORK
Here we highlight various axes along which our work can be compared to others. In this paper we use entropy regularization to ensure exploration in the policy, which is a common practice in policy gradient (Williams & Peng, 1991; Mnih et al., 2016). An alternative is to use KL-divergence instead of entropy as a regularizer, or as a constraint on how much deviation is permitted from a prior policy (Bagnell & Schneider, 2003; Peters et al., 2010; Schulman et al., 2015; Fox et al., 2015). Natural policy gradient can also be interpreted as putting a constraint on the KL-divergence at each step of the policy improvement (Amari, 1998; Kakade, 2001; Pascanu & Bengio, 2013). In Sallans & Hinton (2004) the authors use a Boltzmann exploration policy over estimated Q-values which they update using TD-learning. In Heess et al. (2012) this was extended to use an actor-critic algorithm instead of TD-learning, however the two updates were not combined as we have done in this paper. In Azar et al. (2012) the authors develop an algorithm called dynamic policy programming, whereby they apply a Bellman-like update to the action-preferences of a policy, which is similar in spirit to the update we describe here. In Norouzi et al. (2016) the authors augment a maximum likelihood objective with a reward in a supervised learning setting, and develop a connection that resembles the one we develop here between the policy and the Q-values. Other works have attempted to combine on and off-policy learning, primarily using action-value fitting methods (Wang et al., 2013; Hausknecht & Stone, 2016; Lehnert & Precup, 2015), with varying degrees of success. In this paper we establish a connection between actor-critic algorithms and action-value learning algorithms. In particular we show that TD-actor-critic (Konda & Tsitsiklis, 2003) is equivalent to expected-SARSA (Sutton & Barto, 1998, Exercise 6.10) with Boltzmann exploration where the Q-values are decomposed into advantage function and value function. The algorithm we develop extends actor-critic with a Q-learning style update that, due to the decomposition of the Q-values, resembles the update of the dueling architecture (Wang et al., 2016). Recently, the field of deep reinforcement learning, i.e., the use of deep neural networks to represent action-values or a policy, has seen a lot of success (Mnih et al., 2015; 2016; Silver et al., 2016; Riedmiller, 2005; Lillicrap et al., 2015; Van Hasselt et al., 2016). In the examples section we use a neural network with PGQL to play the Atari games suite.
2 REINFORCEMENT LEARNING
We consider the infinite horizon, discounted, finite state and action space Markov decision process, with state space S, action space A and rewards at each time period denoted by rt ∈ R. A policy π : S × A → R+ is a mapping from state-action pair to the probability of taking that action at that state, so it must satisfy ∑ a∈A π(s, a) = 1 for all states s ∈ S . Any policy π induces a probability distribution over visited states, dπ : S → R+ (which may depend on the initial state), so the probability of seeing state-action pair (s, a) ∈ S ×A is dπ(s)π(s, a). In reinforcement learning an ‘agent’ interacts with an environment over a number of times steps. At each time step t the agent receives a state st and a reward rt and selects an action at from the policy πt, at which point the agent moves to the next state st+1 ∼ P (·, st, at), where P (s′, s, a) is the probability of transitioning from state s to state s′ after taking action a. This continues until the agent encounters a terminal state (after which the process is typically restarted). The goal of the agent is to find a policy π that maximizes the expected total discounted return J(π) = E( ∑∞ t=0 γ
trt | π), where the expectation is with respect to the initial state distribution, the state-transition probabilities, and the policy, and where γ ∈ (0, 1) is the discount factor that, loosely speaking, controls how much the agent prioritizes long-term versus short-term rewards. Since the agent starts with no knowledge
of the environment it must continually explore the state space and so will typically use a stochastic policy.
Action-values. The action-value, or Q-value, of a particular state under policy π is the expected total discounted return from taking that action at that state and following π thereafter, i.e., Qπ(s, a) = E( ∑∞ t=0 γ
trt | s0 = s, a0 = a, π). The value of state s under policy π is denoted by V π(s) = E( ∑∞ t=0 γ
trt | s0 = s, π), which is the expected total discounted return of policy π from state s. The optimal action-value function is denoted Q? and satisfies Q?(s, a) = maxπ Qπ(s, a) for each (s, a). The policy that achieves the maximum is the optimal policy π?, with value function V ?. The advantage function is the difference between the action-value and the value function, i.e., Aπ(s, a) = Qπ(s, a)−V π(s), and represents the additional expected reward of taking action a over the average performance of the policy from state s. Since V π(s) = ∑ a π(s, a)Q
π(s, a) we have the identity ∑ a π(s, a)A
π(s, a) = 0, which simply states that the policy π has no advantage over itself.
Bellman equation. The Bellman operator T π (Bellman, 1957) for policy π is defined as
T πQ(s, a) = E s′,r,b (r(s, a) + γQ(s′, b)),
where the expectation is over next state s′ ∼ P (·, s, a), the reward r(s, a), and the action b from policy πs′ . The Q-value function for policy π is the fixed point of the Bellman operator for π, i.e., T πQπ = Qπ . The optimal Bellman operator T ? is defined as
T ?Q(s, a) = E s′,r (r(s, a) + γmax b Q(s′, b)),
where the expectation is over the next state s′ ∼ P (·, s, a), and the reward r(s, a). The optimal Q-value function is the fixed point of the optimal Bellman equation, i.e., T ?Q? = Q?. Both the π-Bellman operator and the optimal Bellman operator are γ-contraction mappings in the sup-norm, i.e., ‖T Q1−T Q2‖∞ ≤ γ‖Q1−Q2‖∞, for any Q1, Q2 ∈ RS×A. From this fact one can show that the fixed point of each operator is unique, and that value iteration converges, i.e., (T π)kQ → Qπ and (T ?)kQ→ Q? from any initial Q. (Bertsekas, 2005).
2.1 ACTION-VALUE LEARNING
In value based reinforcement learning we approximate the Q-values using a function approximator. We then update the parameters so that the Q-values are as close to the fixed point of a Bellman equation as possible. If we denote by Q(s, a; θ) the approximate Q-values parameterized by θ, then Q-learning updates the Q-values along direction Es,a(T ?Q(s, a; θ)−Q(s, a; θ))∇θQ(s, a; θ) and SARSA updates the Q-values along direction Es,a(T πQ(s, a; θ)−Q(s, a; θ))∇θQ(s, a; θ). In the online setting the Bellman operator is approximated by sampling and bootstrapping, whereby the Q-values at any state are updated using the Q-values from the next visited state. Exploration is achieved by not always taking the action with the highest Q-value at each time step. One common technique called ‘epsilon greedy’ is to sample a random action with probability > 0, where starts high and decreases over time. Another popular technique is ‘Boltzmann exploration’, where the policy is given by the softmax over the Q-values with a temperature T , i.e., π(s, a) = exp(Q(s, a)/T )/ ∑ b exp(Q(s, b)/T ), where it is common to decrease the temperature over time.
2.2 POLICY GRADIENT
Alternatively, we can parameterize the policy directly and attempt to improve it via gradient ascent on the performance J . The policy gradient theorem (Sutton et al., 1999) states that the gradient of J with respect to the parameters of the policy is given by
∇θJ(π) = E s,a Qπ(s, a)∇θ log π(s, a), (1)
where the expectation is over (s, a) with probability dπ(s)π(s, a). In the original derivation of the policy gradient theorem the expectation is over the discounted distribution of states, i.e., over dπ,s0γ (s) = ∑∞ t=0 γ tPr{st = s | s0, π}. However, the gradient update in that case will assign a low
weight to states that take a long time to reach and can therefore have poor empirical performance. In practice the non-discounted distribution of states is frequently used instead. In certain cases this is equivalent to maximizing the average (i.e., non-discounted) policy performance, even whenQπ uses a discount factor (Thomas, 2014). Throughout this paper we will use the non-discounted distribution of states.
In the online case it is common to add an entropy regularizer to the gradient in order to prevent the policy becoming deterministic. This ensures that the agent will explore continually. In that case the (batch) update becomes
∆θ ∝ E s,a Qπ(s, a)∇θ log π(s, a) + αE s ∇θHπ(s), (2)
where Hπ(s) = − ∑ a π(s, a) log π(s, a) denotes the entropy of policy π, and α > 0 is the regularization penalty parameter. Throughout this paper we will make use of entropy regularization, however many of the results are true for other choices of regularizers with only minor modification, e.g., KL-divergence. Note that equation (2) requires exact knowledge of the Q-values. In practice they can be estimated, e.g., by the sum of discounted rewards along an observed trajectory (Williams, 1992), and the policy gradient will still perform well (Konda & Tsitsiklis, 2003).
3 REGULARIZED POLICY GRADIENT ALGORITHM
In this section we derive a relationship between the policy and the Q-values when using a regularized policy gradient algorithm. This allows us to transform a policy into an estimate of the Q-values. We then show that for small regularization the Q-values induced by the policy at the fixed point of the algorithm have a small Bellman error in the tabular case.
3.1 TABULAR CASE
Consider the fixed points of the entropy regularized policy gradient update (2). Let us define f(θ) = Es,aQ π(s, a)∇θ log π(s, a) +αEs∇θH(πs), and gs(π) = ∑ a π(s, a) for each s. A fixed point is one where we can no longer update θ in the direction of f(θ) without violating one of the constraints gs(π) = 1, i.e., where f(θ) is in the span of the vectors {∇θgs(π)}. In other words, any fixed point must satisfy f(θ) = ∑ s λs∇θgs(π), where for each s the Lagrange multiplier λs ∈ R ensures that gs(π) = 1. Substituting in terms to this equation we obtain
E s,a
(Qπ(s, a)− α log π(s, a)− cs)∇θ log π(s, a) = 0, (3)
where we have absorbed all constants into c ∈ R|S|. Any solution π to this equation is strictly positive element-wise since it must lie in the domain of the entropy function. In the tabular case π is represented by a single number for each state and action pair and the gradient of the policy with respect to the parameters is the indicator function, i.e.,∇θ(t,b)π(s, a) = 1(t,b)=(s,a). From this we obtain Qπ(s, a) − α log π(s, a) − cs = 0 for each s (assuming that the measure dπ(s) > 0). Multiplying by π(a, s) and summing over a ∈ A we get cs = αHπ(s) + V π(s). Substituting c into equation (3) we have the following formulation for the policy:
π(s, a) = exp(Aπ(s, a)/α−Hπ(s)), (4)
for all s ∈ S and a ∈ A. In other words, the policy at the fixed point is a softmax over the advantage function induced by that policy, where the regularization parameter α can be interpreted as the temperature. Therefore, we can use the policy to derive an estimate of the Q-values,
Q̃π(s, a) = Ãπ(s, a) + V π(s) = α(log π(s, a) +Hπ(s)) + V π(s). (5)
With this we can rewrite the gradient update (2) as
∆θ ∝ E s,a (Qπ(s, a)− Q̃π(s, a))∇θ log π(s, a), (6)
since the update is unchanged by per-state constant offsets. When the policy is parameterized as a softmax, i.e., π(s, a) = exp(W (s, a))/ ∑ b expW (s, b), the quantity W is sometimes referred to as the action-preferences of the policy (Sutton & Barto, 1998, Chapter 6.6). Equation (4) states that the action preferences are equal to the Q-values scaled by 1/α, up to an additive per-state constant.
3.2 GENERAL CASE
Consider the following optimization problem:
minimize Es,a(q(s, a)− α log π(s, a))2 subject to ∑ a π(s, a) = 1, s ∈ S
(7)
over variable θ which parameterizes π, where we consider both the measure in the expectation and the values q(s, a) to be independent of θ. The optimality condition for this problem is
E s,a
(q(s, a)− α log π(s, a) + cs)∇θ log π(s, a) = 0,
where c ∈ R|S| is the Lagrange multiplier associated with the constraint that the policy sum to one at each state. Comparing this to equation (3), we see that if q = Qπ and the measure in the expectation is the same then they describe the same set of fixed points. This suggests an interpretation of the fixed points of the regularized policy gradient as a regression of the log-policy onto the Q-values. In the general case of using an approximation architecture we can interpret equation (3) as indicating that the error between Qπ and Q̃π is orthogonal to ∇θi log π for each i, and so cannot be reduced further by changing the parameters, at least locally. In this case equation (4) is unlikely to hold at a solution to (3), however with a good approximation architecture it may hold approximately, so that the we can derive an estimate of the Q-values from the policy using equation (5). We will use this estimate of the Q-values in the next section.
3.3 CONNECTION TO ACTION-VALUE METHODS
The previous section made a connection between regularized policy gradient and a regression onto the Q-values at the fixed point. In this section we go one step further, showing that actor-critic methods can be interpreted as action-value fitting methods, where the exact method depends on the choice of critic.
Actor-critic methods. Consider an agent using an actor-critic method to learn both a policy π and a value function V . At any iteration k, the value function V k has parameters wk, and the policy is of the form
πk(s, a) = exp(W k(s, a)/α)/ ∑ b exp(W k(s, b)/α), (8)
where W k is parameterized by θk and α > 0 is the entropy regularization penalty. In this case ∇θ log πk(s, a) = (1/α)(∇θW k(s, a) − ∑ b π(s, b)∇θW k(s, b)). Using equation (6) the parameters are updated as
∆θ ∝ E s,a δac(∇θW k(s, a)− ∑ b πk(s, b)∇θW k(s, b)), ∆w ∝ E s,a δac∇wV k(s) (9)
where δac is the critic minus baseline term, which depends on the variant of actor-critic being used (see the remark below).
Action-value methods. Compare this to the case where an agent is learning Q-values with a dueling architecture (Wang et al., 2016), which at iteration k is given by
Qk(s, a) = Y k(s, a)− ∑ b µ(s, b)Y k(s, b) + V k(s),
where µ is a probability distribution, Y k is parameterized by θk, V k is parameterized by wk, and the exploration policy is Boltzmann with temperature α, i.e.,
πk(s, a) = exp(Y k(s, a)/α)/ ∑ b exp(Y k(s, b)/α). (10)
In action value fitting methods at each iteration the parameters are updated to reduce some error, where the update is given by
∆θ ∝ E s,a δav(∇θY k(s, a)− ∑ b µ(s, b)∇θY k(s, b)), ∆w ∝ E s,a δav∇wV k(s) (11)
where δav is the action-value error term and depends on which algorithm is being used (see the remark below).
Equivalence. The two policies (8) and (10) are identical if W k = Y k for all k. Since X0 and Y 0 can be initialized and parameterized in the same way, and assuming the two value function estimates are initialized and parameterized in the same way, all that remains is to show that the updates in equations (11) and (9) are identical. Comparing the two, and assuming that δac = δav (see remark), we see that the only difference is that the measure is not fixed in (9), but is equal to the current policy and therefore changes after each update. Replacing µ in (11) with πk makes the updates identical, in which case W k = Y k at all iterations and the two policies (8) and (10) are always the same. In other words, the slightly modified action-value method is equivalent to an actor-critic policy gradient method, and vice-versa (modulo using the non-discounted distribution of states, as discussed in §2.2). In particular, regularized policy gradient methods can be interpreted as advantage function learning techniques (Baird III, 1993), since at the optimum the quantity W (s, a) − ∑ b π(s, b)W (s, b) = α(log π(s, a) + H
π(s)) will be equal to the advantage function values in the tabular case.
Remark. In SARSA (Rummery & Niranjan, 1994) we set δav = r(s, a) + γQ(s′, b) − Q(s, a), where b is the action selected at state s′, which would be equivalent to using a bootstrap critic in equation (6) where Qπ(s, a) = r(s, a) + γQ̃(s′, b). In expected-SARSA (Sutton & Barto, 1998, Exercise 6.10), (Van Seijen et al., 2009)) we take the expectation over the Q-values at the next state, so δav = r(s, a)+γV (s′)−Q(s, a). This is equivalent to TD-actor-critic (Konda & Tsitsiklis, 2003) where we use the value function to provide the critic, which is given by Qπ = r(s, a) + γV (s′). In Q-learning (Watkins, 1989) δav = r(s, a) + γmaxbQ(s′, b)−Q(s, a), which would be equivalent to using an optimizing critic that bootstraps using the max Q-value at the next state, i.e., Qπ(s, a) = r(s, a) + γmaxb Q̃
π(s′, b). In REINFORCE the critic is the Monte Carlo return from that state on, i.e., Qπ(s, a) = ( ∑∞ t=0 γ
trt | s0 = s, a0 = a). If the return trace is truncated and a bootstrap is performed after n-steps, this is equivalent to n-step SARSA or n-step Q-learning, depending on the form of the bootstrap (Peng & Williams, 1996).
3.4 BELLMAN RESIDUAL
In this section we show that ‖T ?Qπα −Qπα‖ → 0 with decreasing regularization penalty α, where πα is the policy defined by (4) and Qπα is the corresponding Q-value function, both of which are functions of α. We shall show that it converges to zero by bounding the sequence below by zero and above with a sequence that converges to zero. First, we have that T ?Qπα ≥ T παQπα = Qπα , since T ? is greedy with respect to the Q-values. So T ?Qπα −Qπα ≥ 0. Now, to bound from above we need the fact that πα(s, a) = exp(Qπα(s, a)/α)/ ∑ b exp(Q
πα(s, b)/α) ≤ exp((Qπα(s, a) − maxcQ πα(s, c))/α). Using this we have
0 ≤ T ?Qπα(s, a)−Qπα(s, a) = T ?Qπα(s, a)− T παQπα(s, a) = Es′(maxcQ πα(s′, c)− ∑ b πα(s
′, b)Qπα(s′, b)) = Es′ ∑ b πα(s ′, b)(maxcQ πα(s′, c)−Qπα(s′, b))
≤ Es′ ∑ b exp((Q
πα(s′, b)−Qπα(s′, b?))/α)(maxcQπα(s′, c)−Qπα(s′, b)) = Es′ ∑ b fα(maxcQ πα(s′, c)−Qπα(s′, b)),
where we define fα(x) = x exp(−x/α). To conclude our proof we use the fact that fα(x) ≤ supx fα(x) = fα(α) = αe −1, which yields
0 ≤ T ?Qπα(s, a)−Qπα(s, a) ≤ |A|αe−1
for all (s, a), and so the Bellman residual converges to zero with decreasing α. In other words, for small enough α (which is the regime we are interested in) the Q-values induced by the policy (4) will have a small Bellman residual. Moreover, this implies that limα→0Qπα = Q?, as one might expect.
4 PGQL
In this section we introduce the main contribution of the paper, which is a technique to combine policy gradient with Q-learning. We call our technique ‘PGQL’, for policy gradient and Q-learning. In the previous section we showed that the Bellman residual is small at the fixed point of a regularized
policy gradient algorithm when the regularization penalty is sufficiently small. This suggests adding an auxiliary update where we explicitly attempt to reduce the Bellman residual as estimated from the policy, i.e., a hybrid between policy gradient and Q-learning.
We first present the technique in a batch update setting, with a perfect knowledge of Qπ (i.e., a perfect critic). Later we discuss the practical implementation of the technique in a reinforcement learning setting with function approximation, where the agent generates experience from interacting with the environment and needs to estimate a critic simultaneously with the policy.
4.1 PGQL UPDATE
Define the estimate of Q using the policy as
Q̃π(s, a) = α(log π(s, a) +Hπ(s)) + V (s), (12)
where V has parameters w and is not necessarily V π as it was in equation (5). In (2) it was unnecessary to estimate the constant since the update was invariant to constant offsets, although in practice it is often estimated for use in a variance reduction technique (Williams, 1992; Sutton et al., 1999).
Since we know that at the fixed point the Bellman residual will be small for small α, we can consider updating the parameters to reduce the Bellman residual in a fashion similar to Q-learning, i.e.,
∆θ ∝ E s,a (T ?Q̃π(s, a)− Q̃π(s, a))∇θ log π(s, a), ∆w ∝ E s,a (T ?Q̃π(s, a)− Q̃π(s, a))∇wV (s). (13) This is Q-learning applied to a particular form of the Q-values, and can also be interpreted as an actor-critic algorithm with an optimizing (and therefore biased) critic.
The full scheme simply combines two updates to the policy, the regularized policy gradient update (2) and the Q-learning update (13). Assuming we have an architecture that provides a policy π, a value function estimate V , and an action-value critic Qπ , then the parameter updates can be written as (suppressing the (s, a) notation)
∆θ ∝ (1− η)Es,a(Qπ − Q̃π)∇θ log π + ηEs,a(T ?Q̃π − Q̃π)∇θ log π,
∆w ∝ (1− η)Es,a(Qπ − Q̃π)∇wV + ηEs,a(T ?Q̃π − Q̃π)∇wV, (14)
here η ∈ [0, 1] is a weighting parameter that controls how much of each update we apply. In the case where η = 0 the above scheme reduces to entropy regularized policy gradient. If η = 1 then it becomes a variant of (batch) Q-learning with an architecture similar to the dueling architecture (Wang et al., 2016). Intermediate values of η produce a hybrid between the two. Examining the update we see that two error terms are trading off. The first term encourages consistency with critic, and the second term encourages optimality over time. However, since we know that under standard policy gradient the Bellman residual will be small, then it follows that adding a term that reduces that error should not make much difference at the fixed point. That is, the updates should be complementary, pointing in the same general direction, at least far away from a fixed point. This update can also be interpreted as an actor-critic update where the critic is given by a weighted combination of a standard critic and an optimizing critic. Yet another interpretation of the update is a combination of expected-SARSA and Q-learning, where the Q-values are parameterized as the sum of an advantage function and a value function.
4.2 PRACTICAL IMPLEMENTATION
The updates presented in (14) are batch updates, with an exact critic Qπ . In practice we want to run this scheme online, with an estimate of the critic, where we don’t necessarily apply the policy gradient update at the same time or from same data source as the Q-learning update.
Our proposal scheme is as follows. One or more agents interact with an environment, encountering states and rewards and performing on-policy updates of (shared) parameters using an actor-critic algorithm where both the policy and the critic are being updated online. Each time an agent receives new data from the environment it writes it to a shared replay memory buffer. Periodically a separate learner process samples from the replay buffer and performs a step of Q-learning on the parameters of the policy using (13). This scheme has several advantages. The critic can accumulate the Monte
Carlo return over many time periods, allowing us to spread the influence of a reward received in the future backwards in time. Furthermore, the replay buffer can be used to store and replay ‘important’ past experiences by prioritizing those samples (Schaul et al., 2015). The use of the replay buffer can help to reduce problems associated with correlated training data, as generated by an agent exploring an environment where the states are likely to be similar from one time step to the next. Also the use of replay can act as a kind of regularizer, preventing the policy from moving too far from satisfying the Bellman equation, thereby improving stability, in a similar sense to that of a policy ‘trust-region’ (Schulman et al., 2015). Moreover, by batching up replay samples to update the network we can leverage GPUs to perform the updates quickly, this is in comparison to pure policy gradient techniques which are generally implemented on CPU (Mnih et al., 2016).
Since we perform Q-learning using samples from a replay buffer that were generated by a old policy we are performing (slightly) off-policy learning. However, Q-learning is known to converge to the optimal Q-values in the off-policy tabular case (under certain conditions) (Sutton & Barto, 1998), and has shown good performance off-policy in the function approximation case (Mnih et al., 2013).
4.3 MODIFIED FIXED POINT
The PGQL updates in equation (14) have modified the fixed point of the algorithm, so the analysis of §3 is no longer valid. Considering the tabular case once again, it is still the case that the policy π ∝ exp(Q̃π/α) as before, where Q̃π is defined by (12), however where previously the fixed point satisfied Q̃π = Qπ , with Qπ corresponding to the Q-values induced by π, now we have
Q̃π = (1− η)Qπ + ηT ?Q̃π, (15) Or equivalently, if η < 1, we have Q̃π = (1− η) ∑∞ k=0 η
k(T ?)kQπ . In the appendix we show that ‖Q̃π − Qπ‖ → 0 and that ‖T ?Qπ − Qπ‖ → 0 with decreasing α in the tabular case. That is, for small α the induced Q-values and the Q-values estimated from the policy are close, and we still have the guarantee that in the limit the Q-values are optimal. In other words, we have not perturbed the policy very much by the addition of the auxiliary update.
5 NUMERICAL EXPERIMENTS
5.1 GRID WORLD
In this section we discuss the results of running PGQL on a toy 4 by 6 grid world, as shown in Figure 1a. The agent always begins in the square marked ‘S’ and the episode continues until it reaches the square marked ‘T’, upon which it receives a reward of 1. All other times it receives no reward. For this experiment we chose regularization parameter α = 0.001 and discount factor γ = 0.95.
Figure 1b shows the performance traces of three different agents learning in the grid world, running from the same initial random seed. The lines show the true expected performance of the policy
from the start state, as calculated by value iteration after each update. The blue-line is standard TD-actor-critic (Konda & Tsitsiklis, 2003), where we maintain an estimate of the value function and use that to generate an estimate of the Q-values for use as the critic. The green line is Q-learning where at each step an update is performed using data drawn from a replay buffer of prior experience and where the Q-values are parameterized as in equation (12). The policy is a softmax over the Q-value estimates with temperature α. The red line is PGQL, which at each step first performs the TD-actor-critic update, then performs the Q-learning update as in (14).
The grid world was totally deterministic, so the step size could be large and was chosen to be 1. A step-size any larger than this made the pure actor-critic agent fail to learn, but both PGQL and Q-learning could handle some increase in the step-size, possibly due to the stabilizing effect of using replay.
It is clear that PGQL outperforms the other two. At any point along the x-axis the agents have seen the same amount of data, which would indicate that PGQL is more data efficient than either of the vanilla methods since it has the highest performance at practically every point.
5.2 ATARI
We tested our algorithm on the full suite of Atari benchmarks (Bellemare et al., 2012), using a neural network to parameterize the policy. In figure 2 we show how a policy network can be augmented with a parameterless additional layer which outputs the Q-value estimate. With the exception of the extra layer, the architecture and parameters were chosen to exactly match the asynchronous advantage actor-critic (A3C) algorithm presented in Mnih et al. (2016), which in turn reused many of the settings from Mnih et al. (2015). Specifically we used the exact same learning rate, number of workers, entropy penalty, bootstrap horizon, and network architecture. This allows a fair comparison between A3C and PGQL, since the only difference is the addition of the Q-learning step. Our technique augmented A3C with the following change: After each actor-learner has accumulated the gradient for the policy update, it performs a single step of Q-learning from replay data as described in equation (13), where the minibatch size was 32 and the Q-learning learning rate was chosen to be 0.5 times the actor-critic learning rate (we mention learning rate ratios rather than choice of η in (14) because the updates happen at different frequencies and from different data sources). Each actor-learner thread maintained a replay buffer of the last 100k transitions seen by that thread. We ran the learning for 50 million agent steps (200 million Atari frames), as in (Mnih et al., 2016).
In the results we compare against both A3C and a variant of asynchronous deep Q-learning. The changes we made to Q-learning are to make it similar to our method, with some tuning of the hyperparameters for performance. We use the exact same network, the exploration policy is a softmax over the Q-values with a temperature of 0.1, and the Q-values are parameterized as in equation (12) (i.e., similar to the dueling architecture (Wang et al., 2016)), where α = 0.1. The Q-value updates are performed every 4 steps with a minibatch of 32 (roughly 5 times more frequently than PGQL). For each method, all games used identical hyper-parameters.
The results across all games are given in table 3 in the appendix. All scores have been normalized by subtracting the average score achieved by an agent that takes actions uniformly at random. Each game was tested 5 times per method with the same hyper-parameters but with different ran-
dom seeds. The scores presented correspond to the best score obtained by any run from a random start evaluation condition (Mnih et al., 2016). Overall, PGQL performed best in 34 games, A3C performed best in 7 games, and Q-learning was best in 10 games. In 6 games two or more methods tied. In tables 1 and 2 we give the mean and median normalized scores as percentage of an expert human normalized score across all games for each tested algorithm from random and human-start conditions respectively. In a human-start condition the agent takes over control of the game from randomly selected human-play starting points, which generally leads to lower performance since the agent may not have found itself in that state during training. In both cases, PGQL has both the highest mean and median, and the median score exceeds 100%, the human performance threshold.
It is worth noting that PGQL was the worst performer in only one game, in cases where it was not the outright winner it was generally somewhere in between the performance of the other two algorithms. Figure 3 shows some sample traces of games where PGQL was the best performer. In these cases PGQL has far better data efficiency than the other methods. In figure 4 we show some of the games where PGQL under-performed. In practically every case where PGQL did not perform well it had better data efficiency early on in the learning, but performance saturated or collapsed. We hypothesize that in these cases the policy has reached a local optimum, or over-fit to the early data, and might perform better were the hyper-parameters to be tuned.
6 CONCLUSIONS
We have made a connection between the fixed point of regularized policy gradient techniques and the Q-values of the resulting policy. For small regularization (the usual case) we have shown that the Bellman residual of the induced Q-values must be small. This leads us to consider adding an auxiliary update to the policy gradient which is related to the Bellman residual evaluated on a transformation of the policy. This update can be performed off-policy, using stored experience. We call the resulting method ‘PGQL’, for policy gradient and Q-learning. Empirically, we observe better data efficiency and stability of PGQL when compared to actor-critic or Q-learning alone. We verified the performance of PGQL on a suite of Atari games, where we parameterize the policy using a neural network, and achieved performance exceeding that of both A3C and Q-learning.
7 ACKNOWLEDGMENTS
We thank Joseph Modayil for many comments and suggestions on the paper, and Hubert Soyer for help with performance evaluation. We would also like to thank the anonymous reviewers for their constructive feedback.
A PGQL BELLMAN RESIDUAL
Here we demonstrate that in the tabular case the Bellman residual of the induced Q-values for the PGQL updates of (14) converges to zero as the temperature α decreases, which is the same guarantee as vanilla regularized policy gradient (2). We will use the notation that πα is the policy at the fixed point of PGQL updates (14) for some α, i.e., πα ∝ exp(Q̃πα), with induced Q-value function Qπα .
First, note that we can apply the same argument as in §3.4 to show that limα→0 ‖T ?Q̃πα − T παQ̃πα‖ = 0 (the only difference is that we lack the property that Q̃πα is the fixed point of T πα ). Secondly, from equation (15) we can write Q̃πα − Qπα = η(T ?Q̃πα − Qπα). Combining these two facts we have
‖Q̃πα −Qπα‖ = η‖T ?Q̃πα −Qπα‖ = η‖T ?Q̃πα − T παQ̃πα + T παQ̃πα −Qπα‖ ≤ η(‖T ?Q̃πα − T παQ̃πα‖+ ‖T παQ̃πα − T παQπα‖) ≤ η(‖T ?Q̃πα − T παQ̃πα‖+ γ‖Q̃πα −Qπα‖) ≤ η/(1− ηγ)‖T ?Q̃πα − T παQ̃πα‖,
and so ‖Q̃πα −Qπα‖ → 0 as α→ 0. Using this fact we have
‖T ?Q̃πα − Q̃πα‖ = ‖T ?Q̃πα − T παQ̃πα + T παQ̃πα −Qπα +Qπα − Q̃πα‖ ≤ ‖T ?Q̃πα − T παQ̃πα‖+ ‖T παQ̃πα − T παQπα‖+ ‖Qπα − Q̃πα‖ ≤ ‖T ?Q̃πα − T παQ̃πα‖+ (1 + γ)‖Q̃πα −Qπα‖ < 3/(1− ηγ)‖T ?Q̃πα − T παQ̃πα‖,
which therefore also converges to zero in the limit. Finally we obtain
‖T ?Qπα −Qπα‖ = ‖T ?Qπα − T ?Q̃πα + T ?Q̃πα − Q̃πα + Q̃πα −Qπα‖ ≤ ‖T ?Qπα − T ?Q̃πα‖+ ‖T ?Q̃πα − Q̃πα‖+ ‖Q̃πα −Qπα‖ ≤ (1 + γ)‖Q̃πα −Qπα‖+ ‖T ?Q̃πα − Q̃πα‖,
which combined with the two previous results implies that limα→0 ‖T ?Qπα−Qπα‖ = 0, as before.
B ATARI SCORES

This paper shows how policy gradient and Q-Learning may be combined together, improving learning as demonstrated in particular in the Atari Learning Environment. The core idea is to note that entropy-regularized policy gradient leads to a Boltzmann policy based on Q values, thus linking pi & Q together and allowing both policy gradient and Q-Learning updates to be applied.
I think this is a very interesting paper, not just for its results and the proposed algorithm (dubbed PGQ), but mostly because of the links it draws between several techniques, which I found quite insightful.
That being said, I also believe it could have done a better job at clearly exposing these links: I found it somewhat difficult to follow, and it took me a while to wrap my head around it, even though the underlying concepts are not that complex. In particular:
- The notation \tilde{Q}^pi is introduced in a way that is not very clear, as "an estimate of the Q-values" while eq. 5 is an exact equality (no estimation)
- It is not clear to me what section 3.2 is bringing exactly, I wonder if it could just be removed to expand some other sections with more explanations.
- The links to dueling networks (Wang et al, 2016) are in my opinion not explicit and detailed enough (in 3.3 & 4.1): as far as I can tell the proposed architecture ends up being very similar to such networks and thus it would be worth telling more about it (also in experiments my understanding is that the "variant of asynchronous deep Q-learning" being used is essentially such a dueling network, but it is not clearly stated). I also believe it should be mentioned that PGQ can also be seen as combining Q-Learning with n-step expected Sarsa using a dueling network: this kind of example helps better understand the links between methods
- Overall I wish section 3.3 was clearer, as it draws some very interesting links, but it is hard to see where this is all going when reading the paper for the first time. One confusing point is w.r.t. to the relationship with section 3.2, that assumes a critic outputting Q values while in 3.3 the critic outputs V. The "mu" distribution also comes somewhat out of nowhere.
I hope the authors can try and improve the readability of the paper in a final version, as well as clarify the points raised in pre-review questions (in particular related to experimental details, the derivation of eq. 4, and the issue of the discounted distribution of states).
Minor remarks:
- The function r(s, a) used in the Bellman equation in section 2 is not formally defined. It looks a bit weird because the expectation is on s' and b' but r(s, a) does not depend on them (so either it should be moved out of the expectation, or the expectation should also be over the reward, depending on how r is defined)
- The definition of the Boltzmann policy at end of 2.1 is a bit confusing since there is a sum over "a" of a quantity that does not depend (clearly) on "a"
- I believe 4.3 is for the tabular case but this is not clearly stated
- Any idea why in Fig. 1 the 3 algorithms do not all converge to the same policy? In such a simple toy setting I would expect it to be the case.
Typos:
- "we refer to the classic text Sutton & Barto (1998)" => missing "by"?
- "Online policy gradient typically require an estimate of the action-values function" => requires & value
- "the agent generates experience from interacting the environment" => with the environment
- in eq. 12 (first line) there is a comma to remove near the end, just before the dlog pi
- "allowing us the spread the influence of a reward" => to spread
- "in the off-policy case tabular case" => remove the first case
- "The green line is Q-learning where at the step an update is performed" => at each step
- In Fig. 2 it says A2C instead of A3C
NB: I did not have time to carefully read Appendix A


STRUCTURED ATTENTION NETWORKS
Attention networks have proven to be an effective approach for embedding categorical inference within a deep neural network. However, for many tasks we may want to model richer structural dependencies without abandoning end-to-end training. In this work, we experiment with incorporating richer structural distributions, encoded using graphical models, within deep networks. We show that these structured attention networks are simple extensions of the basic attention procedure, and that they allow for extending attention beyond the standard softselection approach, such as attending to partial segmentations or to subtrees. We experiment with two different classes of structured attention networks: a linearchain conditional random field and a graph-based parsing model, and describe how these models can be practically implemented as neural network layers. Experiments show that this approach is effective for incorporating structural biases, and structured attention networks outperform baseline attention models on a variety of synthetic and real tasks: tree transduction, neural machine translation, question answering, and natural language inference. We further find that models trained in this way learn interesting unsupervised hidden representations that generalize simple attention.
1 INTRODUCTION
Attention networks are now a standard part of the deep learning toolkit, contributing to impressive results in neural machine translation (Bahdanau et al., 2015; Luong et al., 2015), image captioning (Xu et al., 2015), speech recognition (Chorowski et al., 2015; Chan et al., 2015), question answering (Hermann et al., 2015; Sukhbaatar et al., 2015), and algorithm-learning (Graves et al., 2014; Vinyals et al., 2015), among many other applications (see Cho et al. (2015) for a comprehensive review). This approach alleviates the bottleneck of compressing a source into a fixed-dimensional vector by equipping a model with variable-length memory (Weston et al., 2014; Graves et al., 2014; 2016), thereby providing random access into the source as needed. Attention is implemented as a hidden layer which computes a categorical distribution (or hierarchy of categorical distributions) to make a soft-selection over source elements.
Noting the empirical effectiveness of attention networks, we also observe that the standard attentionbased architecture does not directly model any structural dependencies that may exist among the source elements, and instead relies completely on the hidden layers of the network. While one might argue that these structural dependencies can be learned implicitly by a deep model with enough data, in practice, it may be useful to provide a structural bias. Modeling structural dependencies at the final, output layer has been shown to be important in many deep learning applications, most notably in seminal work on graph transformers (LeCun et al., 1998), key work on NLP (Collobert et al., 2011), and in many other areas (Peng et al., 2009; Do & Artiéres, 2010; Jaderberg et al., 2014; Chen et al., 2015; Durrett & Klein, 2015; Lample et al., 2016, inter alia).
In this work, we consider applications which may require structural dependencies at the attention layer, and develop internal structured layers for modeling these directly. This approach generalizes categorical soft-selection attention layers by specifying possible structural dependencies in a soft
∗Equal contribution.
manner. Key applications will be the development of an attention function that segments the source input into subsequences and one that takes into account the latent recursive structure (i.e. parse tree) of a source sentence.
Our approach views the attention mechanism as a graphical model over a set of latent variables. The standard attention network can be seen as an expectation of an annotation function with respect to a single latent variable whose categorical distribution is parameterized to be a function of the source. In the general case we can specify a graphical model over multiple latent variables whose edges encode the desired structure. Computing forward attention requires performing inference to obtain the expectation of the annotation function, i.e. the context vector. This expectation is computed over an exponentially-sized set of structures (through the machinery of graphical models/structured prediction), hence the name structured attention network. Notably each step of this process (including inference) is differentiable, so the model can be trained end-to-end without having to resort to deep policy gradient methods (Schulman et al., 2015).
The differentiability of inference algorithms over graphical models has previously been noted by various researchers (Li & Eisner, 2009; Domke, 2011; Stoyanov et al., 2011; Stoyanov & Eisner, 2012; Gormley et al., 2015), primarily outside the area of deep learning. For example, Gormley et al. (2015) treat an entire graphical model as a differentiable circuit and backpropagate risk through variational inference (loopy belief propagation) for minimium risk training of dependency parsers. Our contribution is to combine these ideas to produce structured internal attention layers within deep networks, noting that these approaches allow us to use the resulting marginals to create new features, as long as we do so a differentiable way.
We focus on two classes of structured attention: linear-chain conditional random fields (CRFs) (Lafferty et al., 2001) and first-order graph-based dependency parsers (Eisner, 1996). The initial work of Bahdanau et al. (2015) was particularly interesting in the context of machine translation, as the model was able to implicitly learn an alignment model as a hidden layer, effectively embedding inference into a neural network. In similar vein, under our framework the model has the capacity to learn a segmenter as a hidden layer or a parser as a hidden layer, without ever having to see a segmented sentence or a parse tree. Our experiments apply this approach to a difficult synthetic reordering task, as well as to machine translation, question answering, and natural language inference. We find that models trained with structured attention outperform standard attention models. Analysis of learned representations further reveal that interesting structures emerge as an internal layer of the model. All code is available at http://github.com/harvardnlp/struct-attn.
2 BACKGROUND: ATTENTION NETWORKS
A standard neural network consist of a series of non-linear transformation layers, where each layer produces a fixed-dimensional hidden representation. For tasks with large input spaces, this paradigm makes it hard to control the interaction between components. For example in machine translation, the source consists of an entire sentence, and the output is a prediction for each word in the translated sentence. Utilizing a standard network leads to an information bottleneck, where one hidden layer must encode the entire source sentence. Attention provides an alternative approach.1 An attention network maintains a set of hidden representations that scale with the size of the source. The model uses an internal inference step to perform a soft-selection over these representations. This method allows the model to maintain a variable-length memory and has shown to be crucially important for scaling systems for many tasks.
Formally, let x = [x1, . . . , xn] represent a sequence of inputs, let q be a query, and let z be a categorical latent variable with sample space {1, . . . , n} that encodes the desired selection among these inputs. Our aim is to produce a context c based on the sequence and the query. To do so, we assume access to an attention distribution z ∼ p(z |x, q), where we condition p on the inputs x and a query q. The context over a sequence is defined as expectation, c = Ez∼p(z | x,q)[f(x, z)] where f(x, z) is an annotation function. Attention of this form can be applied over any type of input, however, we will primarily be concerned with “deep” networks, where both the annotation function
1Another line of work involves marginalizing over latent variables (e.g. latent alignments) for sequence-tosequence transduction (Kong et al., 2016; Lu et al., 2016; Yu et al., 2016; 2017).
and attention distribution are parameterized with neural networks, and the context produced is a vector fed to a downstream network.
For example, consider the case of attention-based neural machine translation (Bahdanau et al., 2015). Here the sequence of inputs [x1, . . . ,xn] are the hidden states of a recurrent neural network (RNN), running over the words in the source sentence, q is the RNN hidden state of the target decoder (i.e. vector representation of the query q), and z represents the source position to be attended to for translation. The attention distribution p is simply p(z = i |x, q) = softmax(θi) where θ ∈ Rn is a parameterized potential typically based on a neural network, e.g. θi = MLP([xi;q]). The annotation function is defined to simply return the selected hidden state, f(x, z) = xz . The context vector can then be computed using a simple sum,
c = Ez∼p(z | x,q)[f(x, z)] = n∑ i=1 p(z = i |x, q)xi (1)
Other tasks such as question answering use attention in a similar manner, for instance by replacing source [x1, . . . , xn] with a set of potential facts and q with a representation of the question.
In summary we interpret the attention mechanism as taking the expectation of an annotation function f(x, z) with respect to a latent variable z ∼ p, where p is parameterized to be function of x and q.
3 STRUCTURED ATTENTION
Attention networks simulate selection from a set using a soft model. In this work we consider generalizing selection to types of attention, such as selecting chunks, segmenting inputs, or even attending to latent subtrees. One interpretation of this attention is as using soft-selection that considers all possible structures over the input, of which there may be exponentially many possibilities. Of course, this expectation can no longer be computed using a simple sum, and we need to incorporate the machinery of inference directly into our neural network.
Define a structured attention model as being an attention model where z is now a vector of discrete latent variables [z1, . . . , zm] and the attention distribution is p(z |x, q) is defined as a conditional random field (CRF), specifying the independence structure of the z variables. Formally, we assume an undirected graph structure withm vertices. The CRF is parameterized with clique (log-)potentials θC(zC) ∈ R, where the zC indicates the subset of z given by clique C. Under this definition, the attention probability is defined as, p(z |x, q; θ) = softmax( ∑ C θC(zC)), where for symmetry we
use softmax in a general sense, i.e. softmax(g(z)) = 1Z exp(g(z)) where Z = ∑ z′ exp(g(z
′)) is the implied partition function. In practice we use a neural CRF, where θ comes from a deep model over x, q.
In structured attention, we also assume that the annotation function f factors (at least) into clique annotation functions f(x, z) = ∑ C fC(x, zC). Under standard conditions on the conditional independence structure, inference techniques from graphical models can be used to compute the forwardpass expectations and the context:
c = Ez∼p(z | x,q)[f(x, z)] = ∑ C Ez∼p(zC | x,q)[fC(x, zC)]
3.1 EXAMPLE 1: SUBSEQUENCE SELECTION
Suppose instead of soft-selecting a single input, we wanted to explicitly model the selection of contiguous subsequences. We could naively apply categorical attention over all subsequences, or hope the model learns a multi-modal distribution to combine neighboring words. Structured attention provides an alternate approach.
Concretely, let m = n, define z to be a random vector z = [z1, . . . , zn] with zi ∈ {0, 1}, and define our annotation function to be, f(x, z) = ∑n i=1 fi(x, zi) where fi(x, zi) = 1{zi = 1}xi. The explicit expectation is then,
Ez1,...,zn [f(x, z)] = n∑ i=1 p(zi = 1 |x, q)xi (2)
Equation (2) is similar to equation (1)—both are a linear combination of the input representations where the scalar is between [0, 1] and represents how much attention should be focused on each input. However, (2) is fundamentally different in two ways: (i) it allows for multiple inputs (or no inputs) to be selected for a given query; (ii) we can incorporate structural dependencies across the zi’s. For instance, we can model the distribution over z with a linear-chain CRF with pairwise edges,
p(z1, . . . , zn |x, q) = softmax ( n−1∑ i=1 θi,i+1(zi, zi+1) ) (3)
where θk,l is the pairwise potential for zi = k and zi+1 = l. This model is shown in Figure 1c. Compare this model to the standard attention in Figure 1a, or to a simple Bernoulli (sigmoid) selection method, p(zi = 1 |x, q) = sigmoid(θi), shown in Figure 1b. All three of these methods can use potentials from the same neural network or RNN that takes x and q as inputs.
In the case of the linear-chain CRF in (3), the marginal distribution p(zi = 1 |x) can be calculated efficiently in linear-time for all i using message-passing, i.e. the forward-backward algorithm. These marginals allow us to calculate (2), and in doing so we implicitly sum over an exponentially-sized set of structures (i.e. all binary sequences of length n) through dynamic programming. We refer to this type of attention layer as a segmentation attention layer.
Note that the forward-backward algorithm is being used as parameterized pooling (as opposed to output computation), and can be thought of as generalizing the standard attention softmax. Crucially this generalization from vector softmax to forward-backward is just a series of differentiable steps,2 and we can compute gradients of its output (marginals) with respect to its input (potentials). This will allow the structured attention model to be trained end-to-end as part of a deep model.
3.2 EXAMPLE 2: SYNTACTIC TREE SELECTION
This same approach can be used for more involved structural dependencies. One popular structure for natural language tasks is a dependency tree, which enforces a structural bias on the recursive dependencies common in many languages. In particular a dependency tree enforces that each word in a source sentence is assigned exactly one parent word (head word), and that these assignments do not cross (projective structure). Employing this bias encourages the system to make a soft-selection based on learned syntactic dependencies, without requiring linguistic annotations or a pipelined decision.
A dependency parser can be partially formalized as a graphical model with the following cliques (Smith & Eisner, 2008): latent variables zij ∈ {0, 1} for all i 6= j, which indicates that the i-th word is the parent of the j-th word (i.e. xi → xj); and a special global constraint that rules out configurations of zij’s that violate parsing constraints (e.g. one head, projectivity).
The parameters to the graph-based CRF dependency parser are the potentials θij , which reflect the score of selecting xi as the parent of xj . The probability of a parse tree z given the sentence
2As are other dynamic programming algorithms for inference in graphical models, such as (loopy and nonloopy) belief propagation.
x = [x1, . . . , xn] is,
p(z |x, q) = softmax 1{z is valid}∑ i 6=j 1{zij = 1}θij  (4) where z is represented as a vector of zij’s for all i 6= j. It is possible to calculate the marginal probability of each edge p(zij = 1 |x, q) for all i, j inO(n3) time using the inside-outside algorithm (Baker, 1979) on the data structures of Eisner (1996).
The parsing contraints ensure that each word has exactly one head (i.e. ∑n i=1 zij = 1). Therefore if we want to utilize the soft-head selection of a position j, the context vector is defined as:
fj(x, z) = n∑ i=1 1{zij = 1}xi cj = Ez[fj(x, z)] = n∑ i=1 p(zij = 1 |x, q)xi
Note that in this case the annotation function has the subscript j to produce a context vector for each word in the sentence. Similar types of attention can be applied for other tree properties (e.g. soft-children). We refer to this type of attention layer as a syntactic attention layer.
3.3 END-TO-END TRAINING
Graphical models of this form have been widely used as the final layer of deep models. Our contribution is to argue that these networks can be added within deep networks in place of simple attention layers. The whole model can then be trained end-to-end.
The main complication in utilizing this approach within the network itself is the need to backpropagate the gradients through an inference algorithm as part of the structured attention network. Past work has demonstrated the techniques necessary for this approach (see Stoyanov et al. (2011)), but to our knowledge it is very rarely employed.
Consider the case of the simple linear-chain CRF layer from equation (3). Figure 2 (left) shows the standard forward-backward algorithm for computing the marginals p(zi = 1 |x, q; θ). If we treat the forward-backward algorithm as a neural network layer, its input are the potentials θ, and its output
after the forward pass are these marginals.3 To backpropagate a loss through this layer we need to compute the gradient of the lossLwith respect to θ,∇Lθ , as a function of the gradient of the loss with respect to the marginals, ∇Lp .4 As the forward-backward algorithm consists of differentiable steps, this function can be derived using reverse-mode automatic differentiation of the forward-backward algorithm itself. Note that this reverse-mode algorithm conveniently has a parallel structure to the forward version, and can also be implemented using dynamic programming.
However, in practice, one cannot simply use current off-the-shelf tools for this task. For one, efficiency is quite important for these models and so the benefits of handoptimizing the reverse-mode implementation still outweighs simplicity of automatic differentiation. Secondly, numerical precision becomes a major issue for structured attention networks. For computing the forward-pass and the marginals, it is important to use the standard log-space semifield over R ∪ {±∞} with binary operations (⊕ = logadd,⊗ = +) to avoid underflow of probabilities. For computing the backward-pass, we need to remain in log-
space, but also handle log of negative values (since ∇Lp could be negative). This requires extending to the signed log-space semifield over [R ∪ {±∞}] × {+,−} with special +/− operations. Table 1, based on Li & Eisner (2009), demonstrates how to handle this issue, and Figure 2 (right) describes backpropagation through the forward-backward algorithm. For dependency parsing, the forward pass can be computed using the inside-outside implementation of Eisner’s algorithm (Eisner, 1996). Similarly, the backpropagation parallels the inside-outside structure. Forward/backward pass through the inside-outside algorithm is described in Appendix B.
4 EXPERIMENTS
We experiment with three instantiations of structured attention networks on four different tasks: (a) a simple, synthetic tree manipulation task using the syntactic attention layer, (b) machine translation with segmentation attention (i.e. two-state linear-chain CRF), (c) question answering using an nstate linear-chain CRF for multi-step inference over n facts, and (d) natural language inference with syntactic tree attention. These experiments are not intended to boost the state-of-the-art for these tasks but to test whether these methods can be trained effectively in an end-to-end fashion, can yield improvements over standard selection-based attention, and can learn plausible latent structures. All model architectures, hyperparameters, and training details are further described in Appendix A.
4.1 TREE TRANSDUCTION
The first set of experiments look at a tree-transduction task. These experiments use synthetic data to explore a failure case of soft-selection attention models. The task is to learn to convert a random formula given in prefix notation to one in infix notation, e.g., ( ∗ ( + ( + 15 7 ) 1 8 ) ( + 19 0 11 ) ) ⇒ ( ( 15 + 7 ) + 1 + 8 ) ∗ ( 19 + 0 + 11 )
The alphabet consists of symbols {(, ),+, ∗}, numbers between 0 and 20, and a special root symbol $. This task is used as a preliminary task to see if the model is able to learn the implicit tree structure on the source side. The model itself is an encoder-decoder model, where the encoder is defined below and the decoder is an LSTM. See Appendix A.2 for the full model.
3Confusingly, “forward” in this case is different than in the forward-backward algorithm, as the marginals themselves are the output. However the two uses of the term are actually quite related. The forward-backward algorithm can be interpreted as a forward and backpropagation pass on the log partition function. See Eisner (2016) for further details (appropriately titled “Inside-Outside and Forward-Backward Algorithms Are Just Backprop”). As such our full approach can be seen as computing second-order information. This interpretation is central to Li & Eisner (2009).
4In general we use∇ab to denote the Jacobian of a with respect to b.
Training uses 15K prefix-infix pairs where the maximum nesting depth is set to be between 2-4 (the above example has depth 3), with 5K pairs in each depth bucket. The number of expressions in each parenthesis is limited to be at most 4. Test uses 1K unseen sequences with depth between 2-6 (note specifically deeper than train), with 200 sequences for each depth. The performance is measured as the average proportion of correct target tokens produced until the first failure (as in Grefenstette et al. (2015)).
For experiments we try using different forms of self -attention over embedding-only encoders. Let xj be an embedding for each source symbol; our three variants of the source representation x̂j are: (a) no atten, just symbol embeddings by themselves, i.e. x̂j = xj ; (b) simple attention, symbol embeddings and soft-pairing for each symbol, i.e. x̂j = [xj ; cj ] where cj = ∑n i=1 softmax(θij)xi is calculated using soft-selection; (c) structured attention, symbol embeddings and soft-parent, i.e. x̂j = [xj ; cj ] where cj = ∑n i=1 p(zij = 1 |x)xi is calculated using parsing marginals, obtained from the syntactic attention layer. None of these models use an explicit query value—the potentials come from running a bidirectional LSTM over the source, producing hidden vectors hi, and then computing
θij = tanh(s > tanh(W1hi + W2hj + b))
where s,b,W1,W2 are parameters (see Appendix A.1).
Depth No Atten Simple Structured
2 7.6 87.4 99.2 3 4.1 49.6 87.0 4 2.8 23.3 64.5 5 2.1 15.0 30.8 6 1.5 8.5 18.2
model is partially reconstructing the arithmetic tree. Figure 3 shows the attention distribution for the simple/structured models on the same source sequence, which indicates that the structured model is able to learn boundaries (i.e. parentheses).
5Thus there are two attention mechanisms at work under this setup. First, structured attention over the source only to obtain soft-parents for each symbol (i.e. self-attention). Second, standard softmax alignment attention over the source representations during decoding.
4.2 NEURAL MACHINE TRANSLATION
Our second set of experiments use a full neural machine translation model utilizing attention over subsequences. Here both the encoder/decoder are LSTMs, and we replace standard simple attention with a segmentation attention layer. We experiment with two settings: translating directly from unsegmented Japanese characters to English words (effectively using structured attention to perform soft word segmentation), and translating from segmented Japanese words to English words (which can be interpreted as doing phrase-based neural machine translation). Japanese word segmentation is done using the KyTea toolkit (Neubig et al., 2011).
The data comes from the Workshop on Asian Translation (WAT) (Nakazawa et al., 2016). We randomly pick 500K sentences from the original training set (of 3M sentences) where the Japanese sentence was at most 50 characters and the English sentence was at most 50 words. We apply the same length filter on the provided validation/test sets for evaluation. The vocabulary consists of all tokens that occurred at least 10 times in the training corpus.
The segmentation attention layer is a two-state CRF where the unary potentials at the j-th decoder step are parameterized as
θi(k) =
{ hiWhj , k = 1
0, k = 0
Here [h1, . . . ,hn] are the encoder hidden states and h′j is the j-th decoder hidden state (i.e. the query vector). The pairwise potentials are parameterized linearly with b, i.e. all together
θi,i+1(zi, zi+1) = θi(zi) + θi+1(zi+1) + bzi,zi+1
Therefore the segmentation attention layer requires just 4 additional parameters. Appendix A.3 describes the full model architecture. We experiment with three attention configurations: (a) standard simple attention, i.e. cj =∑n i=1 softmax(θi)hi; (b) sigmoid attention: multiple selection with Bernoulli random variables,
i.e. cj = ∑n i=1 sigmoid(θi)hi; (c) structured attention, encoded with normalized CRF marginals,
cj = n∑ i=1 p(zi = 1 |x, q) γ hi γ = 1 λ n∑ i=1 p(zi = 1 |x, q)
The normalization term γ is not ideal but we found it to be helpful for stable training.6 λ is a hyperparameter (we use λ = 2) and we further add an l2 penalty of 0.005 on the pairwise potentials b. These values were found via grid search on the validation set.
Simple Sigmoid Structured
CHAR 12.6 13.1 14.6 WORD 14.1 13.8 14.3
For further analysis, Figure 4 shows a visualization of the different attention mechanisms on the character-to-word setup. The simple model generally focuses attention heavily on a single character. In contrast, the sigmoid and structured models are able to spread their attention distribution on contiguous subsequences. The structured attention learns additional parameters (i.e. b) to smooth out this type of attention.
6With standard expectation (i.e. cj = ∑n i=1 p(zi = 1 |x, q)hi) we empirically observed the marginals to quickly saturate. We tried various strategies to overcome this, such as putting an l2 penalty on the unary potentials and initializing with a pretrained sigmoid attention model, but simply normalizing the marginals proved to be the most effective. However, this changes the interpretation of the context vector as the expectation of an annotation function in this case.
4.3 QUESTION ANSWERING
Our third experiment is on question answering (QA) with the linear-chain CRF attention layer for inference over multiple facts. We use the bAbI dataset (Weston et al., 2015), where the input is a set of sentences/facts paired with a question, and the answer is a single token. For many of the tasks the model has to attend to multiple supporting facts to arrive at the correct answer (see Figure 5 for an example), and existing approaches use multiple ‘hops’ to greedily attend to different facts. We experiment with employing structured attention to perform inference in a non-greedy way. As the ground truth supporting facts are given in the dataset, we are able to assess the model’s inference accuracy.
The baseline (simple) attention model is the End-To-End Memory Network (Sukhbaatar et al., 2015) (MemN2N), which we briefly describe here. See Appendix A.4 for full model details. Let x1, . . . ,xn be the input embedding vectors for the n sentences/facts and let q be the query embedding. In MemN2N, zk is the random variable for the sentence to select at the k-th inference step (i.e. k-th hop), and thus zk ∈ {1, . . . , n}. The probability distribution over zk is given by p(zk = i |x, q) = softmax((xki )>qk), and the context vector is given by ck = ∑n i=1 p(zk = i |x, q)oki , where xki ,o k i are the input and output embedding for the i-th sentence at the k-th hop, respectively. The k-th context vector is used to modify the query qk+1 = qk + ck, and this process repeats for k = 1, . . . ,K (for k = 1 we have xki = xi,q
k = q, ck = 0). The K-th context and query vectors are used to obtain the final answer. The attention mechanism for a K-hop MemN2N network can therefore be interpreted as a greedy selection of a length-K sequence of facts (i.e. z1, . . . , zK).
For structured attention, we use an n-state, K-step linear-chain CRF.7 We experiment with two different settings: (a) a unary CRF model with node potentials
θk(i) = (x k i ) >qk
7Note that this differs from the segmentation attention for the neural machine translation experiments described above, which was a K-state (with K = 2), n-step linear-chain CRF.
and (b) a binary CRF model with pairwise potentials
θk,k+1(i, j) = (x k i ) >qk + (xki ) >xk+1j + (x k+1 j ) >qk+1
The binary CRF model is designed to test the model’s ability to perform sequential reasoning. For both (a) and (b), a single context vector is computed: c = ∑ z1,...,zK
p(z1, . . . , zK |x, q)f(x, z) (unlike MemN2N which computes K context vectors). Evaluating c requires summing over all nK possible sequences of length K, which may not be practical for large values of K. However, if f(x, z) factors over the components of z (e.g. f(x, z) = ∑K k=1 fk(x, zk)) then one can rewrite the
above sum in terms of marginals: c = ∑K k=1 ∑n i=1 p(zk = i |x, q)fk(x, zk). In our experiments, we use fk(x, zk) = okzk . All three models are described in further detail in Appendix A.4.
Results We use the version of the dataset with 1K questions for each task. Since all models reduce to the same network for tasks with 1 supporting fact, they are excluded from our experiments. The number of hops (i.e. K) is task-dependent, and the number of memories (i.e. n) is limited to be at most 25 (note that many question have less than 25 facts—e.g. the example in Figure 5 has 9 facts). Due to high variance in model performance, we train 20 models with different initializations for each task and report the test accuracy of the model that performed the best on a 10% held-out validation set (as is typically done for bAbI tasks).
Results of the three different models are shown in Table 4. For correct answer seletion (Ans %), we find that MemN2N and the Binary CRF model perform similarly while the Unary CRF model does worse, indicating the importance of including pairwise potentials. We also assess each model’s ability to attend to the correct supporting facts in Table 4 (Fact %). Since ground truth supporting facts are provided for each query, we can check the sequence accuracy of supporting facts for each model (i.e. the rate of selecting the exact correct sequence of facts) by taking the highest probability sequence ẑ = argmax p(z1, . . . , zK |x, q) from the model and checking against the ground truth. Overall the Binary CRF is able to recover supporting facts better than MemN2N. This improvement is significant and can be up to two-fold as seen for task 2, 11, 13 & 17. However we observed that on many tasks it is sufficient to select only the last (or first) fact correctly to predict the answer, and thus higher sequence selection accuracy does not necessarily imply better answer accuracy (and vice versa). For example, all three models get 100% answer accuracy on task 15 but have different supporting fact accuracies.
Finally, in Figure 5 we visualize of the output edge marginals produced by the Binary CRF model for a single question in task 16. In this instance, the model is uncertain but ultimately able to select the right sequence of facts 5→ 6→ 8.
4.4 NATURAL LANGUAGE INFERENCE
The final experiment looks at the task of natural language inference (NLI) with the syntactic attention layer. In NLI, the model is given two sentences (hypothesis/premise) and has to predict their relationship: entailment, contradiction, neutral.
For this task, we use the Stanford NLI dataset (Bowman et al., 2015) and model our approach off of the decomposable attention model of Parikh et al. (2016). This model takes in the matrix of word embeddings as the input for each sentence and performs inter-sentence attention to predict the answer. Appendix A.5 describes the full model.
As in the transduction task, we focus on modifying the input representation to take into account soft parents via self-attention (i.e. intra-sentence attention). In addition to the three baselines described for tree transduction (No Attention, Simple, Structured), we also explore two additional settings: (d) hard pipeline parent selection, i.e. x̂j = [xj ;xhead(j)], where head(j) is the index of xj’s parent8; (e) pretrained structured attention: structured attention where the parsing layer is pretrained for one epoch on a parsed dataset (which was enough for convergence).
Results Results of our models are shown in Table 5. Simple attention improves upon the no attention model, and this is consistent with improvements observed by Parikh et al. (2016) with their intra-sentence attention model. The pipelined model with hard parents also slightly improves upon the baseline. Structured attention outperforms both models, though surprisingly, pretraining the syntactic attention layer on the parse trees performs worse than training it from scratch—it is possible that the pretrained attention is too strict for this task.
We also obtain the hard parse for an example sentence by running the Viterbi algorithm on the syntactic attention layer with the non-pretrained model:
$ The men are fighting outside a deli .
8The parents are obtained from running the dependency parser of Andor et al. (2016), available at https://github.com/tensorflow/models/tree/master/syntaxnet
Despite being trained without ever being exposed to an explicit parse tree, the syntactic attention layer learns an almost plausible dependency structure. In the above example it is able to correctly identify the main verb fighting, but makes mistakes on determiners (e.g. head of The should be men). We generally observed this pattern across sentences, possibly because the verb structure is more important for the inference task.
5 CONCLUSION
This work outlines structured attention networks, which incorporate graphical models to generalize simple attention, and describes the technical machinery and computational techniques for backpropagating through models of this form. We implement two classes of structured attention layers: a linear-chain CRF (for neural machine translation and question answering) and a more complicated first-order dependency parser (for tree transduction and natural language inference). Experiments show that this method can learn interesting structural properties and improve on top of standard models. Structured attention could also be a way of learning latent labelers or parsers through attention on other tasks.
It should be noted that the additional complexity in computing the attention distribution increases run-time—for example, structured attention was approximately 5× slower to train than simple attention for the neural machine translation experiments, even though both attention layers have the same asymptotic run-time (i.e. O(n)).
Embedding differentiable inference (and more generally, differentiable algorithms) into deep models is an exciting area of research. While we have focused on models that admit (tractable) exact inference, similar technique can be used to embed approximate inference methods. Many optimization algorithms (e.g. gradient descent, LBFGS) are also differentiable (Domke, 2012; Maclaurin et al., 2015), and have been used as output layers for structured prediction in energy-based models (Belanger & McCallum, 2016; Wang et al., 2016). Incorporating them as internal neural network layers is an interesting avenue for future work.
ACKNOWLEDGMENTS
We thank Tao Lei, Ankur Parikh, Tim Vieira, Matt Gormley, André Martins, Jason Eisner, Yoav Goldberg, and the anonymous reviewers for helpful comments, discussion, notes, and code. We additionally thank Yasumasa Miyamoto for verifying Japanese-English translations.
A MODEL DETAILS
A.1 SYNTACTIC ATTENTION
The syntactic attention layer (for tree transduction and natural language inference) is similar to the first-order graph-based dependency parser of Kipperwasser & Goldberg (2016). Given an input sentence [x1, . . . , xn] and the corresponding word vectors [x1, . . . ,xn], we use a bidirectional LSTM to get the hidden states for each time step i ∈ [1, . . . , n],
hfwdi = LSTM(xi,h fwd i−1) h bwd i = LSTM(xi,h bwd i+1) hi = [h fwd i ;h bwd i ]
where the forward and backward LSTMs have their own parameters. The score for xi → xj (i.e. xi is the parent of xj), is given by an MLP
θij = tanh(s > tanh(W1hi + W2hj + b))
These scores are used as input to the inside-outside algorithm (see Appendix B) to obtain the probability of each word’s parent p(zij = 1 |x), which is used to obtain the soft-parent cj for each word xj . In the non-structured case we simply have p(zij = 1 |x) = softmax(θij).
A.2 TREE TRANSDUCTION
Let [x1, . . . , xn], [y1, . . . , ym] be the sequence of source/target symbols, with the associated embeddings [x1, . . . ,xn], [y1, . . . ,ym] with xi,yj ∈ Rl. In the simplest baseline model we take the source representation to be the matrix of the symbol embeddings. The decoder is a one-layer LSTM which produces the hidden states h′j = LSTM(yj ,h ′ j−1), with h ′ j ∈ Rl. The hidden states are combined with the input representation via a bilinear map W ∈ Rl×l to produce the attention distribution used to obtain the vector mi, which is combined with the decoder hidden state as follows,
αi = expxiWh ′ j∑n
k=1 expxkWh ′ j
mi = n∑ i=1 αixi ĥj = tanh(U[mi;h ′ j ])
Here we have W ∈ Rl×l and U ∈ R2l×l. Finally, ĥj is used to to obtain a distribution over the next symbol yj+1,
p(yj+1 |x1, . . . , xn, y1, . . . , yj) = softmax(Vĥj + b)
For structured/simple models, the j-th source representation are respectively
x̂i = [ xi;
n∑ k=1 p(zki = 1 |x)xk
] x̂i = [ xi;
n∑ k=1 softmax(θki)xk ] where θij comes from the bidirectional LSTM described in A.1. Then αi and mi changed accordingly,
αi = exp x̂iWh ′ j∑n
k=1 exp x̂kWh ′ j
mi = n∑ i=1 αix̂i
Note that in this case we have W ∈ R2l×l and U ∈ R3l×l. We use l = 50 in all our experiments. The forward/backward LSTMs for the parsing LSTM are also 50-dimensional. Symbol embeddings are shared between the encoder and the parsing LSTMs.
Additional training details include: batch size of 20; training for 13 epochs with a learning rate of 1.0, which starts decaying by half after epoch 9 (or the epoch at which performance does not improve on validation, whichever comes first); parameter initialization over a uniform distribution U [−0.1, 0.1]; gradient normalization at 1 (i.e. renormalize the gradients to have norm 1 if the l2 norm exceeds 1). Decoding is done with beam search (beam size = 5).
A.3 NEURAL MACHINE TRANSLATION
The baseline NMT system is from Luong et al. (2015). Let [x1, . . . , xn], [y1, . . . , ym] be the source/target sentence, with the associated word embeddings [x1, . . . ,xn], [y1, . . . ,ym]. The encoder is an LSTM over the source sentence, which produces the hidden states [h1, . . . ,hn] where
hi = LSTM(xi,hi−1)
and hi ∈ Rl. The decoder is another LSTM which produces the hidden states h′j ∈ Rl. In the simple attention case with categorical attention, the hidden states are combined with the input representation via a bilinear map W ∈ Rl×l and this distribution is used to obtain the context vector at the j-th time step,
θi = hiWh ′ j cj = n∑ i=1 softmax(θi)hi
The Bernoulli attention network has the same θi but instead uses a sigmoid to obtain the weights of the linear combination, i.e.,
cj = n∑ i=1 sigmoid(θi)hi
And finally, the structured attention model uses a bilinear map to parameterize one of the unary potentials
θi(k) =
{ hiWh ′ j , k = 1
0, k = 0
θi,i+1(zi, zi+1) = θi(zi) + θi+1(zi+1) + bzi,zi+1
where b are the pairwise potentials. These potentials are used as inputs to the forward-backward algorithm to obtain the marginals p(zi = 1 |x, q), which are further normalized to obtain the context vector
cj = n∑ i=1 p(zi = 1 |x, q) γ hi γ = 1 λ n∑ i p(zi = 1 |x, q)
We use λ = 2 and also add an l2 penalty of 0.005 on the pairwise potentials b. The context vector is then combined with the decoder hidden state
ĥj = tanh(U[cj ;h ′ j ])
and ĥj is used to obtain the distribution over the next target word yj+1
p(yj+1 |x1, . . . , xn, y1, . . . yj) = softmax(Vĥj + b)
The encoder/decoder LSTMs have 2 layers and 500 hidden units (i.e. l = 500).
Additional training details include: batch size of 128; training for 30 epochs with a learning rate of 1.0, which starts decaying by half after the first epoch at which performance does not improve on validation; dropout with probability 0.3; parameter initialization over a uniform distribution U [−0.1, 0.1]; gradient normalization at 1. We generate target translations with beam search (beam size = 5), and evaluate with multi-bleu.perl from Moses.9
A.4 QUESTION ANSWERING
Our baseline model (MemN2N) is implemented following the same architecture as described in Sukhbaatar et al. (2015). In particular, let x = [x1, . . . , xn] represent the sequence of n facts with the associated embeddings [x1, . . . ,xn] and let q be the embedding of the query q. The embeddings
9https://github.com/moses-smt/mosesdecoder/blob/master/scripts/generic/ multi-bleu.perl
are obtained by simply adding the word embeddings in each sentence or query. The full model with K hops is as follows:
p(zk = i |x, q) = softmax((xki )>qk)
ck = n∑ i=1 p(zk = i |x, q)oki
qk+1 = qk + ck
p(y |x, q) = softmax(W(qK + cK))
where p(y |x, q) is the distribution over the answer vocabulary. At each layer, {xki } and {oki } are computed using embedding matrices Xk and Ok. We use the adjacent weight tying scheme from the paper so that Xk+1 = Ok,WT = OK . X1 is also used to compute the query embedding at the first hop. For k = 1 we have xki = xi,q k = q, ck = 0.
For both the Unary and the Binary CRF models, the same input fact and query representations are computed (i.e. same embedding matrices with weight tying scheme). For the unary model, the potentials are parameterized as
θk(i) = (x k i ) >qk
and for the binary model we compute pairwise potentials as
θk,k+1(i, j) = (x k i ) >qk + (xki ) >xk+1j + (x k+1 j ) >qk+1
The qk’s are updated simply with a linear mapping, i.e.
qk+1 = Qqk
In the case of the Binary CRF, to discourage the model from selecting the same fact again we additionally set θk,k+1(i, i) = −∞ for all i ∈ {1, . . . , n}. Given these potentials, we compute the marginals p(zk = i, zk+1 = j |x, q) using the forward-backward algorithm, which is then used to compute the context vector:
c = ∑
z1,...,zK p(z1, . . . , zK |x, q)f(x, z) f(x, z) = K∑ k=1 fk(x, zk) fk(x, zk) = o k zk
Note that if f(x, z) factors over the components of z (as is the case above) then computing c only requires evaluating the marginals p(zk |x, q). Finally, given the context vector the prediction is made in a similar fashion to MemN2N:
p(y |x, q) = softmax(W(qK + c))
Other training setup is similar to Sukhbaatar et al. (2015): we use stochastic gradient descent with learning rate 0.01, which is divided by 2 every 25 epochs until 100 epochs are reached. Capacity of the memory is limited to 25 sentences. The embedding vectors are of size 20 and gradients are renormalized if the norm exceeds 40. All models implement position encoding, temporal encoding, and linear start from the original paper. For linear start, the softmax(·) function in the attention layer is removed at the beginning and re-inserted after 20 epochs for MemN2N, while for the CRF models we apply a log(softmax(·)) layer on the qk after 20 epochs. Each model is trained separately for each task.
A.5 NATURAL LANGUAGE INFERENCE
Our baseline model/setup is essentially the same as that of Parikh et al. (2016). Let [x1, . . . , xn], [y1, . . . , ym] be the premise/hypothesis, with the corresponding input representations [x1, . . . ,xn], [y1, . . . ,ym]. The input representations are obtained by a linear transformation of the 300-dimensional pretrained GloVe embeddings (Pennington et al., 2014) after normalizing the GloVe embeddings to have unit norm.10 The pretrained embeddings remain fixed but the linear layer
10We use the GloVe embeddings pretrained over the 840 billion word Common Crawl, publicly available at http://nlp.stanford.edu/projects/glove/
(which is also 300-dimensional) is trained. Words not in the pretrained vocabulary are hashed to one of 100 Gaussian embeddings with mean 0 and standard deviation 1.
We concatenate each input representation with a convex combination of the other sentence’s input representations (essentially performing inter-sentence attention), where the weights are determined through a dot product followed by a softmax,
eij = f(xi) >f(yj) x̄i = xi; m∑ j=1 exp eij∑m k=1 exp eik yj  ȳj = [yj ; n∑ i=1 exp eij∑n k=1 exp ekj xi ]
Here f(·) is an MLP. The new representations are fed through another MLP g(·), summed, combined with the final MLP h(·) and fed through a softmax layer to obtain a distribution over the labels l,
x̄ = n∑ i=1 g(x̄i) ȳ = m∑ j=1 g(ȳj)
p(l |x1, . . . , xn, y1, . . . , ym) = softmax(Vh([x̄; ȳ]) + b)
All the MLPs have 2-layers, 300 ReLU units, and dropout probability of 0.2. For structured/simple models, we first employ the bidirectional parsing LSTM (see A.1) to obtain the scores θij . In the structured case each word representation is simply concatenated with its soft-parent
x̂i = [ xi;
n∑ k=1 p(zki = 1 |x)xk ] and x̂i (and analogously ŷj) is used as the input to the above model. In the simple case (which closely corresponds to the intra-sentence attention model of Parikh et al. (2016)), we have
x̂i = [ xi;
n∑ k=1 exp θki∑n l=1 exp θli xk ] The word embeddings for the parsing LSTMs are also initialized with GloVe, and the parsing layer is shared between the two sentences. The forward/backward LSTMs for the parsing layer are 100- dimensional.
Additional training details include: batch size of 32; training for 100 epochs with Adagrad (Duchi et al., 2011) where the global learning rate is 0.05 and sum of gradient squared is initialized to 0.1; parameter intialization over a Gaussian distribution with mean 0 and standard deviation 0.01; gradient normalization at 5. In the pretrained scenario, pretraining is done with Adam (Kingma & Ba, 2015) with learning rate equal to 0.01, and β1 = 0.9, β2 = 0.999.
B FORWARD/BACKWARD THROUGH THE INSIDE-OUTSIDE ALGORITHM
Figure 6 shows the procedure for obtaining the parsing marginals from the input potentials. This corresponds to running the inside-outside version of Eisner’s algorithm (Eisner, 1996). The intermediate data structures used during the dynamic programming algorithm are the (log) inside tables α, and the (log) outside tables β. Both α, β are of size n×n×2×2, where n is the sentence length. First two dimensions encode the start/end index of the span (i.e. subtree). The third dimension encodes whether the root of the subtree is the left (L) or right (R) index of the span. The fourth dimension indicates if the span is complete (1) or incomplete (0). We can calculate the marginal distribution of each word’s parent (for all words) in O(n3) using this algorithm.
Backward pass through the inside-outside algorithm is slightly more involved, but still takes O(n3) time. Figure 7 illustrates the backward procedure, which receives the gradient of the loss L with respect to the marginals, ∇Lp , and computes the gradient of the loss with respect to the potentials ∇Lθ . The computations must be performed in the signed log-space semifield to handle log of negative values. See section 3.3 and Table 1 for more details.

This is a solid paper that proposes to endow attention mechanisms with structure (the attention posterior probabilities becoming structured latent variables). Experiments are shown with segmental atention (as in semi-Markov models) and syntactic attention (as in projective dependency parsing), both in a synthetic task (tree transduction) and real world tasks (neural machine translation and natural language inference). There is a small gain in using structured attention over simple attention in the latter tasks. A clear accept.
The paper is very clear, the approach is novel and interesting, and the experiments seem to give a good proof of concept. However, the use of structured attention in neural MT seems doesn't seem to be fully exploited here: segmental attention could be a way of approaching neural phrase-based MT, and syntactic attention offers a way of incorporating latent syntax in MT -- these seem very promising directions. In particular it would be interesting to try to add some (semi-)supervision on these attention mechanisms (e.g. posterior marginals computed by an external parser) to see if that helps learning the attention components of the network, or at least help initializing them.
This seems to be the first interesting use of the backprop of forward-backward/inside-outside (Stoyanov et al. 2011). As stated in sec 3.3., for general probabilistic models the forward step over structured attention corresponds to the computation of first-order moments (posterior marginals) while the backprop step corresponds to second-order moments (gradients of marginals wrt log-potentials, i.e., Hessian of log-partition function). This extends the applicability of the proposed approach to arbitrary graphical models where these quantities can be computed efficiently. E.g. is there a generalized matrix-tree formula that allows to do backprop for non-projective syntax? On the negative side, I suspect the need for second-order statistics may bring some numerical instability in some problems, caused by the use of the signed log-space field. Was this seen in practice?
Minor comments/typos:
- last paragraph of sec 1: "standard attention attention"
- third paragraph of sec 3.2: "the on log-potentials"
- sec 4.1, Results: "... as it has no information about the source ordering" -- what do you mean here?