Introduction

Evaluation is a critical component of learning contextual bandit policies that can be deployed in high-risk settings. One way to perform this evaluation is to directly deploy it in the setting of interest, or use an accurate simulator evaluate how well this policy will do when deployed. Unfortunately, directly deploying this policy may be unsafe or unethical (i.e. deploying in a hospital environment), and we may not have accurate simulators of the deployment setting. In this case, we turn to Off-policy Evaluation (OPE), which estimates the value of the learned policy using an offline or cached dataset of samples that arise from a behavior policy. The challenge is that this behavior policy is distinct from the target policy, and there may be limited coverage of the behavior dataset in areas of the state-action space where the target policy frequents. Several strategies have been proposed to mitigate this dataset distribution shift.

Categories of OPE estimators

There are three broad categories of OPE estimators. The first is the direct method (DM)[1], which in the contextual bandit setting, approximates the reward function, and uses that reward model to simulate the target policy. Typically, if the reward function estimate is good (as in the reward model is fully realizable, and there is enough data to estimate it), DM is a great option for OPE. When the reward estimate is bad, DM estimators typically have high bias.

The second option is importance sampling (IS)[2], in which you re-weight samples in the behavior dataset as if they appeared from the target policy. The weight used is an importance sampling ratio, typically written as $\rho$, which compares how likely a sample is to occur in the target policy and the behavior policy. IS-based OPE estimators typically have high variance.

Notation

$\pi_e$ : target, or learned policy. Our goal is to evaluate this policy.

$\pi_b$ : behavior policy. Our offline dataset contains samples that arise from this behavior policy.

$\rho = \frac{\pi_e(a \vert s)}{\pi_b(a \vert s)}$ : importance sampling ratio comparing the likelihood of observing a sample in the target policy to the likelihood of observing it in the behavior policy

$S,A$ : state and action spaces

$s, a, r$ : observed state, action, and reward

$N$: number of samples

$d$ : state distribution

$R(s,a)$: the true reward function used for the factual dataset, with mean $\bar{R}(s,a)$ and variance $\sigma^2(s,a)$.

$D_1 \sim \mathcal{D}$: dataset used to fit the OPE estimate

$D_{\hat{R}} \sim \mathcal{D}$: dataset used to fit the reward model

$V = V(\pi_e) = \mathbb{E}_{s \sim d, a\sim\pi_e}[R(s, a)]$: the true return of the model. $\hat V$ is an estimator for this estimand.

Direct Method

The direct method is $\hat{V}^{DM} = \frac{1}{N}\sum_{i=1}^N \sum_{a \in A}\pi_e(a \vert s_i) \hat{R}(s_i, a_i)$.

Bias Derivation

First we use the linearity of expectation and definition of expectation over $D_1$. Then, we use the mean of the distribution $R$. Finally, we use the definition of the value function and the policy value.

\begin{align}\mathbb{E}_{D_1 \sim\mathcal{D}, D_{\hat{R}} \sim \mathcal{D}}[\hat{V}^{\text{DM}}] &= \mathbb{E}_{D_1 \sim\mathcal{D}} \mathbb{E}_{D_{\hat{R}} \sim \mathcal{D}} \left[ \frac{1}{N}\sum_{s_i \in D_1} \sum_{a \in A} \pi_e(a|s_i) \hat{R}(s_i,a) \right] \\
&= \mathbb{E}_{D_1 \sim \mathcal{D}} \left[ \frac{1}{N} \sum_{j=1}^{N} \sum_{a \in A} \pi_e(a|s_i) \mathbb{E}_{D_{\hat{R}} \sim \mathcal{D}} [\hat{R}(s_i,a) ] \right] \\
&= \frac{1}{N} \sum_{i=1}^{N} \mathbb{E}_{D_1 \sim\mathcal{D}} \left[ \sum_{a \in A} \pi_e(a|s_i) \mathbb{E}_{D_{\hat{R}} \sim \mathcal{D}} [\hat{R}(s_i,a) ] \right] \\
&= \frac{1}{N} \sum_{i=1}^{N} \sum_{s \in \mathcal{S}} d(s) \left( \sum_{a \in A} \pi_e(a|s) \mathbb{E}_{D_{\hat{R}} \sim \mathcal{D}} [\hat{R}(s,a) ] \right) \\
&= \frac{1}{N} \sum_{i=1}^{N} \sum_{s \in \mathcal{S}} d(s) \left( \sum_{a \in A} \pi_e(a|s) \bar{R}(s,a) \right) \\
&= \frac{1}{N} \sum_{i=1}^{N} \sum_{s \in \mathcal{S}} d(s) V^{\pi_e}(s_i) \\
&= \frac{1}{N} \sum_{i=1}^{N} V(\pi_e) \\
&= V(\pi_e)\end{align}

Variance Derivation

By law of total variance, we can decompose the variance of the direct method estimator into two terms:

\begin{align}
& \mathbb{V}_{D_{\hat{R}} \sim\mathcal{D}, D_{1} \sim \mathcal{D}}[\hat{V}^{\text{DM}}] = \underbrace{\mathbb{V}_{D_{\hat{R}} \sim\mathcal{D}} \mathbb{E}_{D_{\hat{R}} \sim \mathcal{D}} \left[\hat{V}^{\text{DM}}\right]}{(1)} + \underbrace{\mathbb{E}_{D_{\hat{R}} \sim\mathcal{D}} \mathbb{V}_{D_{\hat{R}} \sim \mathcal{D}} \left[ \hat{V}^{\text{DM}} \right]}_{(1')}
\end{align}

We can substitute the intermediate result from the proof of the bias of DM into (1):

\begin{align}
(1) &= \mathbb{V}_{D_{\hat{R}} \sim\mathcal{D}} \left[ \mathbb{E}_{D_{\hat{R}} \sim \mathcal{D}}\big[ \hat{V}^{\text{DM}} \big] \right] = \mathbb{V}_{D_{\hat{R}} \sim\mathcal{D}} \left[ \hat{R}(d_0,\pi_e) \right]
\end{align}

For (1’), we first consider the inner variance with respect to $D_{\hat{R}}$ assuming $\hat{R}$ is given:

\begin{align}
\mathbb{V}_{D_{\hat{R}} \sim \mathcal{D}} \left[ \hat{V}^{\text{DM}} \right] &= \mathbb{V}_{D_{\hat{R}} \sim \mathcal{D}} \left[ \frac{1}{N} \sum_{i=1}^{N} \sum_{a \in A} \pi_e(a|s_i) \hat{R}(s_i,a) \right] \\
&= \frac{1}{N^2} \sum_{i=1}^{N} \mathbb{V}_{s_i \sim d_0} \left[ \hat{R}(s_i,\pi_e) \right] \\
&= \frac{1}{N} \mathbb{V}_{s \sim d_0} \left[ \hat{R}(s,\pi_e) \right]
\end{align}

Substituting this into (1’):

\begin{align}
& (1') = \mathbb{E}_{D_{\hat{R}} \sim\mathcal{D}} \mathbb{V}_{D_{\hat{R}} \sim \mathcal{D}} \left[ \hat{V}^{\text{DM}} \right] \\
&= \mathbb{E}_{D_{\hat{R}} \sim\mathcal{D}} \left[ \frac{1}{N} \mathbb{V}_{s \sim d_0} \left[ \hat{R}(s,\pi_e) \right] \right] \\
&= \frac{1}{N} \mathbb{E}_{D_{\hat{R}} \sim\mathcal{D}} \left[ \mathbb{E}_{s \sim d_0} \left[ \hat{R}(s,\pi_e)^2 \right] - \mathbb{E}_{s \sim d_0} \left[ \hat{R}(s,\pi_e) \right]^2 \right]  && \\
&= \frac{1}{N} \Bigg( \underbrace{\mathbb{E}_{D_{\hat{R}} \sim\mathcal{D}} \mathbb{E}_{s \sim d_0} \left[ \hat{R}(s,\pi_e)^2 \right]}_{(2)} - \underbrace{\mathbb{E}_{D_{\hat{R}} \sim\mathcal{D}}\left[\mathbb{E}_{s \sim d_0} \left[ \hat{R}(s,\pi_e) \right]^2 \right]}_{(2')} \Bigg)
\end{align}

\begin{align}
(2) &= \mathbb{E}_{D_{\hat{R}} \sim\mathcal{D}} \mathbb{E}_{s \sim d_0} \left[ \hat{R}(s,\pi_e)^2 \right] \\
&= \mathbb{E}_{s \sim d_0} \mathbb{E}_{D_{\hat{R}} \sim\mathcal{D}} \left[ \hat{R}(s,\pi_e)^2 \right] \\
&= \mathbb{E}_{s \sim d_0} \left[ v^{\pi_e}(s)^2 + \mathbb{V}_{D_{\hat{R}}\sim\mathcal{D}}[\hat{R}(s,\pi_e)] \right] \\
&= \mathbb{E}_{s \sim d_0} \left[ {v^{\pi_e}(s)}^2 \right] + \mathbb{E}_{s \sim d_0} \left[ \mathbb{V}_{D_{\hat{R}} \sim\mathcal{D}} [ \hat{R}(s,\pi_e) ] \right] \\
&= \mathbb{E}_{s \sim d_0} \left[ {v^{\pi_e}(s)} \right]^2 + \mathbb{V}_{s \sim d_0} \left[ {v^{\pi_e}(s)} \right] + \mathbb{E}_{s \sim d_0} \left[ \mathbb{V}_{D_{\hat{R}} \sim\mathcal{D}} [ \hat{R}(s,\pi_e) ] \right] &&  \\
&= {v(\pi_e)}^2 + \mathbb{V}_{s \sim d_0} \left[ {v^{\pi_e}(s)} \right] + \mathbb{E}_{s \sim d_0} \left[ \mathbb{V}_{D{\hat{R}} \sim\mathcal{D}} [ \hat{R}(s,\pi_e) ] \right]  \\
\\
(2') &= \mathbb{E}_{D_{\hat{R}} \sim\mathcal{D}}\left[\mathbb{E}_{s \sim d_0} \left[ \hat{R}(s,\pi_e) \right]^2 \right] \\
&= \mathbb{E}_{D_{\hat{R}} \sim\mathcal{D}}\left[\hat{R}(d_0,\pi_e)^2 \right] \\
&= v(\pi_e)^2 + \mathbb{V}_{D_{\hat{R}}\sim\mathcal{D}}[\hat{R}(d_0,\pi_e)]
\end{align}

Thus,

\begin{align}
(2) - (2') &= \mathbb{V}_{s \sim d_0} \left[ {v^{\pi_e}(s)} \right] + \mathbb{E}_{s \sim d_0} \left[ \mathbb{V}_{D_{\hat{R}} \sim\mathcal{D}} [ \hat{R}(s,\pi_e) ] \right] - \mathbb{V}_{D_{\hat{R}} \sim\mathcal{D}}\left[\hat{R}(d_0,\pi_e) \right] \\
\end{align}

Putting everything together, we have:

\begin{align}
& \mathbb{V}_{D_{\hat{R}} \sim\mathcal{D}, D_{\hat{R}} \sim \mathcal{D}}[\hat{V}^{\text{DM}}] = (1) + (1') \\
&= \mathbb{V}_{D_{\hat{R}} \sim\mathcal{D}} \left[ \hat{R}(d_0,\pi_e) \right] \\
&+ \frac{1}{N} \left(\mathbb{V}_{s \sim d_0} \left[ {v^{\pi_e}(s)} \right] + \mathbb{E}_{s \sim d_0} \left[ \mathbb{V}_{D{\hat{R}} \sim\mathcal{D}} [ \hat{R}(s,\pi_e) ] \right] - \mathbb{V}_{D_{\hat{R}} \sim\mathcal{D}}\left[\hat{R}(d_0,\pi_e) \right]\right) \\
&= \frac{1}{N}\mathbb{V}_{s \sim d_0} \left[ {v^{\pi_e}(s)} \right] + \frac{1}{N}\mathbb{E}_{s \sim d_0} \left[ \mathbb{V}_{D_{\hat{R}} \sim\mathcal{D}} [ \hat{R}(s,\pi_e) ] \right] + \big(1-\frac{1}{N}\big) \mathbb{V}_{D_{\hat{R}} \sim\mathcal{D}}\left[\hat{R}(d_0,\pi_e) \right] \\
&= \frac{1}{N}\mathbb{V}_{s \sim d_0} \left[ {v^{\pi_e}(s)} \right] \frac{1}{N}\mathbb{E}_{s \sim d_0} \bigg[ \mathbb{E}_{a \sim \pi_e} \Big[\pi_e(a|s) \, \sigma_R^2(s,a) \, \mathbb{E}_{D{\hat{R}} \sim \mathcal{D}}\big[\frac{1}{N_{s,a}(D_{\hat{R}})}\big] \Big] \bigg] \\
& \qquad \qquad \qquad \qquad \;\; + \big(1-\frac{1}{N}\big) \mathbb{E}_{s \sim d_0}\mathbb{E}_{a \sim \pi_e} \left[d_0(s) \, \pi_e(a|s) \, \sigma_R^2(s,a) \, \mathbb{E}_{D_{\hat{R}} \sim \mathcal{D}}\Big[\frac{1}{N_{s,a}(D_{\hat{R}})}\Big] \right] \\
&= \frac{1}{N}\mathbb{V}_{s \sim d_0} \left[ {v^{\pi_e}(s)} \right] \\
&+ \mathbb{E}_{s \sim d_0}\mathbb{E}_{a \sim \pi_e} \left[\Big(\frac{1}{N} + \big(1-\frac{1}{N}\big)d_0(s) \Big) \, \pi_e(a|s) \, \sigma_R^2(s,a) \, \mathbb{E}_{D_{\hat{R}} \sim \mathcal{D}}\Big[\frac{1}{N_{s,a}(D_{\hat{R}})}\Big] \right]
\end{align}

Importance Sampling

The importance sampling (IS) estimator is: $\hat V^{\text{IS}} = \frac{1}{N}\sum_{i=1}^N \frac{\pi_e(a_i \vert s_i)}{\pi_b(a_i \vert s_i)} r_i$

Bias Derivation

Importance sampling estimator is unbiased.

\begin{align}\mathbb{E}_{D \sim \mathcal{D}}[\hat V^{\text{IS}}] &= \mathbb{E}_{s_i \sim d, a_i \sim \pi_b} \Bigg[\frac{1}{N}\sum_{i=1}^N \frac{\pi_e(a_i|s_i)}{\pi_b(a_i|s_i)} r_i \Bigg] \\
&= \frac{1}{N}\sum_{i=1}^N \mathbb{E}_{s_i \sim d, a_i \sim \pi_b} \Bigg[ \frac{\pi_e(a_i|s_i)}{\pi_b(a_i|s_i)} r_i \Bigg] \\
&= \frac{1}{N}\sum_{i=1}^N \mathbb{E}_{s_i \sim d} \Bigg[ \pi_b(a_i|s_i) \frac{ \pi_e(a_i|s_i)}{\pi_b(a_i|s_i)} r_i \Bigg] \\
&= \frac{1}{N}\sum_{i=1}^N \mathbb{E}_{s_i \sim d} \Bigg[  \pi_e(a_i|s_i) r_i \Bigg] \\
&= \frac{1}{N}\sum_{i=1}^N \mathbb{E}_{s_i \sim d, a_i \sim \pi_e(s_i)} \Big[ r_i \Big] \\
&= \frac{1}{N}\sum_{i=1}^N V(\pi_e) \\
&= V(\pi_e)
\end{align}

Variance Derivation

However, Importance sampling estimator is known for having a relatively large variance.

\begin{align}
\mathbb{V}_{D \sim \mathcal{D}}[\hat{V}^{\text{IS}}] &= \mathbb{V}_{s_i \sim d, a_i \sim \pi_b} \Bigg[ \frac{1}{N}\sum_{i=1}^N \frac{\pi_e(a_i|s_i)}{\pi_b(a_i|s_i)} R(s_i, a_i) \Bigg] \\
&= \frac{1}{N^2} \mathbb{V}_{s_i \sim d, a_i \sim \pi_b} \Bigg[ \sum_{i=1}^N \frac{\pi_e(a_i|s_i)}{\pi_b(a_i|s_i)} R(s_i, a_i) \Bigg] \\
&= \frac{1}{N} \mathbb{V}_{s_i \sim d, a_i \sim \pi_b} \Bigg[  \frac{\pi_e(a_i|s_i)}{\pi_b(a_i|s_i)} R(s_i, a_i) \Bigg]\\
\end{align}

Step 24 to 25 is because we know each sample is i.i.d. Let $w_i = \frac{\pi_e(a_i \vert s_i)}{\pi_b(a_i \vert s_i)}$, the empirical variance estimator is:

$$
\hat{\mathbb{V}}_{D \sim \mathcal{D}}[\hat{V}^{\text{IS}}] = \frac{1}{N} \sum_{i=1}^N \big (w_i R(s_i, a_i) - \hat V^{\text{IS}} \big)^2
$$

Bonus Content (Weighted Importance Sampling)

Weighted importance sampling estimator is also called a self-normalizing importance sampling estimator. It is a biased estimator but has smaller variance compared to IS.

WIS can be defined as (using $w_i = \frac{\pi_e(a_i \vert s_i)}{\pi_b(a_i \vert s_i)}$)

$$
\hat V^{\text{WIS}} = \frac{\sum_{i=1}^N w_i R(s_i, a_i)}{\sum_{i=1}^N w_i} = \sum_{i=1}^N \frac{w_i}{\sum_{i=1}^N w_i} R(s_i, a_i)
$$

We show a quick derivation of their bias and variance.

Bias Derivation

\begin{align}\mathbb{E}_{D \sim \mathcal{D}}[\hat V^{\text{WIS}}] &= \mathbb{E}_{s_i \sim d, a_i \sim \pi_b} \Bigg[\sum_{i=1}^N \frac{w_i}{\sum_{i=1}^N w_i} R(s_i, a_i) \Bigg] \\
&= \sum_{i=1}^N \mathbb{E}_{s_i \sim d, a_i \sim \pi_b} \Bigg[ \frac{w_i}{\sum_{i=1}^N w_i} R(s_i, a_i) \Bigg] \\
&= \sum_{i=1}^N \frac{1}{\mathbb{E}_{s_i \sim d, a_i \sim \pi_b} [\sum_{i=1}^N w_i]} \mathbb{E}_{s_i \sim d, a_i \sim \pi_b} \Bigg[ w_i R(s_i, a_i) \Bigg] \\
&=\sum_{i=1}^N \frac{1}{\sum_{i=1}^N\mathbb{E}_{s_i \sim d, a_i \sim \pi_b} [w_i]} \mathbb{E}_{s_i \sim d, a_i \sim \pi_b}\Bigg[ \frac{\pi_e(a_i|s_i)}{\pi_b(a_i|s_i)} R(s_i, a_i) \Bigg] \\
&=\sum_{i=1}^N \frac{1}{\sum_{i=1}^N\mathbb{E}_{s_i \sim d, a_i \sim \pi_b} [w_i]} V(\pi_e) \\
&= \frac{1}{\sum_{i=1}^N\mathbb{E}_{s_i \sim d, a_i \sim \pi_b} [w_i]} \sum_{i=1}^N V(\pi_e)
\end{align}

Step (29) to step (30) followed the IS bias derivation. We can see that WIS estimator is only unbiased when

$\sum_{i=1}^N\mathbb{E}_{s_i \sim d, a_i \sim \pi_b} [w_i] = N$,

which means $\mathbb{E}_{s_i \sim d, a_i \sim \pi_b} [w_i] = 1$. If the evaluation policy $\pi_e$ and behavior policy $\pi_b$ takes the exact same action, then this estimator is unbiased.

Consistency Corollary

This is a biased estimator

$$
\hat V^{\text{WIS}}  =  \underbrace{\frac{1}{\frac{1}{N}\sum_{i=1}^N w_i}}_{a} \cdot \underbrace{\frac{1}{N} \sum_{i=1}^N  w_iR(s_i, a_i)}_{b}
$$

We can show that $\hat V^{\text{WIS}}$ is consistent.

$$
\lim_{N \rightarrow \infty }\frac{1}{N}\sum_{i=1}^N w_i \stackrel{a.s.}{\rightarrow} \mathbb{E}_{s_i \sim d, a_i \sim \pi_b} \Big[ \frac{\pi_e(a_i|s_i)}{\pi_b(a_i|s_i)} \Big] = \mathbb{E}_{s_i \sim d} \Big[ \pi_e(a_i|s_i)\Big] = 1
$$

As we can see, $a \rightarrow 1$ as $N \rightarrow \infty$. We can also apply the strong law of large numbers on $b$ and see $b \rightarrow V(\pi_e)$ as $N \rightarrow \infty$. Therefore, $\hat V^{\text{WIS}} \rightarrow V(\pi_e)$ as $N \rightarrow \infty$.

Variance Derivation

Variance of an WIS estimator requires derivation over $\frac{w_i}{\sum_{i}w_i}$, which in its essence, a ratio — the variance of this ratio through the sampled dataset $D$ has an effect on the overall variance of WIS. In order to work this part out, we will first very quickly discuss the multivariate delta method:

For ease of discussion, we define two dummy variables $X = \frac{1}{N} \sum_{i} w_i R(s_i, a_i)$ and $Y = \frac{1}{N} \sum_i w_i$. It’s easy to see $\hat V^{\text{WIS}} = X/Y$. Let $\mu_x = \mathbb{E}[X]$ and $\mu_y = \mathbb{E}[Y]$. We can define a two-variable function $f(x, y) = x/y$ and its first-order Taylor Expansion to be:

$$
f(X, Y) \approx f(\mu_x, \mu_y) + \nabla_X f(\mu_x, \mu_y)(X - \mu_x) + \nabla_Y f(\mu_x, \mu_y)(Y-\mu_y)
$$

If we take variance on both side, we get:

$$
\mathbb{V}_{D \sim \mathcal{D}} \Bigg[\frac{X}{Y} \Bigg] \approx \Big(\nabla_X f(\mu_x, \mu_y) \Big)^2 \mathbb{V}[X] + \Big(\nabla_Y(\mu_x, \mu_y) \Big)^2 \mathbb{V}[Y] + 2 \nabla_X f(\mu_x, \mu_y) \nabla_Y f(\mu_x, \mu_y) \mathrm{Cov}(X, Y)
$$

Note that $\nabla_X f = \frac{1}{Y}$, therefore, $\nabla_X f(\mu_x, \mu_y) = \frac{1}{\mu_y}$. Similarly, $\nabla_Y f = \frac{X}{Y^2}$, therefore $\nabla_Y f(\mu_x, \mu_y) = \frac{\mu_x}{\mu_y^2}$. Using that, we can write out the equation above as below. We also assume $w_i \sim W, R(s_i, a_i) \sim R$ (where both $W$ and $R$ have randomness over $D$).

\begin{align}
\mathbb{V}_{D \sim \mathcal{D}}\Bigg[\frac{X}{Y} \Bigg] &\approx \frac{1}{\mu_x^2}\frac{1}{N}\mathbb{E}\Big[(WR - \mu_x)^2 \Big] + \frac{\mu_x^2}{\mu_y^4} \frac{1}{N} \mathbb{E} \Big[(W - \mu_y)^2 \Big] \\ &- 2\frac{1}{\mu_y} \frac{\mu_x}{\mu_y^2} \mathbb{E} \Big[W^2 R \Big] + 2 \frac{1}{\mu_y} \frac{\mu_x}{\mu_y^2} \mathbb{E} \Big[WR \Big] \mathbb{E} \Big[W \Big] \\
&=\frac{1}{\mu_y^2} \frac{1}{N} \Bigg[ \mathbb{E} \Big[ W^2 R^2\Big] + \frac{\mu_x^2}{\mu_y^2} \Big[ W^2\Big] - 2 \frac{\mu_x}{\mu_y} \mathbb{E} \Big[ W^2 R\Big] \Bigg] \\
&= \frac{1}{\mu_y^2} \frac{1}{N} \mathbb{E} \Big[ (WR - \frac{\mu_x}{\mu_y}W)^2\Big] \\
&= \frac{1}{N} \frac{\mathbb{E} \Big[ W^2 (R - \frac{\mu_x}{\mu_y})^2 \Big]}{\mu_y^2}
\end{align}

Note that $\mu_x / \mu_y = V^{\text{WIS}}$, therefore, we can write this as:

$$
\mathbb{V}_{D \sim \mathcal{D}} [\hat V^{\text{WIS}}] \approx \frac{1}{N}\frac{\mathbb{E}_{D \sim \mathcal{D}} \Bigg[ w_i^2 \Big(R(s_i, a_i) - \hat V^{\text{WIS}} \Big)^2 \Bigg]}{\mathbb{E}_{D \sim \mathcal{D}}\big[w_i \big]^2}
$$

The empirical variance estimator for this is:

\begin{align}
\hat{\mathbb{V}}_{D \sim \mathcal{D}}[\hat V^{\text{WIS}}] &= \frac{1}{N} \frac{\frac{1}{N} \sum_{i=1}^N w_i^2 \Big(R(s_i, a_i) - \hat V^{\text{WIS}} \Big)^2 }{(\frac{1}{N} \sum_{i=1}^N w_i)^2} \\
&= \sum_{i=1}^N \Big(\frac{w_i}{\sum_{i=1}^N w_i}\Big)^2 \Big( R(s_i, a_i) - \hat V^{\text{WIS}} \Big)^2
\end{align}

Note that unlike other variance estimator, we do not have the $1/N$ term after summation.

Reference: https://www.stat.cmu.edu/~hseltman/files/ratio.pdf

Introduction

Massive Online Coding Education has reached striking success over the past decade. Fast internet speed, improved UI design, code editors that are embedded in a browser window allow educational platforms such as Code.org to build a diverse set of courses tailored towards students of different coding experiences and interest levels (for example, Code.org offers “Star War-themed coding challenge,” and “Elsa/Frozen themed for-loop writing”). As a non-profit organization, Code.org claims to have reached over 60 million learners across the world ¹. Such organizations typically provide a variety of carefully constructed teaching materials such as videos and programming challenges.

A challenge faced by these platforms is that of grading assignments. It is well known that grading is critical to student learning ², in part because it motivates students to complete their assignments. Sometimes manual grading can be feasible in small settings, or automated grading used in simple settings such as when assignments are multiple choice or adopt a fill-in-the-blink modular coding structure. Unfortunately, many of the most exciting assignments, such as developing games or interactive apps, are also much more difficult to automatically evaluate. For such assignments, human teachers are currently needed to provide feedback and grading. This requirement, and the corresponding difficulty with scaling up manual grading, is one of the biggest obstacles of online coding education. Without automated grading systems, students who lack additional teacher resources cannot get useful feedback to help them learn and advance through the materials provided.

Programming a game that is playable is exciting for students who are learning to code. Code.org provides many game development assignments in their curriculum. In these assignments, students write JavaScript programs in a code editor embedded in the web browser. Game assignments are great for teachers to examine student’s progress as well: students not only need to grasp basic concepts like if-conditionals and for-loops but use these concepts to write the physical rules of the game world — calculate the trajectories of objects, resolve inelastic collision of two objects, and keep track of game states. To deal with all of these complexities, students need to use abstraction (functions/class) to encapsulate each functionality in order to manage this complex set of logic.

Automated grading on the code text alone can be an incredibly hard challenge, even for introductory level computer science assignments. As examples, two solutions which are only slightly different in text can have very different behaviors, and two solutions that are written in very different ways can have the same behaviors. As such, some models that people develop for grading code can be as complex as those used to understand paragraphs of natural language. But, sometimes, grading code can be even more difficult than grading an essay because coding submissions can be in different programming languages. In this situation, one must not only develop a program that can understand many programming languages, but guard against the potential that the grader is more accurate for some languages than others. A Finally, these programs must be able to generalize to new assignments because correct grading is just as necessary for the first student working on an assignment as the millionth – the collect-data, train, deploy cycle is not quite suitable in this context. We don’t have the luxury of collecting a massive amount of labeled dataset to train a fully supervised learning algorithm for each and every assignment.

In a recent paper, we circumvent these challenges by developing a method that grades assignments by playing them, without needing to look at the source code at all. Despite this different approach, our method still manages to provide scalable feedback that potentially can be deployed in a massively-online setting.

The Play to Grade Challenge

Our solution to these problems is to ignore the code text entirely and to grade an assignment by having a grading algorithm play it. We represent the underlying game of each program submission as a Markov Decision Process (MDP), which defines a state space, action space, reward function, and transition dynamics. By running each student’s program, we can build the MDP directly without needing to read or understand the underlying code. You can read more about the MDP framework here: ³.

Since all student programs are written for the same assignment, these programs should generate MDP with a lot of commonalities, such as shared state and action space. After playing the game and fully constructing the MDP for an assignment, all we need is to compare the MDP specified by the student’s program (student MDP) to the teacher’s solution (reference MDP) and determine if these two MDPs are the same. What sets this challenge apart from any other reinforcement learning problems is the fact that a classification needs to be made at the end of this agent’s interaction with this MDP — the decision of whether the MDP is the same as the reference MDP or not.

In order to solve this challenge, we present an algorithm with two components: an agent that plays the game and can reliably reach bug states, and a classifier that can recognize bug states (i.e., provide the probability of the observed state being a bug). Both components are necessary for accurate grading: an agent that reaches all states but cannot determine if any represents bugs is just as bad as a perfect classifier paired with an agent that is bad at taking actions which might cause bugs. Imagine a non-optimal agent that never catches the ball (in the example above) – this agent will never be able to test if the wall, or paddle, or goal does not behave correctly.

An ideal agent needs to produce differential trajectories, i.e., sequences of actions that can be used to differentiate two MDPs, and must contain at least one bug-triggering state if the trajectory is produced from the incorrect MDP. Therefore, we need both a correct MDP and a few incorrect MDPs to teach the agent and the classifier. These incorrect MDPs are incorrect solutions that can either be provided by the teacher, or come from manually grading a few student programs to find common issues. Although having to manually label incorrect MDPs is an annoyance, we show that the total amount of effort is generally significantly lower than grading each assignment: in fact, we show that for the task we solve in the paper, you only need 5 incorrect MDPs to reach a decent performance (see the appendix section of our paper).

Recognizing Bugs from Scratch

Here are three incorrect programs and what they look like when played. Each incorrect program behaves differently from the correct program:

• One program’s wall does not allow the ball to bounce on it.
• Another program’s goal post does not let the ball go through.
• The last program spawns 2 new balls whenever the ball bounces on the wall.

A challenge with building differential trajectories is that one must know which state is a bug triggering state. Previous works ⁴⁵⁶ have made the strong assumption that one would automatically know when they encountered a bug, potentially because they expect the game program to crash after encountering a bug. Because of this assumption, they focus their efforts on building pure-exploration agents that try to visit as many states as possible. However, in reality, bugs can be difficult to identify and do not all cause the game to crash. For example, a ball that is supposed to bounce off of a wall is now piercing through it and flying off into oblivion. These types of behavioral anomalies motivate the use of a predictive model that can take in the current game state and determine whether it is anomalous or not.

Unfortunately, training a model to predict if a state is a bug state is non-trivial. This is because, although we have labels for some MDPs, these labels are not on the state-level (i.e., not all states in an incorrect MDP are bug states). Put another way, our labels can tell us when a bug has been encountered but cannot tell us what specific action caused the bug. The determination of whether bugs exist in a program can be framed as a chicken-and-egg problem where, if bug states could be unambiguously determined one would only need to explore the state space, and if the exploration was optimal one would only need to figure out how to determine if the current state exhibited a bug.

Collaborative Reinforcement Learning

Fortunately, these types of problems can generally be solved through the expectation-maximization framework, which involves an intimate collaboration between the neural network classifier and the reinforcement learning agent. We propose collaborative reinforcement learning, an expectation-maximization approach, where we use a random agent to produce a dataset of trajectories from the correct and incorrect MDP to teach the classifier. Then the classifier would assign a score to each state indicating how much the classifier believes the state is a bug-triggering state. We use this score as reward and train the agent to reach these states as often as possible for a new dataset of trajectories to train the classifier.

After using the RL agent to interact with the MDP to produce trajectories, we can try out different ways to learn a classifier that can classify a state as a bug or not (a binary label). Choosing the right label is important because this label will become the reward function for the agent, so it can learn to reach bug states more efficiently. However, we only know if the MDP (the submitted code) is correct or broken, but we don’t have labels for the underlying states. Learning state-level labels becomes a challenge!

We tested out several types of classifiers: (1) a noisy classifier that classifies all states in a broken MDP as broken, (2) a Gaussian Mixture Model that treats all states independently, (3) a Variational Autoencoder that also treats all states independently but directly models non-linear interactions among the features, or (4) an LSTM that jointly models the teacher program as an MDP (HoareLSTM) and an LSTM that models the student program as an MDP (Contrastive HoareLSTM) – with a distance function that compares the two MDPs, borrowing distance notions from literature in MDP homomorphism⁷⁸⁹¹⁰.

In this toy environment, the agent drives a car on a 2D plane. Whenever the agent drives the car into the outer rim of this area (space between the boundary and red dotted line), a bug will cause the car to get stuck (Leftmost panel in Figure 5). Being stuck means the car’s physical movement is altered, resulting in back-and-forth movement around the same location. The bug classifier needs to recognize the resulting states (position and velocity) of the car being “stuck”, by correctly outputting a binary label (bug) for these states.

In this setting, there is only one type of bug. Most classifiers do well when the agent only drives a straight line (single-direction agent). However, when the agent randomly samples actions at each state, simpler classifiers can no longer differentiate between bug and non-bug states with high accuracy.

We can increase the difficulty of this setting to see if collaborative training can make the agent operate in the environment with an intention to trigger bugs. In this toy environment, now the bugs will only be triggered in red boxes (Leftmost panel in Figure 6 below). We can see that with only one round of collaborative training (“CT Round 1”), the performances of ALL classifiers are improved, including weaker classifiers. This is understandable, as the agent learns to gradually collect better datasets to train classifiers – and higher quality datasets lead to stronger classifiers. For example, variational auto-encoder started only with 32% precision, but it increased to 74.8% precision after 2 rounds of collaborative training.

We can also visualize how the collaborative training quickly allows the agent to learn to explore states that most-likely contain bugs by visualizing the trajectories (see figure below). Initially the agent just explores the space uniformly (blue curves), but after one round of collaborative training (CT), it learns to focus on visiting the potential bug areas (regions marked by red boxes) (red curves).

Grading Bounce

Next, we returned to the motivating example for this type of approach: grading real student submissions. With help from Code.org, we are able to verify the algorithm’s performance on a massive amount of unlabeled, ungraded student submissions. The game Bounce, from Code.org’s Course3 for students in 4th and 5th grade, provides a real-life dataset of what variations of different bugs and behaviors in student programs should look like. The dataset is compiled of 453,211 students who made an attempt on this assignment. In total, this dataset consists of 711,274 programs.

We train our agent and classifier on 10 broken programs that we wrote without looking at any of the student’s submissions. The 10 programs contain bugs that we “guess” to be most likely to occur, and we use them to train 10 agents that learn to reach bug states in these 10 programs. This means that in our training dataset, we have 1 correct program and 10 broken programs. Even with only 11 labeled programs, our agent and classifier can get 99.5% precision at identifying a bug program and 93.4-94% accuracy overall – the agent is able to trigger most of the bugs and the classifier recognizes the bug states using only 10 broken programs. Though for other games, especially more complicated games, the number of training programs will vary. We strongly believe the number is still in magnitude smaller than training supervised code-as-text algorithms. This demonstration shows the promise of reformulating code assignment grading as the Play to Grade.

What is Next?

We started this project by making the argument that sometimes it is far easier to grade a complex coding assignment not by looking at the code text but by playing it. Using Bounce, we demonstrated that in the simple task of identifying if a program has a bug or not (a binary task, nonetheless), we are able to achieve striking accuracy with only 11 labeled programs. We provide a simulator and all of the student’s programs on this Github repo.

Multi-label Bounce

One promising direction for future work is to expand beyond pass/fail binary feedback, and actually identify which bug is in the student’s program and provide that information. Our Bounce dataset enables this by providing multi-error labels, as shown in the table below. The multi-error label setting was not solved by our current algorithm and remains an open challenge!

More than One Correct Solution

Oftentimes, students create solutions that are creative. Creative solutions are different, but not wrong. For example, students can change the texture pattern of the ball or paddle; or they can make the paddle move much faster. How to set the boundary between “being creative” and “being wrong”? This is not a discussion that happens often in AI, but is of huge importance in education. Though we didn’t use the Bounce dataset to focus on the problem of understanding creativity, our work can still use distance measures to set a “tolerance threshold” to account for creativity.

For Educators

We are interested in collecting a suite of interactive coding assignments and creating a dataset for future researchers to work on this problem. Feel free to reach out to us and let us know what you would consider as important in grading and giving students feedback on their coding assignments!

Conclusion

Providing automated feedback for coding is an important area of research in computational education, and an important area for building fully autonomous coding education pipeline (that can generate coding assignment, grade assignment, and teach interactively). Providing a generalizable algorithm that can play interactive student programs in order to give feedback is an important problem for education and an exciting intellectual challenge for the reinforcement learning community. In this work, we introduce the challenge and a dataset, set up the MDP distance framework that is highly data efficient, algorithms that achieve high accuracy, and demonstrate this is a promising direction of applying machine learning to assist education.

This blog post is based on the following paper:

• “Play to Grade: Testing Coding Games as Classifying Markov Decision Process.” Allen Nie, Emma Brunskill, and Chris Piech. Advances in Neural Information Processing Systems 34 (2021).

Acknowledgements

Many thanks to Emma Brunskill, Chris Piech for their guidance on the project. Many thanks to Mike Wu, Ali Malik, Yunsung Kim, Lisa Yan, Tong Mu, and Henry Zhu for their discussion and feedback. Special thanks to code.org, and Baker Franke, for many years of collaboration and generously providing the research community with data. Thanks to Stanford Hoffman-Yee Human Centered AI grant for supporting AI in education. Thanks for the numerous rounds o edits from Megha Srivastava and Jacob Schreiber.

In the past two years, controlling generative models (such as GANs and VAEs) have been widely studied in the Computer Vision literature. The idea is that once these large capacity neural network models learn the data manifold of millions of images, it has internalized some knowledge about the world. The “knowledge” learned is often in the form of abstract concepts such as “skin tone”, “hair style” from faces, or “brightness”, “rotation angles” for general images. We can then explicitly control images generated from GAN or VAE with respect to these abstract concepts through some post-training manipulations.

Figure 1: (a) from OpenAI GLOW demo; (b) from ICLR 2020 paper: On the "steerability" of generative adversarial networks.

Despite such an impressive result, controlling generative models in text has been far more difficult. Limited successes have been achieved for properties like sentiment or tense (Hu et al.¹). Further work from Lample et al.² leveraged the domain where the data are collected to add additional attributes such as gender, race, age, and interests (like music, book, movie, etc.), and trained a neural network from scratch conditioned on these attributes. However, for vagueness, politeness, genericness, sarcasm, humor, and formality, we have made little progress in developing neural network based decoding method to control these aspects of text.

In this post, we will focus on one: Question-under-Discussion (QuD). When we communicate with each other, there is often an implicit question under discussion (hence the name). For example, when two strangers meet, the QuD would often be “tell something about yourself”, so people often start with self-introduction. However, instead of listing all the facts about oneself (starting from age, height), people usually choose to communicate more relvant facts about themselves given the situation. In a blind date, people talk about their interests, income, attitudes towards starting a family, dating preferences etc. In a business metting, people talk about their professional history.

The QuD model has been formalized in 1984 as a model of discourse. It emphasizes on the fact that we communicate the state of the world up to a partition. The partition often refers to context (or in this case, the QuD that specifies the context).

As a general discourse model, QuD model goes beyond simple conversation. It suggests that for every sentence we say or write, there is a corresponding underlying question; even for paragraphs of text that we write. QuD Model turns any paragraph of text into a self-conversation in a writer’s head.

Here is an example, if I start with the sentence “A few of us arrive at the classroom and the door is locked.” What should I write next? As you can see, QuD dictates how I write my next sentence. QuD decides what’s relevant in this situation and what’s not.

Even though all of my next sentences are direct followup of the previous sentence (often referred to as “context” or “scenario”), they are all quite different because I’m addressing different QuDs. In another word, if our generative model for text is QuD-sensitive, we can control text generation by swapping QuDs.

Maybe you can already see what we want to do: an image can serve as scenario/context, questions from VQA (Visual Question Answering dataset) can serve as QuDs , and sentence generated by the image captioner are the “next sentence” we want to generate. By changing the question, our method will change the generated caption. Cool, huh? There is no attribute classifier, no additional annotated training data, just a question and an image.

If we are able to achive this – using a QuD (natural language question) to ask the model to focus on different features of the image, we open the door to truly flexible controllable text generation. For the same image, I can ask: “Can the man run far?”, and the caption can be “This is a runner who jogs through the entire park.” Or I can ask “Does this man look tired?”, and the caption can be “The man is energetically running.”

So the idea seems fun, but is this technologically possible? Remember, there is no additional training involved. How can an image captioner that is only trained on image caption dataset be sensitive to questions from a VQA dataset? To answer this question, we need to rethink what question is and introduce Rational Speech Act (RSA) framework. A more technical overview of this is described in my new paper with with Reuben Cohn-Gordon and Chris Potts³.

QuD in Image Captioning

Figure 2: Examples from applying our method to a SoTA 6-layer encoder-decoder Transformer with 55M parameters image captioner trained on the MS COCO dataset. The base caption shows what the Transformer originally would have outputted(last column). The penultimate column (Issue-sensitive caption) shows what our method would produce. We can produce captions that try to address a question.

VQA style questions focus on three things: object, attribute of the object, relationship between objects. If we think more carefully about a useful representation of an image (let’s say if we want an image to someone who can’t see), we will describe the first picture in Figure 2 as “There’s grass on the ground. The photo is captured through a net. There is a guy wearing orange jersey, white pants, and he’s throwing a baseball. This guy is a pitcher. He’s standing on a field”. We can write this more structurally into a list, where we would list every object, its every attribute (color, weight, texture, etc.), and its relationship to other object in the picture. Needless to say, this list will be incredibly long and I personally doubt if it can ever be “complete”.

This is how we write down captions for each image – we write down a small amount of features in the image that catch our attention. That is to say: human captions attempt to address an indeterminate question (or “all” questions). The captioning system learns to map features in an image to a grammatical sequence of words. To change the caption with respect to a question, what we mean is that we want the caption to reflect a particular feature the question is seeking. If the question is “What position is this man playing?”, then the caption needs to mention a feature which is this guy’s position (i.e., “pitcher”). We call this Issue-Sensitive Image Captioning (ISIC) because we want the generated caption to resolve an issue (question).

We show some preliminary success in Figure 2. By asking questions such as “what color is the sky”, we control the image captioner to focus on a particular nature of this photo (that this is a black & white photo). Instead of describing what is the scene of the picture (“airplane taking off”), the caption model describes the meta-level information of this picture: (“a black and white photo of an airplane”).

Examining controllable text generation in a multi-modal domain (image, text) allows us to clearly define the question (through a pre-collected VQA dataset). Since image datasets are often densely annotated (with object information, attributes of objects, etc.), we can automatically evaluate whether our generated caption addresses the question (from here on, we refer to this as “resolve the issue”). However, even with dense annotations, evaluating issue resolution is still a major challenge. We evaluate our proposed method on a restricted domain of bird images (CUB-2011), but we showed some promising examples in a more general dataset such as MSCOCO (in Figure 2).

Controlling via RSA: A Simple Tour

In the previous section, we have established that in order to be question-aware, captions need to describe a feature in the image that is the answer to the question. Then we can think of the process to make our model generate question-aware captions as the process of picking out the right features that will answer the question. This is when RSA comes in.

The Rational Speech Act (RSA) is a framework of Bayesian models for cognition and language. It turns out a simple model in this family is highly revelant for machine learning tasks, since it captures the idea of referential language generation (which I will describe below). Computing this model turns out to require only simple matrix operations, which is extremely suitable and modular to apply to any large neural network model.

Source: New Yorker; University of Michigan, Law School.

Here’s the intuition on what human thought process RSA is trying to capture. If you were asked to say a word to the police so that they can pick out number 5 (but you are not allowed to say numbers), you probably would not say “a man,” nor would you say “khaki pants”. The best word you can pick (assuming you are rational) is “baseball cap” because it uniquely identifies number 5. RSA is created to emulate this thought process – when you have a group of images, you want to pick the word that best represents it (so that it’s distinguishing the target image from the rest).

The computational process for the simplest RSA model can be easily achieved through Bayes rule – probability matrix normalization. In this formulation, a simplified version of our pre-trained image captioner model is $S_0(\mathbf{w} \vert \mathbf{i})$, a row-stochastic probability matrix (probabilities in a row sum up to 1) where the rows are “images”, represented as a list of features, and the columns are the features we can pick to represent the image. How is this a captioning model? Well, you can imagine a captioning model that only produces one word (represented as an emoji in our schema).

Figure 3: A simple RSA computing process.

Without considering the goal of contrasting two other images, for the first image, we will choose either baseball cap or mountain to describe it. However, if we want to uniquely identify this image (to a police officer, or to anyone who’s “listening”) against the other two images, we would realize that the other two rows also have mountains. What’s unique about the first row is the baseball cap.

The process described in Figure 3 is extremely simple but captures the process of reasoning about plausible alternatives and then decide what to select to achieve its best unambiguous outcome. We can see in bold number that $S_1$ will choose baseball cap to describe the first image, contrasting the other two.

More formally, these two computations can be described as:

\begin{align*}
L_1(\mathbf{i}|\mathbf{w}) &= \frac{S_0(\mathbf{w}|\mathbf{i}) P(\mathbf{i})}{P(\mathbf{w})} = \frac{S_0(\mathbf{w}|\mathbf{i}) P(\mathbf{i})}{\sum_{i \in \mathcal{I}} S_0(\mathbf{w}|\mathbf{i}) P(\mathbf{i})} \\
S_1(\mathbf{w}|\mathbf{i}) &= \frac{L_1(\mathbf{i}|\mathbf{w}) P(\mathbf{w})}{P(\mathbf{i})} = \frac{L_1(\mathbf{i}|\mathbf{w}) P(\mathbf{w})}{\sum_{w \in \mathcal{V}} L_1(\mathbf{i}|\mathbf{w}) P(\mathbf{w})} \\
\end{align*}

In RSA books/papers, you often see the simplified version (skipping the normalization constant):

\begin{align*}
L_1(\mathbf{i}\vert\mathbf{w}) &\propto S_0(\mathbf{w}|\mathbf{i}) P(\mathbf{i}) \\
S_1(\mathbf{w}\vert\mathbf{i}) &\propto L_1(\mathbf{i}|\mathbf{w})P(\mathbf{w})
\end{align*}

More detailed tutorial of RSA can be found in here.

RSA + VQA: Is Our Caption Question-Aware?

Equipped with the knowledge of RSA, in order to achieve our goal, all we need to do is to pick out a set of images so that for the target image, the caption we produce will contain the feature that is the answer to the question. We will now describe a procedure to select a set of images so that our caption model and our RSA re-weighting will put high probability on features that can answer our question.

Given a target image and a question: $(\mathbf{i}, \mathbf{q})$, we can directly apply any pre-trained VQA model to get an answer. Given a lot of images and the same question, we can get a lot of answers. Some of these answers are different from the answer for the target image, some are the same. For example, with our target image: {blue baseball cap, mountain}, we can ask the following question: Does the person wear a baseball cap?, where the answer is Yes or No.

Given this answer, we can partition a list of images into two groups: images where the person is wearing a baseball cap, and images where the person is NOT wearing a baseball cap. Note that the question can be extremely general and ask about various aspects of the image.

Suppose we happen to select these 6 images from a larger group of images. First three images have baseball cap, the next three do not. Can we just naively apply RSA and hope the item we pick will be about the baseball cap? The answer is unfortunately no. Before RSA, the caption model will randomly choose between baseball cap and mountain. After RSA, it will choose mountain, which is worse.

Figure 5: Directly applying RSA is NOT question-aware (or issue-sensitive, as defined in our paper).

OK. We ran into a problem: even though our VQA model partitioned 6 images into two cells (top 3 rows and bottom 3 rows), the RSA computation is unaware of this (cell structure). What it does is treating all 5 other images (rows) as distractors and try to find what’s unique about the target image (first row) against all else, which is the mountain.

Luckily, a solution has already been worked out by Kao et al. ⁴ The idea is pretty simple, why not just add up the probability within the cell (across the column) after computing $L_1$ probability matrix? More formally, this corresponds a different $S_1$:

\begin{align*}
U_1^{\mathbf{C}}(\mathbf{i}, \mathbf{w}, \mathbf{C}) &= \log \Big( \sum_{\mathbf{i}' \in \mathcal{I}}\delta_{\mathbf{C}(\mathbf{i})=\mathbf{C}(\mathbf{i}')} L_1(\mathbf{i}'|\mathbf{w}) \Big) \\
S_1^{\mathbf{C}}(\mathbf{w} \vert \mathbf{i}, \mathbf{C}) &\propto \text{exp} \big(\alpha U_1^{\mathbf{C}}(\mathbf{i}, \mathbf{w}, \mathbf{C}) - \text{cost}(\mathbf{w}) \big)
\end{align*}

This formula redefines the pragmatic listener matrix $L_1$ as an informative utility $U_1^{\mathbf{C}}$, and compute $S_1$ probability matrix proportional to it. In the RSA literature, this is often referred to as the QuD-RSA (QuD: Question-under-Discussion). If we visualize this process with actual probability numbers, here’s the result:

Figure 6: We show the computational process of QuD-RSA.

As you can see, what we really did is just add up along the column for the original $L_1$ matrix, and then normalize over the row for the target image. This allows our RSA output to be issue-sensitive (question-aware). But wait wait wait, something is not right here! If you actually looked very closely at the $S_1^\mathbf{C}$ table, you’d realize what the RSA picked out is still wrong – it would randomly choose between skiing and mountain. What the heck is going on?

So, what QuD-RSA actually does is that, it creates an “equivalence class” between all images in the same cell. It converts the original objective, which is to “pick an item to best describe target image” to “pick an item to best describe the target cell”. QuD-RSA is designed to ignore the differences between within-cell images. Since picking mountain or skiing (neither appeared in the distractor images) would already best identify the target cell, there is no additional incentive to pick baseball cap. Suffice to say this is not what we want.

This last ingredient allows us to add a pressure to $S_1$ to select items that are shared amongst all images within the cell. Intuitively, since all images within the target cell share the same answer to the VQA question, whatever attribute/object in the image allows that answer will have a higher chance to appear in the resulting caption. This ingredient we choose to add is called information entropy, where flatter distribution (more uniformly distributed) will have a higher entropy, and peakier distribution will have a lower entropy. Since baseball cap is shared among all three images in the target cell, it will have the highest entropy.

More formally, we can write it out to combine both $U_1$ and $U_2$, with a balancing hyper-parameter $\beta \in [0, 1]$ that decides between how much weight we want to put on either utility:

\begin{align*}
U_2(\mathbf{i}, \mathbf{w}, \mathbf{C}) &= H(L_1(\mathbf{i}'|\mathbf{w}) \cdot \delta_{\mathbf{C}(\mathbf{i})=\mathbf{C}(\mathbf{i}')}) \\
S_1^{\mathbf{C}+H}(\mathbf{w} \vert \mathbf{i}, \mathbf{C}) &\propto \text{exp} \big( \alpha ((1-\beta)U_1 + \beta U_2) -\text{cost}(\mathbf{w}) \big)
\end{align*}

And computationally it can be visualized as:

Figure 6: We show the computational process of QuD-Entropy-RSA.

Now the story for generating issue-sensitive (question-aware) captions is complete. With the added entropy reward, the $S_1$ matrix will finally pick baseball cap as the answer to What is the person wearing?.

We apply this RSA re-weighting process to every decoder step, which is called Incremental RSA, described in ⁵. In a seperate blog post, I will share how to write RSA computation as a modular component to add to any image captioning algorithm.

Evaluating on Birds (CUB)

Verifying whether we are successful at “controlling” the caption is difficult. We do have a question and an answer in VQA, but how do we determine if the answer is mentioned in the caption? Even for simple attributes like color, for example, “What color is the flower?” with an answer “red”, what captions would satisfy our definition of “addressing the question”?

1. There is a scarlett-colored flower.
2. There is a red humming bird flying over a flower.

The first caption addresses the question (resolves the issue), but the second doesn’t (because “red” is modifying the bird, not the flower). This highlights the type of challenges we need to address for general domain image datasets such as MSCOCO and VQA. So, can we try an automated evaluation for a restricted domain? Turns out, yes.

We use CalTech-UC-San-Diego Bird Dataset (CUB-2011), which has 312 features of birds annotated for each image (11788 images in total), and since it’s a restricted domain, these 312 features are exhaustive for each bird. These features include 26 body parts, such as “belly color”, “bill length”, etc. We can imagine that each body part can be the focus of one question: “What is the belly color of this bird?” and “What is the bill length of this bird”?

We are able to build a simple keyword based classifier that can identify body part mentions in caption as well as attributes (modifiers) of these body parts. We use this classifier to evaluate our method to control caption generation. We refer to this as the “feature-in-text” classifier.

Instead of using VQA, we use the feature matrix annotation as guide to select birds that share similar features (that describe the body parts) with the target image, and birds that don’t share similar features. We show some generated captions below:

Figure 7: Generated captions for CUB birds. Left bracket contains images that share the same feature (under discussion) as the target image. Right bracket contains images that don't.

We begin by assessing the extent to which our issues ensitive pragmatic models produce captions that are more richly descriptive of the target image than a base neural captioner. For CUB, we can simply count how many attributes the caption specifies according to our feature-in-text classifier. More precisely, for each image and each model, we generate captions under all resolvable issues, concatenate those captions, and then use the feature-in-text classifier to obtain a list of attributes, which we can then compare to the ground truth for the image as given by the CUB dataset.

Table 1

Table 1 reports on this evaluation. Precision for all models is very high; the underlying attributes in CUB are very comprehensive, so all high-quality captioners are likely to do well by this metric. In contrast, the recall scores vary substantially, and they clearly favor the issue-sensitive models, revealing them to be substantially more descriptive than $S_0$.

In Table 2, we provide a breakdown of these scores by body part. The issue-sensitive models are clear winners for all categories. It is noteworthy that the entropy term seems to help for some categories but not others, suggesting underlying variation in the categories themselves. It’s safe to say that some categories are rarely discussed in ground-truth captions thus leading to difficulty for the captioning model to generate them.

Table 2

Our previous evaluation shows that varying the issue has a positive effect on the captions generated by our issue-sensitive models, but it does not assess whether these captions resolve individual issues in an intuitive way. We now report on an assessment that quantifies issue-sensitivity in this sense.

The question posed by this method is as follows: for a given issue, does the produced caption precisely resolve the issue? We can divide this into two sub-questions. First, does the caption resolve the issue, which is a notion of recall. Second, does the caption avoid addressing issues that are different from what we want to address, which is a notion of precision. The recall pressure is arguably more important, but the precision one can be seen as assessing how often the caption avoids irrelevant and potentially distracting information.

Overall, the scores reveal that this is a very challenging problem, which traces to the fine-grained issues that CUB supports. Entropy in here proves to be incredibly important to resolve the issue.

Table 3

Conclusion

Obviously there are a lot left to be done. We defined the task of Issue-Sensitive Image Captioning (ISIC) and developed a Bayesian pragmatic model that allows us to address this task successfully using existing datasets and pretrained image captioning systems. We see two natural extensions of this approach that might be explored.

First, one might collect a dataset that exactly matched the structure of ISIC. This could allow for more free-form, naturalistic issues to arise, and would facilitate end-to-end training of models for ISIC. Such models could complement and extend the ones we can create using existing datasets and our issue-sensitive pragmatic captioning agents.

Second, one could extend our notion of issuesensitivity to other domains. As we saw in Figure 2, questions (as texts) naturally give rise to issues in our sense where the domain is sufficiently structured, so these ideas might find applicability in the context of question answering and other areas of controllable natural language generation.

Robustness is a central concern in engineering. Our suspension bridges need to stand against strong wind and so it won’t collapse like the Tacoma Narrows Bridge [video]. Our nuclear reactors need to be highly fault tolerant so that Fukushima Daiichi incident won’t happen in the future [link].

In the second post, we will focus on this paper:

Ribeiro, Marco Tulio, Sameer Singh, and Carlos Guestrin. “Semantically equivalent adversarial rules for debugging nlp models.” Proceedings of the 56th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers). Vol. 1. 2018.

When we increasingly become reliant on a type of technology – suspension bridges, nuclear power, or in this case: NLP models, we must raise the level of trust we have in this technology. Robustness is precisely the requirement we need to place on such systems.

Early work from Jia & Liang (2017)¹ shows that NLP models are not immune to small negligible-by-human perturbation in text – a simple addition or deletion can break the model and force it to produce nonsensical answers. Other work such as Belinkov & Bisk ², Ebrahimi et al.³ showed a systematic perturbation that is to simply drop or replace a character is sufficient to break a model. Introducing noise to sequence data is not always bad: ealier work done by Xie et al.⁴ shows that training machine translation or language model with word/character-level perturbation (noising) actually improves performance.

However, it is hard to call these perturbed examples “adversarial examples” in the original conception of Ian Goodfellow⁵. This paper proposed a way to define an adversarial examples in text through two properties:

Semantic equivalence of two sentences: $\text{SemEq}(x, x’)$

Perturbed label prediction: $f(x) \not= f(x’)$

In our discussion, people point out that from a linguistic point of view, it is very difficult to define “semantic equivalence” because we don’t have a precise and objective definition of “meaning”. This is to say that even though two sentences might elicit the same effect for a particular task, they do not need to be synonymous. A more nuanced discussion of paraphrases in English can be found in What Is a Paraphrase? [link] by Bhagat & Hovy (2012). In this paper, semantic equivalence is operationalized as what humans (MTurkers) judged to be “equivalent”.

Semantically Equivalent Adversaries (SEAs)

Ribeiro et al. argue that only a sequence that satisfies both conditions is a true adversarial example in text. They translate this criteria into a conjunctive form using an indicator function:

$\text{SEA}(x, x’) = \unicode{x1D7D9}[\text{SemEq}(x, x’) \wedge f(x) \not= f(x’)] \label{1}$

In this paper, semantic equivalence is measured by the likelihood of paraphrasing, defined in multilingual multipivot paraphrasing paper from Lapata et al. (2017)⁶. Pivoting is a technique in statistical machine translation proposed by Bannard and Callison-Burch (2005)⁷: if two English strings $e_1$ and $e_2$ can be translated into the same French string $f$, then they can be assumed to have the same meaning.

The pivot scheme is depicted by the generative model on the left, which assumes conditional independence between $e_1$ and $e_2$ given $f$: $p(e_2 \vert e_1, f) = p(e_2 \vert f)$ . Multipivot is depicted by the model on the right: it translates one English sentence into multiple French sentences, and translate back to generate the parphrase. The back-translation of multipivoting can be a simple decoder average – each decoder takes a French string, and the overall output probability for the next English token is the weighted sum of the probability of every decoder.

Paraphrase Probability Reweighting

Assuming the unnormalized logit from the paraphrasing model is $\phi(x’ \vert x)$, and suppose $\prod_x$ is the set of paraphrases that the model could generate given $x$, then the probability of a particular paraphrase can be written as below:

$p(x’ \vert x) = \frac{\phi(x’ \vert x)}{\sum_{i \in \prod_x} \phi(i \vert x)}$

Note in the denominator, all sentences being generated (including generating the original sentence) share the probability mass. If a sentence has many easy-to-generate paraphrases (indicated by high $\phi$ value), then $p(x \vert x)$ will be small, as well as all other $p(x’ \vert x)$. Dividing $p(x’ \vert x)$ by $p(x \vert x)$ will get a large value (closer to 1). As for a sentence that is difficult to paraphrase, $p(x \vert x)$ should be rather large compared to $p(x’ \vert x)$, then this ratio will provide a much smaller value.

Based on this intuition, Ribeiro et al. proposed to compute a semantic score $S(x, x’)$ as a measure of the paraphrasing quality:

$S(x, x’) = \min(1, \frac{p(x’ \vert x)}{p(x \vert x)})$

$\text{SemEq}(x, x’) = \unicode{x1D7D9}[S(x, x’) \geq \tau]$

A simple schema to generate adversarial sentences that satisfy the Equation 1 is: ask the paraphrase model to generate paraphrases of a sentence $x$. Try these paraphrases if they can change the model prediction: $f(x’) \not = f(x)$.

Semantically Equivalent Adversarial Rules (SEARs)

SEAs are paraphrases that are generated for each text sequence. In this step, authors lay out steps to convert these local SEAs to global rules (SEARs). The rule defined in this paper is a simple discrete transformation $r = (a \rightarrow c)$. The example for $r = (movie \rightarrow film)$ can be $r$(“Great movie!”) = “Great film!”.

Given a pair of text $(x, x’)$ where $\text{SEA}(x, x’) = 1$, Ribeiro et al. select the minimal contiguous span of text that turn $x$ into $x’$, include the immediate context (one word before and/or after the span), and annotate the sequence with POS (Part of Speech) tags. The last step is to generate the product of combinations between raw words and their POS tags. A step-wise example is the follow:

Step 1: (What -> Which)

Step 2: (What color -> Which color)

Step 3: (What color -> Which color), (What NOUN -> Which NOUN), (WP color -> Which color), (What color -> WP color)

Since this process is applied for every pair of $(x, x’)$, and if we assume humans are only willing to go through $B$ rules, then Ribeiro et al. propose to filter the candidates such that $\vert R \vert \leq B$. The criteria would be:

1. High probability of producing semantically equivalent sentences: this is measured by a population statistic $E_{x \sim p(x)}[\text{SemEq(x, r(x))}] \geq 1 - \delta$. Simply put, by applying this rules, most $x$ in the corpus can be translated to semantically equivalent paraphrases. In the paper, $\delta = 0.1$.
2. High adversary count: rule $r$ must also generate paraphrases that will alter the prediction of the model. Additionally, the semantic similarity should be high between paraphrases. This can be measured by $\sum_{x \in X} S(x, r(x)) \text{SEA}(x, r(x))$.
3. Non-redundancy: rules should be diverse and cover as many $x$ as possible.

To satisfy criteria 2 and 3, Ribeiro et al. proposed a submodular optimization objective, which can be solved with a greedy algorithm with a theoretical guarantee to a constant factor off of the optimum.

$\max_{R, \vert R \vert <B} \sum_{x \in X} \max_{r \in R} S(x, r(x)) \text{SEA}(x, r(x))$

The overall algorithm is described below:

Experiment and Validation

The key metric Ribeiro et al. measure is the percentapge of Flips, defined as in the validation set, how many instances are predicted correctly on the validation data, but predicted incorrectly after the application of the rule.

The comment on this metric during discussion is that it does not indicate how many examples are there that are affected by this rule. For example, a rule that changes “color” to “colour” might have a Flips rate of 2.2% in VQA dataset, but this might only be that in the validation set of VQA, only 2.2% of instances contain the word color, so this rule has a 100% rate of success at generating adversarial examples.

The paper shows some really good discrete rules that can generate adversarial text examples:

Human-in-the-loop

Ribeiro et al. conducted experiments on humans. Bringing humans into the loop can serve two purposes: humans can judge if rules can actually generate paraphrases (beyond the semantic scoring model provided by Lapata et al.); humans can judge if the perturbations incurred by rules are actually meaningful.

They first judge the quality of SEA: For 100 correctly predicted instances in the validation set, they create three sets of comparison: 1). completely created by human MTurkers, referred as humans; 2). purely generated by the paraphrasing model described above as SEA; 3). Generate SEA by the algorithm, but replace the $S(x, x’)$ criteria with human judgment.

They show that SEA narrowly beats human (18% vs. 16%), but combining with human judgments, HSEA outperforms human by a large margin (24% vs. 13%).

Then they evaluate the global rules SEARs. This time, they invite “experts” to use an interactive web interface to create global rules. They define experts as students, faculties who have taken one graduate-level NLP or ML class. Stricly speaking, experts should have been linguistic students.

Experts are allowed to see immediate feedback on their rule creation: they know how many instances (out of 100) are perturbed by their rule, and how many instances have their prediction label perturbed. In order to have a fair comparison, they are asked to create as many rules as they want, but select 10 as the best. Also each expert is given roughly 15 minutes to create rules. They were also asked to evaluate SEARs and select 10 rules that most perserve semantic equivalence.

The results are not surprising. SEARs are much better at getting high flip percentage. Combined effort between human and machine is higher than the individual. They also compared the number of seconds on average it takes an expert to create rules vs. evaluating rules created by the machine.

Finally, the paper shows to fix those bugs, they can simply perturb the training set using these human-accepted rules, and they are able to reduce the percentage of error from 12.6% to 1.4% on VQA, and from 12.6% to 3.4$ on sentiment analysis.

Wrap up

This paper uses paraphrasing models as a way to measure semantic similarity and generating semantically equivalent sentences. As is mentioned in text, machine translation based paraphrasing perturbs the sentence only locally, while humans generate semantically equivalent adversaries with more significant perturbations.

Another limitation is that gradient-based adversarial example generation is directional and precise, while method proposed by this paper seems to be simple trial-and-error approach (keep generating paraphrases until one paraphrase perturbs the model prediction).

After careful human evaluation, only 4 rules for VQA and 16 rules for sentimen are accepted by humans as acceptable rules – a rather small yield compared to the effort spent at generating these rules. In the algorithm’s defense, these rules do cover more than 10% of the validation set, a nontrivial coverage.

This paper provides a clear framework and proposes clear properties that adversarial text examples should abide. This definition is very compatible with adversarial examples in computer vision. However, this framework only covers a specific type of adversarial examples. An obvious adversarial example not covered by this method would be operations such as adding or deleting sentences, which is important at attacking QA models.

There has been a lot interests on sentence representation learning, similar to the explosion of word embedding.

Based on my limited education in linguistics, I do not believe linguistis have an agreed upon and exhaustive set of “properties” that a sentence representation¹ needs to capture. As this is still an ongoing area of research, I am very excited to see more linguistic theories on discourse and how to better understand sentential relations.

This post tries to take a closer look at the principled work produced by Sanjeev Arora and his students: Linear Algebraic Structure of Word Senses, with Applications to Polysemy, that took a generative approach to computationally infer the distribution of sentence representation as a random vector.

Inspired by their theorem, I examined if a statistical test that asks this question: “can a word be recovered from its context?” can be used to provide a quantitative measure of sentence embedding quality. The experiments are conducted in the later sections.

A Generative Model of Sentence

Figure 1

Arora et al. assume a discourse random vector² $c \sim \mathcal{N}(0, \Sigma)$ that represents the meaning of a sentence, then every word of the sentence is generated by this random variable $c$. The probabilistic graphical representation is illustrated above. Then we can write the overall model as words in a window of size $n$ generated by $c$ as:

$$
P(w_1, w_2, ..., w_n \vert c) = \prod_{i=1}^n P(w_i \vert c)
$$

The probability of a word $w_i$ occur in the sequence can be described as a log-linear model:

$$
P(w_i \vert c) = \frac{\exp(c \cdot v_{w_i})}{Z_c}
$$

Note that $v_{w_i}$ indicates the actual vector representation of $w_i$, and this probability model uses dot-product distance to capture the probability of word $w_i$ appear in the sequence, normalized by a partition function $Z_c = \sum_w \exp(v_w \cdot c)$. This dictates that the discourse vector $c$ will be close to words it produces in vector space.

Based on this very simple model, a very interesting theorem can be proved. Here, I write out the actual proof (with more details than the one provided in the paper, as well as an easy illustration on what’s going on).

Theorem 1 of Arora et al. (2018)³ paper can be understood by introducing a vector $u$. For a given word $w$ and its corresponding word vector $v_w$, we can compute a vector $u$.

Figure 2

For this word $w$, it must appear in different spans of words across the entire document. A random variable of a window of $n$ words can be introduced as $s$, a span. Computationally, the vector $u$ for the word $w$ can be computed as follow:

$$
u = \frac{1}{k} \sum_{s \in \{s_1, ..., s_k\}} \frac{1}{n} \sum_{w_i \in s} v_{w_i}
$$

To even make this statement simpler, assume the above figure represents a tensor $S \in \mathcal{R}^{n \times k \times d}$, we can easily run the following Numpy operation to obtain $u$: u = np.mean(np.mean(S, axis=0), axis=1). After knowing how $u$ is computed, then we can understand Theorem 1:

$$
v_w = A u
$$

For any word, if we compute the corresponding vector $u$, the word embedding of this word can be obtained through a linear transformation (matrix multiplication) by a fixed matrix $A$. I provide some algebra walk through the proof of Theorem 1 in the paper. Readers who find it elementary or advanced can skip this block straight to the next section.

(Optional)

The proof stands as long as the generative model in Figure 1 holds. We set to show that $v_w \approx A \mathbb{E}(\frac{1}{n} \sum_{w_i \in s} v_{w_i} \vert w \in s)$. By “iterated expectation” or “law of total expectation”, we can expand the $\mathbb{E}[c_s \vert w \in s]$ as:

\begin{equation}
\mathbb{E}[c_s \vert w \in s] = \mathbb{E}[\mathbb{E}[c_s \vert s = w_1...w...w_n \vert w \in s]]
\label{first-eq}
\tag{1}
\end{equation}

The following step is to find the probability density function (pdf) of $c \vert w$: $p(c \vert w)$. In the earlier portion of the paper, we have the following equalities that we can substitute: $Z_c \approx Z \exp(|c |^2)$⁴, the probability density function of a multivariate normal distribution $c \sim (0, \Sigma)$ is $p(c) = \exp(-\frac{1}{2} c^T \Sigma^{-1}c)$, $| c |^2 = c^Tc = c^T I c$, and the log-linear model we assumed: $p(w \vert c) = \exp(c \cdot v_w)$. We can expand $p(c\vert w)$ using Bayes rule and substitute these terms in and obtain:

\begin{align*}
p(c \vert w) &\propto p(w \vert c)p(c) \\
&\propto \frac{1}{Z} \exp(v_w \cdot c - c^T(\frac{1}{2} \Sigma^{-1} + I)c) \\
\end{align*}

After obtaining the probability density function of $c \vert w$, we can think about what kind of random variable this pdf suggests, because eventually we want to know what is $\mathbb{E}(c \vert w)$, the left hand side of equation (1). Since there is a covariance matrix inverse $\Sigma^{-1}$ invovled, we can try to re-arrange the terms to make it look more like a multivariate Gaussian distribution. Since we do want to know $\mathbb{E}(c \vert w)$, we need to know what is the mean of this new distribution.

First, we ignore the covariance determinant term as it is a constant and in Arora’s setting, the covariance matrix is invertible – “if the word vectors as random variables are isotropic to some extent, it will make the covariance matrix identifiable” (identifiable = invertible). The assumption “isotropic word embedding” here means that word embedding dimensions should not be correlated with each other.

Then, all we need to do is to rearrange the terms in $p(c \vert w)$ to appear in the form of $\exp(-\frac{1}{2} (x-\mu)^T \Sigma^{-1} (x-\mu))$. By doing so, we will be able to find our $\mu$, the expectation of this pdf. Since the form $\frac{1}{2} c^T( \Sigma^{-1} + 2I)c$ looks very similar to the quadratic form that we need, we can let $B^{-1} = \Sigma^{-1} + 2I$ and let $B$ be our new covariance matrix for $c \vert w$. We can work out the equations from both sides. We let $\mu$ be the mean we want and solve for it:

\begin{align*}
p(c \vert w) &\propto \exp(-\frac{1}{2} (c-\mu)^T B^{-1} (c-\mu)) \\
&= \exp(-\frac{1}{2}(c^T B^{-1} c - c^TB^{-1}\mu - \mu^TB^{-1}c + \mu^TB^{-1}\mu))\\
p(c \vert w) &\propto \frac{1}{Z} \exp(v_w \cdot c - c^T(\frac{1}{2} \Sigma^{-1} + I)c) \\
&= \frac{1}{Z} \exp(-\frac{1}{2}(-2 v_w \cdot c + c^TB^{-1}c))
\end{align*}

Now we have two expressions of $p(c \vert w)$. We can match the terms between two equations, one term $c^TB^{-1}c$ already appears in both, but not $-2 v_w \cdot c$. However, there are two terms with negative signs in the top expansion. A trick that applies here is to just make them equal and hope things to work out – we solve for $\mu$:

$$
-2 v_w \cdot c = - c^TB^{-1}\mu - \mu^TB^{-1}c \\
-2 v_w \cdot c = -2 \mu^T B^{-1}c
$$

It is somewhat transparent that on the RHS (right hand side), $B$ needs to disappear since the LHS (left hand side) does not contain any $B$. To do that, $\mu$ should at least contain $B$ so that it cancels out with $B^{-1}$. Also the LHS has $v_w$ while RHS has none. Then the answer should be apparent: $\mu = Bv_w$. If you plug this in, the above equality works, shows that this is our $\mu$.

My stats PhD friend told me, if I saw a pdf in the form of $w^Tx - \frac{1}{2} x^TB^{-1}x$, then I can actually skip the above algebra and directly “see” this distribution of $x$ as mean $Bw$, with variance $B$.

So now, we know that $c \vert w \sim \mathcal{N}(B^{-1}v_w, B)$ where $B = (\Sigma^{-1} + 2I)^{-1}$, the posterior distribution of $c$ after conditioning on a single word in the sequence. Thus $\mathbb{E}(c \vert w) = (\Sigma^{-1} + 2I)^{-1} v_w$.

Then we want to get the pdf that describes $c|w_1, ..., w_n​$. This part is relatively straightforward, no algebra trick / insight is required. The work mostly hinges on the following expression:

$$
p(c \vert w_1, ..., w_n) \propto p(w_1,...,w_n \vert c) p(c) \propto p(c) \prod_{i=1}^n p(w_i \vert c) \\
= \frac{1}{Z^n} \exp(\sum_{i=1}^n v_{w_i}^Tc - \frac{1}{2} c^T(\Sigma^{-1} + 2nI)c)
$$

The generation of words are independent with each other conditioned on $c​$. We already know the expression of $p(w \vert c)​$. So the above equation evaluates to a form that we have already worked out before. We can skip the algebra and know that $\mathbb{E}[c \vert w_1, ..., w_n] \approx (\Sigma^{-1} + 2nI)^{-1} \sum_{i=1}^n v_{w_i}​$.

If you still recall the LHS and RHS of the Equation (1), then what we have left to conclude the proof is to plug what we have derived into the LHS and RHS. Feel free to refer to the paper since it offers a cleaner/shorter presentation.

$$
\mathbb{E}[c_s \vert w \in s] = \mathbb{E}[\mathbb{E}[c_s \vert s = w_1...w...w_n \vert w \in s]] \\
 (\Sigma^{-1} + 2I)^{-1} v_w \approx (\Sigma^{-1} + 2nI)^{-1} \sum_{i=1}^n v_{w_i}
$$

Therefore, we know that the matrix $A$ that we set out to find is now solvable by re-arranging the terms in above equations: $A = n(\Sigma^{-1} + 2I) (\Sigma^{-1} + 2nI)^{-1}$.

(Optional End)

Estimation of Linear Transformation

Now we have an analytic form of $A$, it seems that there are two ways to finding what $A$ is. The first way is to directly estimate $c$'s prior distribution in the hopes to getting $\Sigma$. The problem is that $p(c) = \sum_w p(c \vert w)p(w)$. We can easily compute $c|w$ but it's not very easy for us to compute $p(c)$.

Then we also know that $c \vert w \sim \mathcal{N}(B^{-1}v_w, B)$, and $A = n B^{-1}B$, if we can somehow estimate $B$, we can also compute $A$. Let's take a closer look at $c \vert w$. The mean of the distribution $B^{-1} v_w$ is word-dependent, but the covariance is not. This means if we want to find this distribution $c|w$ for any word, then we need to: for every word $w$, fit the posterior $c \vert w$ with a multivariate Gaussian distribution, and share the same covariance matrix across all these pdfs. This seems possible but computing matrix inverse: $ B^{-1}$ is expensive ($B$ is a 300 x 300 matrix, for a 300d word vector).

The last choice is to do monte carlo estimation on the Theorem 1’s original equation: $v_w \approx A \mathbb{E}(\frac{1}{n} \sum_{w_i \in s} v_{w_i} \vert w \in s)$, replace the expectation with sampling over window $s$ and sampling over word $w$, and find $A$ as a linear regression problem. This is also the method the original paper has chosen.

If we suppose that $u$ is the averaged discourse vectors for word $w$, then iterating through the vocabulary, we should be able to find matrix $A$ by solving the following optimization:

$$
\arg\min_A \sum_w \| A u_w - v_w \|_2^2
$$

The paper tried to fit GloVe embeddings using linear transformation and use SIF⁵ to generate context vector $c_w$. As we can see the practical fit is good in Table 2.

Figure 3

Application to Word Senses

Intuitively, Theorem 1 dictates that a word has a linear relationship (fulfilled by matrix $A$) to the average of all the context vectors this word appears in: $v_w = A \sum_s c_s$. This relationship is fully specified by $\Sigma$, the covariance of discourse random vector $c$. Theorem 2 of the paper can be interpreted as: $v_w \approx \alpha v_{s_1} + \beta v_{s_2}$, a linear combination of 2 senses (each sense is represented as a vector). We can see the parallel between this linear decomposition and the transformed average of all context that word $w$ appears in. The proof of Theorem 2 essentially expresses that if there are 2 senses for a given word, then with high probability, $\alpha$% of the context vectors should be similar to each other as they are all this word expressed in sense 1, $\beta$% of the context vectors should be similar/grouped together as they are all expressed as sense 2.

Since we do not observe the frequency ($\alpha$, $\beta$), nor do we know how many senses are in there, Arora proposed to discover senses using sparse coding⁶, finding a set of unit vectors $A_1, …,A_m$, that for any word $v_w$, it can be expressed by a small number of these unit vectors. These unit vectors are referred as the Atoms of Discourse.

Sparse coding objective and description can be found in the paper, overall, given a set of word vectors, two integers k, m with k « m, find a set of unit vectors $A_1, …,A_m$ such that:

$$
v_w = \sum_{j=1}^m \alpha_{w,j}A_j + \eta_w
$$

where at most k of the coefficients $\alpha$ are nonzero. The goal is to minimize the reconstruction error term $\sum_w \eta_w$.

$$
\sum_w \vert\vert v_w - \sum_{j=1}^m \alpha_{w,j} A_j \vert\vert_2^2
$$

Do Non-linear Sentence Embeddings Conform to this?

Arora et al.’s proof above assumed SIF sentence embeddings. SIF algorithm uses a linear combination of word embeddings to create sentence embeddings. One clear downside of SIF is that since GloVE/Word2Vec word embeddings do not contain order information, SIF sentence embeddings don’t have order-information either.

We can list out the assumptions used in Arora et al.’s model:

1. $p(w \vert c)$ is a log-linear model
2. Words $w$ in a window are generated independently by $c$.
3. $p(c)$, the prior of discourse vector $c$ is Gaussian with mean 0 and invertible covariance matrix $\Sigma$.

Each assumption has some flaws. Assumption (1) assumes a very simplistic model on how words are generated from meaning. By extending to a more complex model, such as bilinear transformation: $p(w \vert c) \propto \exp(v_w^T H c)$, we gain more expressivity and we still have an analytical expression of $c\vert w$, however we might lose the concentration of the partition function $Z_c$.

Assumption (2) is a very serious offense for syntax in language. Words definitely depend on one another – grammatical structure naturally emerges from word-to-word dependencies. The drink is cold, the choice of is is clearly influenced by the plurality of the subject. However, if we drop this assumption, we won’t be able to get the following factorization:

$$
 p(w_1,...,w_n \vert c) \propto \prod_{i=1}^n p(w_i \vert c)
$$

Instead, what we get is:

$$
p(w_1,...,w_n \vert c) \propto \prod_{i=1}^n p(w_i \vert c, w_{<i})
$$

This will not enable us to reuse the $p(w\vert c)$ expression that we derived, breaking the equality in Equation (1).

Assumption (3) dictates a Gaussian random walk model, which seems least objectionable. The discussion of this assumption is left for future posts.

So what does Theorem 1 really provide for us? Well, it obviously leads to Theorem 2, word embeddings under these assumptions are additive combination of senses. Also, learning matrix $A$ can find a semantic meaning for $u$, the averaged context vectors. Arora et al. described an algorithm that is referenced from Reisinger and Mooney (2010)⁷: compute $c_1, …, c_m$, for a word $w$ appears in $m$ contexts. Cluster these vectors and average them. The cluster center originally are not near meaningful words that suggest the sense this cluster tries to represent, but by applying $A$ to the cluster center, we obtain meaningful nearest words again:

Figure 4

What about Sentence Embedding Training Objectives?

In many cases, we want the theory to have suggestive power for applications. Many sentence embedding models have been proposed and many of them have different objectives: InferSent, DisSent⁸ trained on a specific semantic task (predicting inference or discourse relation); OpenAI GPT⁹ and ELMo¹⁰ trains with language modeling; BERT¹¹ trains with word cloze task (using context to predict center word).

From a very high-level view, Theorem 1 seems to be proposing a relationship between averaged context and a word that appeared in those contexts: a word can be “recovered” through applying a linear transformation to the averaged context vectors that contains this word. This is very different from the language-modeling objective:

$$
p(x_1,...,x_n) = \prod_{t=1}^n p(x_t | x_{<t})
$$

A cursory look at the graphical model of Arora’s model and language model reminded me of the confounding variable diagram below:

Figure 5

When there is a hidden variable (confounding variable) $c$ that governs the behavior of $a$ and $b$, without knowing the existence of $c$, we (the algorithm) can often make the mistake of inferring a correlational relationship between variable $a$ and $b$. Does the left figure look quite similar to a language model, and the right figure similar to what Arora et al. has proposed? This is not to suggest that there is no dependence between words – there clearly is, but modeling language well by conditioning on previous words should not exclude us from hypothesizing a different kind of generative model.

The language modeling objective is a decomposition of joint probability over words $p(x_1,...,x_n)$, and the decomposition itself does not include any structural bias. However, what I suggest is to model the joint probability of $p(c, x_1, ..., x_n)$, with the inclusion of an additional random variable.

BERT objective, using context to predict the prescence of a word seems quite similar to what Theorem 1 seems to suggest, allowing a word to be recovered through its context. However BERT does mask the word itself (as it makes sense), and the prediction happens on each sentence, instead of averaged context.

DisSent is very similar to BERT due to the nature of predicting word using context, even though it is narrow in scope (only predicting a small number of words).

There is one last type of sentence embedding objective that has not been thoughly investigated – sequence auto-encoding. Sequence auto-encoding objective compresses the whole sequence into a context representation, and the model is required to reconstruct each word from this context representation¹². I have not seen many paper on this objective, but maybe it is worth a shot given its approximity to Arora et al.’s proposal model.

Word Recovering Through Context

On a high-level, Theorem 1 inspires a statistical test: “can a word be recovered from its context?”. Obviously sentence embddings, unlike word embeddings, should capture both semantic and syntactic information. Word recover test can only measure the former. However, it is a task-agnostic, global measure of the quality of sentence embedding.

We consider 4 models: BERT, InferSent⁸, DisSent¹³, and SIF (Arora et al.’s original model). We evaluate on News Crawl Shuffled 2011 dataset, which contains roughly 2M sentences from news sources, and is part of the LM1B dataset. BERT uses SentencePiece to tokenize and produces its own subword unit for infrequent words. InferSent, DisSent, and SIF purely relies on GloVe embeddings. Thus, we choose vocabularies that overlap both GloVe and BERT.

Since our testing models are large, we subsample the overlapped vocabulary (12.5%) and construct a training/testing set (1000 words and their corresponding sentence embedding pairs for training, 100 words and their sentence embedding pairs for test) to learn the transformation matrix $A$. These words jointly appear in >99% of the sentences in the corpus.

For BERT, we select the [CLS] token position as the embedding of the sentence because it is trained on the next-sentence prediction task, albeit not fine-tuned on other tasks.

The y-axis plots the cosine distance between $Ax$ and $x$ on the test set, and x-axis shows the number of epochs during training. Unfortunately this is not the result I was expecting, at least not according to the theorem and experiment reported by the paper.

There are a few key differences between my experiment here and the paper:

1. I only sampled 1000 high frequency words. The paper probabily fitted for all vocabulary (global fit).
2. Data source is very different. Wikipedia allows words that have different senses to be talked about in very different context (“crane” will be talked as if it’s an animal or a heavy lifting machine). This is a special artifact of the Wikipedia.
3. I use the entire sentence to construct context embedding (throw away sentences that are longer than 30 words), while the paper used a simple window around the word.

BERT embedding has a much higher dimension (768-dim) compared to SIF embedding (300-dim), but it performed much worse than SIF embedding. The idea that BERT is simply averaging word vectors does not hold up under this test, but since BERT also mixes in positional embeddings, it complicates any simple analysis like this.

InferSent and DisSent (4096-dim) both did very well generalizing to the test set words. This test might still be flawed and it’s possible and would appreciate different inputs.

Closing Thoughts

Unwittingly at first, Word2Vec is quickly shown to be an implicit solution to a non-convex co-occurence matrix decomposition. GloVE and other word embedding methods followed the lead and grounded these methods in theory. Are sentence embedding models, let it be InferSent, DisSent, ELMo, BERT, implicit solutions to a more principled discourse model? So far, we are still expecting great work and interesting answers.

Acknowledgement

Special thanks to Jaime Roquero who discussed this paper in very detailed manner with me, and pointed out an algebra error I made in an earlier version of this draft.

Last week, there was a paper deadline, and I was tasked to build a multiclass text classifier at the same time. I was slightly overwhelmed. I have started using PyTorch on and off during the summer. So I decided to give Torchtext another chance.

Turns out, it is possible to use this library within 3 hours if you are willing to dig into the codebase and not afraid to use code analysis tools like PyCharm. Besides a slightly outdated and unfinished “tutorial” I can find on Google, there’s no other tutorial or explanatory documentation for this library.

This means it’s a perfect opportunity to write a blog post :) that will save many people the trouble of going through the source code!

Basics

Torchtext is a very powerful library that solves the preprocessing of text very well, but we need to know what it can and can’t do, and understand how each API is mapped to our inherent understanding of what should be done. An additional perk is that Torchtext is designed in a way that it does not just work with PyTorch, but with any deep learning library (for example: Tensorflow).

Let’s compile a list of tasks that text preprocessing must be able to handle. All checked boxes are functionalities provided by Torchtext.

• Train/Val/Test Split: seperate your data into a fixed train/val/test set (not used for k-fold validation)
• File Loading: load in the corpus from various formats
• Tokenization: break sentences into list of words
• Vocab: generate a vocabulary list
• Numericalize/Indexify: Map words into integer numbers for the entire corpus
• Word Vector: either initialize vocabulary randomly or load in from a pretrained embedding, this embedding must be “trimmed”, meaning we only store words in our vocabulary into memory.
• Batching: generate batches of training sample (padding is normally happening here)
• Embedding Lookup: map each sentence (which contains word indices) to fixed dimension word vectors

Here is a visual explanation on what we are doing in this process:

"The quick fox jumped over a lazy dog."
-> (tokenization)
["The", "quick", "fox", "jumped", "over", "a", "lazy", "dog", "."]

-> (vocab)
{"The" -> 0, 
"quick"-> 1, 
"fox" -> 2,
...}

-> (numericalize)
[0, 1, 2, ...]

-> (embedding lookup)
[
  [0.3, 0.2, 0.5],
  [0.6, 0., 0.1],
  ...
]

This process is actually quite easy to mess up, especially for tokenization. Researchers would spend a lot of time writing custom code for this, and in Tensorflow (not Keras), this process is excruiating because you would create some preprocessing script that handles everything before the Batching step, which was the recommended way in Tensorflow.

Torchtext standardizes this procedure and makes it very easy for people to use. Let’s assume we have split our text corpus into three tsv files (I tried to use the JSON loader but I was not able to figure out the format specifications): train.tsv, val.tsv, test.tsv.

We first define a Field, this is a class that contains information on how you want the data preprocessed. It acts like an instruction manual that data.TabularDataset will use. We define two fields:

import spacy
spacy_en = spacy.load('en')

def tokenizer(text): # create a tokenizer function
    return [tok.text for tok in spacy_en.tokenizer(text)]

TEXT = data.Field(sequential=True, tokenize=tokenizer, lower=True)
LABEL = data.Field(sequential=False, use_vocab=False)

I have preprocessed the label to be integers, so I used use_vocab=False. Torchtext probably allows you use to text as labels, but I have not used it yet. As you can see, we can let the input to be variable length, and TorchText will dynamically pad each sequence to the longest in the batch.

Then we load in all our corpuses at once using this great class method splits:

train, val, test = data.TabularDataset.splits(
        path='./data/', train='train.tsv',
        validation='val.tsv', test='test.tsv', format='tsv',
        fields=[('Text', TEXT), ('Label', LABEL)])

This is quite straightforward, in fields, the amusing part is that tsv file parsing is order-based. In my raw tsv file, I do not have any header, and this script seems to run just fine. I imagine if there is a header, the first element Text might need to match the column header. (check out comment from Keita Kurita, who pointed out current TabularDataset does not align features with column headers)

Then we load in pretrained word embeddings:

TEXT.build_vocab(train, vectors="glove.6B.100d")

Note you can directly pass in a string and it will download pre-trained word vectors and load them for you. You can also use your own vectors by using this class vocab.Vectors. The downloaded word embeddings will stay at ./.vector_cache folder. I have not yet discovered a way to specify a custom location to store the downloaded vectors (there should be a way right?).

Next we can start the batching process, where we use Torchtext’s function to build an iterator for our training, validation and testing data split. This is also made into an easy step:

train_iter, val_iter, test_iter = data.Iterator.splits(
        (train, val, test), sort_key=lambda x: len(x.Text),
        batch_sizes=(32, 256, 256), device=-1)

Note that if you are runing on CPU, you must set device to be -1, otherwise you can leave it to 0 for GPU. Note that you can easily examine the result, and realize that Torchtext already uses dynamic padding, meaning that the padded length of your batch is dependent on the longest sequence in your batch. This is a no-brainer must-do, but I’ve seen quite a few github repos/tutorials fail to cover.

TorchText Iterator is different from a normal Python iterator. They accept several keywords which we will walk through in the later advanced section. Keep in mind that each batch is of type torch.LongTensor, they are the numericalized batch, but there is no loading of the pretrained vectors.

Don’t worry. This step is still very easy to handle. We can obtain the Vocab object easily from the Field (there is a reason why each Field has its own Vocab class, because of some pecularities of Seq2Seq model like Machine Translation, but I won’t get into it right now.). We use PyTorch’s nice Embedding Layer to solve our embedding lookup problem:

vocab = TEXT.vocab
self.embed = nn.Embedding(len(vocab), emb_dim)
self.embed.weight.data.copy_(vocab.vectors)

When we call len(vocab), we are getting the total vocabulary size, and then we just copy over pretrained word vectors by calling vocab.vectors! It would normally take me half a day to write a preprocessing script that handles all of these, but with Torchtext, I was able to finish the whole classifier in 1 hour.

Reversible Tokenization

In quite many situations, you would want to examine your output, and try to interpret your model’s actions. Then you need to map a batch of your serialized tokens like [[0, 5, 15, ...], [152, 0, 50, ...] ] back to strings! Typically this is difficult to do, because standard tokenization tools like Spacy or NLTK only tokenize but do not map back for you. Even though mapping these batches would not be too difficult, Torchtext has an easy solution for it: RevTok.

RevTok is a simple Python tokenizer, which (may) run relatively slower compared to industry-ready implementations like Spacy, but it’s integrated into TorchText. You can download the source code from: https://github.com/jekbradbury/revtok and install as follows:

cd revtok/
python setup.py install

You cannot install revtok from pip install revtok because it is not the latest version of revtok (at the time this article is written), and it will break.

Once you installed revtok, you can easily use it in TorchText using:

TEXT = data.ReversibleField(sequential=True, lower=True, include_lengths=True)

When you need to revert back to your string tokens from numericalized tokens, you can:

    for data in valid_iter:
        (x, x_lengths), y = data.Text, data.Description
        orig_text = TEXT.reverse(x.data)

Here’s a bit of juicy / advanced material for you, ReversibleField can only be used jointly with revtok. This is perhaps a bug in Torchtext. If you believe your discrete data only needs simple tokenization or normal splitting, then you can build your own customized Reversible Field like the following:

from torchtext.data import Field
class SplitReversibleField(Field):

    def __init__(self, **kwargs):
        if kwargs.get('tokenize') is list:
            self.use_revtok = False
        else:
            self.use_revtok = True
        if kwargs.get('tokenize') not in ('revtok', 'subword', list):
            kwargs['tokenize'] = 'revtok'
        if 'unk_token' not in kwargs:
            kwargs['unk_token'] = ' UNK '
        super(SplitReversibleField, self).__init__(**kwargs)

    def reverse(self, batch):
        if self.use_revtok:
            try:
                import revtok
            except ImportError:
                print("Please install revtok.")
                raise
        if not self.batch_first:
            batch = batch.t()
        with torch.cuda.device_of(batch):
            batch = batch.tolist()
        batch = [[self.vocab.itos[ind] for ind in ex] for ex in batch]  # denumericalize

        def trim(s, t):
            sentence = []
            for w in s:
                if w == t:
                    break
                sentence.append(w)
            return sentence

        batch = [trim(ex, self.eos_token) for ex in batch]  # trim past frst eos

        def filter_special(tok):
            return tok not in (self.init_token, self.pad_token)

        batch = [filter(filter_special, ex) for ex in batch]
        if self.use_revtok:
            return [revtok.detokenize(ex) for ex in batch]
        return [' '.join(ex) for ex in batch]

TorchText Iterators for masked BPTT

In the basic part of the tutorial, we have already used Torchtext Iterators, but the customizable parts of the Torchtext Iterator that are truly helpful.

We talk about three main keywords: sort, sort_within_batch and repeat.

First, PyTorch’s current solution for masked BPTT is slightly bizzare, it requires you to pack the PyTorch variables into a padded sequences. Note that not all PyTorch RNN libraries support padded sequence, for example, SRU does not, and even though I haven’t seen issues being raised, but possibly current implementation of QRNN doesn’t support padded sequence class either.

The code to incorporate padded sequence for RNN is simple:

TEXT = data.ReversibleField(sequential=True, lower=True, include_lengths=True)

We first add include_lengths argument for the Field. Then, we build our iterator like:

train, val, test = data.TabularDataset.splits(
        path='./data/', train='train.tsv',
        validation='val.tsv', test='test.tsv', format='tsv',
        fields=[('Text', TEXT), ('Label', LABEL)])
        
train_iter, val_iter, test_iter = data.Iterator.splits(
        (train, val, test), sort_key=lambda x: len(x.Text), 
        batch_sizes=(32, 256, 256), device=args.gpu, 
        sort_within_batch=True, repeat=False)

So in here, we look at a couple of arguments: sort_key is the sorting function Torchtext will call when it attempts to sort your dataset. If you leave this blank, no sorting will happen (I could be wrong, but on my simple “experiment”, it seems to be the case).

sort argument sorts through your entire dataset. You might want to avoid this for training set because the order of your traning examples will have an impact on your network, and my experiments have shown negative impact.

sort_within_batch must be flagged True if you want to use PyTorch’s padded sequence class.

So in your model’s forward method, you can write:

 def forward(self, input, lengths=None):
        embed_input = self.embed(input)

        packed_emb = embed_input
        if lengths is not None:
            lengths = lengths.view(-1).tolist()
            packed_emb = nn.utils.rnn.pack_padded_sequence(embed_input, lengths)

        output, hidden = self.encoder(packed_emb)  # embed_input

        if lengths is not None:
            output = unpack(output)[0]

And in your main training loop you can pass your variable in like this:

(x, x_lengths), y = data.Text, data.Description
output = model(x, x_lengths)

repeat is an interesting argument which is set to None by default, which means the iterator you get will repeat if it is the train iterator, and will not for validation and test iterator. What is the advantage of repeat=None? If you don’t want an epoch / iteration difference (epoch = a full traversal of your dataset, iteration = processed one batch from your dataset), and you want a simple while loop and a stop condition, then you should use repeat. Some pseudo code for train_iter that is generated by default:

# Note this loop will go on FOREVER
for val_i, data in enumerate(train_iter):
  (x, x_lengths), y = data.Text, data.Description
  
  # terminate condition, when loss converges or it reaches 50000 iterations
  if loss converges or val_i == 50000:
  	break

I personally love the concept of epoch. I adjust my learning rate according to epoch, I run on validation once the algorithm finishes every epoch. If you use repeat=None, then you have to validate every 1000 iterations (or any number you think is appropriate), and adjust your learning rate accordingly.

When you set repeat=False, which is what I provided up there, the loop will break automatically when the entire dataset has been traversed. So you write can write this:

# Note this loop will stop when training data is traversed once
epochs = 10

for epoch in range(epochs):
  for data in train_iter:
    (x, x_lengths), y = data.Text, data.Description
    # model running...

So this will train for precisely 10 epochs.

Initialize Unknown Words Randomly

If most of your vocabulary are OOV, then you might want to initialize them into a random vector. There isn’t a good way to do this, at least not in the current version of Torchtext, which only initializes every unk token to 0. I wrote a custom function that initializes unknown word embeddings to be a random vector with norm equal to the average norm of pretrained vectors, which means your variance should be set as average_norm / dim_vector.

def init_emb(vocab, init="randn", num_special_toks=2):
    emb_vectors = vocab.vectors
    sweep_range = len(vocab)
    running_norm = 0.
    num_non_zero = 0
    total_words = 0
    for i in range(num_special_toks, sweep_range):
        if len(emb_vectors[i, :].nonzero()) == 0:
            # std = 0.05 is based on the norm of average GloVE 100-dim word vectors
            if init == "randn":
                torch.nn.init.normal(emb_vectors[i], mean=0, std=0.05)
        else:
            num_non_zero += 1
            running_norm += torch.norm(emb_vectors[i])
        total_words += 1
    logger.info("average GloVE norm is {}, number of known words are {}, total number of words are {}".format(
        running_norm / num_non_zero, num_non_zero, total_words))

I hope people would find this post useful, and here are a list of references I’ve used for basic tutorial:

An (old) Torchtext Tutorial: https://github.com/mjc92/TorchTextTutorial/blob/master/01.%20Getting%20started.ipynb

Randomly Initializing Word Embeddings: https://github.com/pytorch/text/issues/32

Update 1: Added the section of masked BPTT

Update 2: Fixed typos, etc., add some code to randomly initialize unknown vectors.

I’m also including a list of tools that I find helpful and can be used in conjunction with TorchText:
Ftfy: https://ftfy.readthedocs.io/en/latest/ (apparently this fixes your text encoding problems for you, and normalize many strange artifacts in text)
BPE (Byte-pair encoding): https://github.com/rsennrich/subword-nmt (use this to generate BPE for reduced vocabulary)

This is an article providing another perspective on understanding Lagrangian and dual problem. These two topics are essential to convex and non-convex optimization. Since it is a blog post, the proper background to understand this article is kept rather low. If you need to brush up on convex knowledge, or wish to be introduced to the subject, there is a handout from Stanford CS229 website, and it’s written in a very concise and clear way: Convex Optimization (I) and Convex Optimization (II). However, you don’t need to understand convex sets and functions in order to follow this post. If you are like me, who took Stanford EE364A, then you might find this post more interesting. The post randomly jumps between Chapter 5, Chapter 10, and 11 in the book (Convex Optimization book by Stephen Boyd), but mostly focus on presenting Chapter 5. This post is also inspired by an office hour conversation with Nicholas Moehle.

Lagrangian

In convex optimization, unlike other fields of machine learning, we can add constraints. Convex problems have a form of the following:

\begin{equation*}
\begin{aligned}
& \text{minimize}
& & f_0(x) \\
& \text{subject to}
& & f_i(x) \leq 0, \; i = 1, \ldots, m. \\
&
& & h_i(x) = 0, \; i = 1, \ldots, p.
\end{aligned}
\end{equation*}

We define $f_i(x)​$ being convex functions, and $h_i(x)​$ being affine functions. If you are familiar with convex optimization or other optimization problem, this formula should look familar. The constraints listed under subject to are considered hard constraints, and those constraints do not normally appear in traditonal machine learning cost functions. One way to distinguish them is that, cost functions in Machine Learning does not need to have a realistic meaning (objective can be an abstract measurement of progress), where in convex optimization, objective is something you want to find (like how much fuel the rocket will need).

However, it has been quite difficult to optimize with these hard constraints. Typical methods include solving for a Newton step:

\begin{equation*}
\begin{bmatrix}
    \nabla^2f(x)  & A \\
    A & 0
  \end{bmatrix}
    \begin{bmatrix}
    \Delta x_{nt} \\
    w
  \end{bmatrix} =
  \begin{bmatrix}
    - \nabla f(x) \\
    0
  \end{bmatrix}
\end{equation*}

Solving for $\Delta x_{nt}$ and we get our central-path (feasible) Newton method of satisfying equality constraint as well as optimizing our overall objective. As to why it has such effect, refer to the book.

As to satisfy the inequality constraint, we can use barrier method, which adds the inequality constraints as a logarithmic barrier, so instead of optimizing the original problem, we can have:

\begin{equation*}
\begin{aligned}
& \text{minimize}
& & f_0(x) - (1/t) \sum_{i=1}^m \text{log}(-f_i) \\
& \text{subject to}
& & h_i(x) = 0, \; i = 1, \ldots, p.
\end{aligned}
\end{equation*}

The logarithmic barrier punishes the objective logarithmically when it’s approaching 0 from the negative side.

So what do those two “implementation” details have to do with Lagrangian? The reason is simple: Lagrangian is simply a way, not unlike equality Newton step or barrier method, to incoroporate constraints into the overall objective. It turns the hard constraints into soft constraints. Soft constraints that you (or the optimization) can violate.

Let’s look back to the first equation: we can formulate its Lagrangian as:

\begin{equation*}
L(x, \lambda, \nu) = f_0(x) + \sum_{i=1}^m \lambda_i f_i(x) + \sum_{i=1}^p \nu_i h_i(x)
\end{equation*}

Now it’s as if we are telling the God of optimization that, you can go ahead and violate our constraints as much as you can, and we are going to introduce two new variables $\lambda$ and $\nu$ to help you do that. These two variables are called dual variables. $\lambda$ is referred as inequality constraint dual variable, and not surprisingly $\nu$ is the equality constraint dual variable.

Dual Problem

After all these long and tedious definition of things (hopefully you aren’t too bored with them), we get to one last bit of information: dual problem. We can form a dual problem from the Lagrangian. What is a dual problem? Dual problem, simply put, is like you looking into the mirror. It’s you in the mirror, but the image is flipped. A deeper interpretation of dual problem is that it allows us to explore the effects of constraints on our original objective. Keep in mind that by forming Lagrangian, we allow the constraints to be soft, to add artibitrary large or small “penalty” to the original objective, but how would these constraints add penalization to the objective? dual variables provide the answer.

We first must get rid of the effect of $x$. Since we want to minimize the new objective (the Lagrangian), we can use the operator inf, it can give us the minimum value of the function over a chosen variable, and in here we choose $x$:

\begin{equation*}
g(\lambda, \nu) = \inf_{x \in D} L(x, \lambda, \nu) = \inf_{x \in D} \Big(f_0(x) + \sum_{i=1}^m \lambda_i f_i(x) + \sum_{i=1}^p \nu_i h_i(x) \Big)
\end{equation*}

By doing this, we reduce the first term $f_0(x)$, $f_i(x)$, and $h_i(x)$ to constants, and the infimum over a family of affine functions is a concave function (go through chapter 2 of the convex book for more details on this). Thus our dual objective is concave. Then we can quite awesomely form the following dual problem:

\begin{equation*}
\begin{aligned}
& \text{maximize}
& & g(\lambda, \nu) \\
& \text{subject to}
& & \lambda \succeq 0
\end{aligned}
\end{equation*}

A nicer understanding of this is that: by fixing $x$ (we shall now call it primal variable out of no reason…it’s actually to draw contrast to dual variables), we are driving the Lagrangian’s value down (as well as $f_0(x)$’s value), however, we must be violating constraints! So now, maybe it’s the time for the constraints to punish our objective, and we want the penalty to be as high as possible. Thus, by maximizing $g(\lambda, \nu)$, we can figure out a set of values for $\lambda$ and $\nu$ that will make the penalty as large as possible for the Lagrangian objective, hence the maximizing over $\lambda$ and $\nu$ on function $g$.

Geometric Analysis of Lagrangian

Now we can start our geometric interpretation of the Lagrangian. The problem rises when we alter the primal problem to form the lagrangian. We repeatedly ask ourself this question: does the newly formed Lagrangian become different from the primal problem in any way? If so, how can we avoid such primal-dual discordance by adding constraints? (Spoiler alert, these constraints become the famous KKT conditions).

We begin with the simplest example (as many explaination would start). You can easily extend them (mathematically, and mentally) to more complex problems. We first assume the simplest convex optimization problem, and we start with one equality constraint.

\begin{equation*}
\begin{aligned}
& \text{minimize}
& & f(x) \\
& \text{subject to}
& & h(x) = 0
\end{aligned}
\end{equation*}

We form the following Lagrangian: $L(x, \nu) = f(x) + \nu h(x)$. Notice the condition of optimality for the original problem is when $\nabla f(x) = 0$, and the optimality condition for the new Lagrangian is seemingly changed: $\nabla f(x) + \nu \nabla h(x) = 0$. We ask the question: will the change in optimality condition affect us from reaching the same goal (objective) as the original objective?

From the Lagrangian optimality condition, we can easily get $\nabla f(x) = -\nu \nabla h(x)$. Assume our final udpate to $x$ is $\Delta x$, meaning we move from our curent position $x$ to $x + \Delta x$, then in order to stay on the plane, $\Delta x$ must be orthogonal to the gradient vector of the plane ($h(x) = 0$), thus we get $\Delta x^T \nabla h(x) = 0$, and $\Delta x^T \nabla f(x) < 0$. Keep doing this, eventually, at the optimality point, we naturally have two gradient vectors being parallel (or anti-parallel, meaning pointing at different directions) becoming the end result when optimal solution is found. Notice we don’t require $\nabla f(x)$ and $\nabla h(x)$ to point at different directions, meaning $\nu$ doesn’t need to be positive or negative. Per our analysis, Lagrangian optimal condition captured what we would have wanted in the primal problem very nicely, we don’t need to add other constraints/conditions.

We get slightly more interesting result from examining the inequality constraint case. We can see the rise of KKT conditions and complementary slackness (two concepts that are often hard to grasp as first-timers) from the seemingly discordance between the original and Lagrangian optimality conditions.

\begin{equation*}
\begin{aligned}
& \text{minimize}
& & f(x) \\
& \text{subject to}
& & g(x) \leq 0
\end{aligned}
\end{equation*}

We formulate the simplest convex problem with one inequality constraint. The Lagrangian is $L(x, \lambda) = f(x) + \lambda g(x)$. The optimality condition for Lagrangian is $\nabla f(x) + \lambda \nabla g(x) = 0$. Now this case is slightly different from the equality case. We should first take a look at an illustration:

The shaded area including the surface when $g(x) = 0$ are the regions $x$ is allowed to be, satisfying our constraint. Like before, we can slightly reformulate the optimality condition to be: $\nabla f(x) = - \lambda \nabla g(x)$. We consider two cases. First, when $x$ is on the surface of $g(x) = 0$, meaning $x$ is on the boundary condition. When this is the case, it can only happen if $\nabla f(x)$ and $\nabla g(x)$ are pointing at opposite directions (notice gradient vector always points at the direction that increases function values), our desired direction is $-\nabla f(x)$. If not, meaning $\nabla f(x)$ and $\nabla g(x)$ are pointing at the same direction, all we need to do is to go to direction $-\nabla f(x)$, and we will decrease primal objective value as well as decreasing the value of $g(x)$, staying within $g(x) \leq 0$. To conclude, in order to have inequality constraint working, we must specify $\lambda > 0$.

Now we consider a second case, when $x$ is at $x^*$ position (below constraint surface). If the inequality constraint is satisfied, it should be effectively removed from the lagrangian equation, so that we are only optimizing under the primal optimal condition: $\nabla f(x) = 0$, thus $\lambda = 0$.

This is normally called “complementary slackness”, if $g(x) = 0$, then $\lambda > 0$, like our first case, and if $g(x) < 0$, then $\lambda = 0$, like our second case.

Now we are free to use what we have analyzed before, the conditions that will keep primal objective and Lagrangian the same:

1. Primal constraints: $f_i(x) \leq 0, i=1,…,m$, $h_i(x) = 0, i=1,…,p$
2. Dual constraints: $\lambda \succeq 0$
3. Complementary Slackness: $\lambda_i f_i(x) = 0, i =1,…,m $
4. Gradient of Lagnragian w.r.t $x$ vanishes: $\nabla f_0(x) + \sum_{i=1}^m \lambda_i \nabla f_i(x) + \sum_{i=1}^p \nu_i \nabla h_i(x) = 0$

These are also known as the KKT conditions. As we have pointed out before, if these conditions are satisfied, then Lagrangian and the primal problem are equivalent, meaning the primal problem has strong duality (refer to the book for more details).

Another fun fact is that, dual variables $\lambda$ and $\nu$, if KKT conditions are satisfied, (or if the primal problem has strong duality), then the dual variables are Lagrange multipliers.

Sensitivity Analysis

At last, we briefly mention sensitivity analysis. This portion was presented in a very confusing manner from the book (including powerpoint slides). If we take derivative of the Lagrangian with respect to the constraint:

\begin{equation*}
\lambda_i = - \frac{\partial g(\lambda^\star, \nu^\star)}{\partial f_i}
\text{,     }
\nu_i = - \frac{\partial g(\lambda^\star, \nu^\star)}{\partial h_i}
\end{equation*}

Notice that we are taking derivative of $g(\lambda^\star, \nu^\star)$. So the dual variable is the partial derivative of dual objective with respect to the constraints. Based on what we’ve learned from the last section, it’s not too difficult to reason that if $\lambda_i = 0$, the $i$ th inequality constraint is not constraining the objective. The larger $\lambda$ or $\nu$ is, the more constraining such constraint will be, a small change to tighten the constraint will cause a negative effect on the overall objective.

Dual variable $\lambda$ and $\nu$ are also known as shadow prices in economics.

Conclusion

This is the geometric intuition that slowly builds up the complementary slackness and the KKT condition. On first sight, KKT condition could be very daunting or confusing, same with complementary slackness. But a nice and solid grasp on directional gradient vectors will help provide a much cleaner understanding.

Trick 1: Collection1 == Collection2

The == in Numpy, when applied to two collections mean element-wise comparison, and the returned result is an array. This trick can be neatly combined with other operations that take in array result.

>> X_train, y_train, X_test, y_test = load_CIFAR10(cifar10_dir)
# Training labels (y_train) shape:  (50000,)
# to pick out correct y examples for one class (assuming that class is indexed at 5)
>> print(y_train == 5)
# array([False, False, False, ..., False, False, False], dtype=bool)
>> idxs = np.flatnonzero(y_train == 5)

np.flatnonzero() takes an array as input. Boolean array in Python has the property of True mapping to 1 and False mapping to 0, so flatnonzero() can easily pick out the index of those examples that are of class 5.

Trick 2: Vectorization/Loop-elimination with np.tile()

Sometimes when you want to add two matrices together, and have to manually align/broadcast them, you can use np.tile(). This copy the current matrix and broadcast into the shape you want. Keep the lowest dimension as 1 so your matrix/vector stay the same.

Any for-loop can be written in np.tile() and enjoy the advantage of vectorized operation:

# for a matrix shape of (10,5), assume this is our training example
>> train = np.random.randn(10,5)
# for another matrix shape of (5,5), assume this is our training example
# 5 examples, each example has 5 data points
>> test = np.random.randn(5,5)
# We want to form a matrix of (5, 10), each data point on every row is the
# difference between sum of test-data[j] and train-data[i].
>> train = np.sum(train, axis=1)
>> test = np.sum(test, axis=1)

# duplicate train data of shape (10,) 5 times
>> expanded_train = np.tile(train, [test.shape[0],1])
# (5,10)
>> expanded_test = np.tile(test, [train.shape[0],1])
# (10, 5)
>> expanded_test = expanded_test.T
# (5,10) -> transpose it so we can do subtraction
>> result = expanded_train - expanded_test
# (5,10)

np.tile() can also be used to expand arrays and allow element-wise multiplication/division/addision, however such functionality can be easily (or better) replaced by np.expand_dims() because expand_dims() will automatically figure out which dimension to expand on, and does not require a final transpose (which is often the case for np.tile).

Trick 3: Using Array for Pair-wise Slicing

Numpy array’s slicing often offers many pleasant surprises. This suprise comes from the context of SVM’s max() hinge-loss vectorization. SVM’s multi-class loss function requires wrong class scores to subtract correct class scores. When you dot product weight matrix and training matrix, you get a matrix shape of (num_train, num_classes). However, how do you get the score of correct classes out without looping (given y_labels of shape (num_train,))? At this situation, pair-wise selection could be helpful:

>> X = np.random.randn(500, 3073)
>> num_train = X.shape[0]
>> y_label = np.random.randint(10, (500,))
>> W = np.random.randn(3073, 10)
>> scores = X.dot(W)
>> correct_scores = scores[range(num_train),y_label]

If a numpy array’s [] operation has two arguments which both are arrays, then Numpy will run this pair-wise operations, like the example below:

>> a = np.asarray([1,2,3])
>> b = np.asarray([1,2,3])
>> correct_scores = scores[a, b]

It will select the [1,1], [2,2], [3,3] value from scores matrix. Another way to reversely use this is to assign value. So if we want to clean out the correct_label’s score for the loss function, we can:

>> scores[range(X.shape[0]), y_label] = 0

With this, all the correct label’s score are set to be 0, then when we sum them up, the correct_label will not affect the overall cost at all.

Trick 4: Smart use of ‘:’ to extract the right shape

Sometimes you encounter a 3-dim array that is of shape (N, T, D), while your function requires a shape of (N, D). At a time like this, reshape() will do more harm than good, so you are left with one simple solution:

for t in xrange(T):
  x[:, t, :] = # ...

You can use it to extract values or assign values!

Trick 5: Use Array as Slicing index

In previous posts, we already explored how Numpy array takes slicing of pairs (such as x[range(x.shape[0]), y]), however, Numpy can also take another array as slicing. Assume x is an index array of shape (N, T), each element index
of x is in the range 0 <= idx < V, and we want to convert such index array into array with real weights, from a weight matrix w of shape (V, D), we can simply do:

N, T, V, D = 2, 4, 5, 3

x = np.asarray([[0, 3, 1, 2], [2, 1, 0, 3]])  # (N, T)
W = np.linspace(0, 1, num=V*D).reshape(V, D)  # (V, D)

# this is the only required line
out = W[x]  # (N, T, D)

Numpy uses the underlying value of x as index to extract values from W.

Trick 6: Unfunc at

Numpy has a special group of “functions” called unfunc. Turns out numpy.add, numpy.subtract and so on belong to this special group of functions.

Numpy has predefined methods for those functions such as at! This function performs the operation (let it be add or subtract) on a specific location (or locations).

unfunc.at() takes 3 arguments, first one is the matrix you intend to modify, second argument is the indices of the matrix you want to modify, third argument is the value you want to change (depending on what the unfunc function is).

A simple example is:

>>> a = np.array([1, 2, 3, 4])
>>> np.add.at(a, [0, 1, 2, 2], 1)
>>> print(a)
array([2, 3, 5, 4])

However, this example is too simple and almost useless. This kind of useless examples permeate the whole Numpy documentation. Let’s look at a more advanced example: use this trick to update a word embedding matrix!

# dout: Upstream gradients of shape (N, T, D)
# x: from example in Trick 5 (N, T), indices
# V: the length of total vocabulary
# D: weight matrix dimension
# task here is to build a dW to modify
# the original weight matrix

dW = np.zeros((V, D), dtype=dout.dtype)
np.add.at(dW, x, dout)

Notice that Numpy converts x, a matrix into individual indicies, and use it to assign values from dout to dW. np.add.at() flattened dout so the dimension becomes (N*T, D). It will be checked against x’s dimension (N, T), and see if the product of x’s dimension and N*T will match. Only when this happens, you will assign the same amount of D arrays as you are instrcuted in x, to dW, which also happens to take D arrays.

for a tensor of shape (4,3,2,4) is a 4D tensor. Think of a 4D tensor as a tree, root connects with 4 leaves, and each leaf has 3 children, and each such children has 2 children, and each such children has 4 concrete elements: yes, the last number of the shape indicates the length of a list/array, such as array([1., 2., 3., 4.]).

So basically, you have 4 * 3 * 2 = 24 copies of array([1., 2., 3., 4.]). The axis in numpy, more intuitively, I think should be called “layer” (or level), corresponding to the tree analogy.

The use of axis is throughout Numpy. It’s an important concept because you can avoid using loops, thus resulting in better optimization (or vectorized operations). For example, with an array of shape (3, 20, 3), a thought would be, we have 3 layers of (20, 3), then we have 20 layers of array length of (3). So to sum up axis 1 (which is the 20 layer), we write np.sum(a, axis=1), and it will sum up lower dimensions in vector operation (i.e., adding all the 3-dim arrays together), and give us an output shape of (3,3).

This type of representation does make many operations every simple. No one wants to use 3 layers of for-loop to operate on the base layer. Even though both Numpy and Theano have broadcast, operating on two arrays with different shapes would be difficult and sometimes can’t act in the same way that we hope it would.

So various ways to effectively change the shape of arrays were developed. Before any further discussion, I want to point out that Theano’s dimshuffle is functionally more convenient than Numpy, and transpose + reshape = dimshuffle.

Let’s look at reshape first:

numpy.reshape(a, newshape, order='C')

a is the array, and newshape can be an int or a tuple like (3,2,5). When you are reshaping, the total number of elements can’t be altered, as explained above. If you are too lazy to calculate the what the remaining of this tuple should look like, you can just put -1, and Numpy will calculate for you.

>> rng = np.random.RandomState(234)
>> a = rng.randn(2,3,10)
[[[-0.19527893 -0.55857688  1.81404088 -0.96837131  0.24091401 -0.73500459
    0.2109724  -0.02775907 -0.3963966   0.59031683]
  [-0.50886026  0.4946645  -1.40719802 -1.11904803 -0.57989281  1.32724684
   -0.59643449 -1.23477173 -0.29232262 -0.1781048 ]
  [-0.77729182  0.12059733  0.04879302  0.0432387   0.36680183 -0.49865206
    1.98657646  1.49187182 -0.72115552 -0.13696597]]

 [[ 0.45780119 -0.21050629 -0.51760964 -0.10392791  1.3585057   2.47169207
   -0.2112615  -1.1344434  -1.094633    0.37917997]
  [-1.27183452 -0.45772402  0.28187467 -1.5637492  -0.9033226   0.32506839
    0.17751023 -0.91124188 -0.52497156  1.80199851]
  [ 1.07582354  1.34339255 -1.24913382  0.36885077 -0.42382145 -0.39417196
    1.28506918  1.40770487 -0.66003583 -0.93787805]]]
>> np.reshape(a, (3,5,-1)).shape
(3, 5, 4)

Reshape can also be used (or it’s commonly used) to flatten some dimension but not all of them. For certain algorithms (like simple KNN), layered informations aren’t needed, so we can flat out those information:

>> print(X_train.shape)
(5000, 32, 32, 3)
# we have 5000 examples, pixel 32X32, RGB color 3 channels
# but for KNN, we don't need this layered information, so we just flatten out the later 3 dimensions:
>> X_train = np.reshape(X_train, (X_train.shape[0], -1))
>> print(X_train.shape)
(5000, 3072)