KL divergence minimization

While reading Unsupervised Data Augmentation for Consistency Training, I came across an equation that describes the minimization of KL divergence.

$$\min_\theta \mathscr{J}_{UDA}(\theta) = \mathbb{E}_{x \in U}\mathbb{E}_{\hat{x}\in q(\hat{x} | x)}\left[D_{KL}\left(p_{\tilde{\theta}}(y | x) \left|\right| p_\theta(y | \hat{x} )\right)\right]$$

I'm trying to deconstruct this equation so that I can understand what's happening but having a lot of trouble doing so. Can anyone help me deconstruct and understand what this equation means?

The loss basically uses $$p_{\bar{\theta}}(y|x)$$ as the target/label for $$p_{\theta}(y|\hat{x})$$, where $$\hat{x}$$ is the augmented data. The goal is to let the output of augmented data stay close to the output of the original data, hence to enhance the consistency of the prediction function.

The following pseudocode should help understanding the loss

for x in unlabeled_data:  # expectation over the unlabeled data set
for x_hat in augment(x):  # expectation over the augmentation distribution
target = p_theta(x)  # don't optimize theta here
prediction = p_theta(x_hat)
loss_uda += kl_divergence(target, prediction)


When combining this loss with another supervised loss, it helps improve model generalizability by leveraging unlabeled data.

• Thank you for the answer; your pseudocode made it very clear. Just to clarify and better understand the mathematical notations, the equation has two expectations. I thought KL divergence in itself can be expressed as an expectation of the difference between the two probability distributions. So what are those two expectations outside the KL part? Commented Aug 6, 2019 at 20:20
• @BrianKo the fist expectation corresponds to the outer loop of the code, which is the expectation over the unlabeled data set, the second expectation corresponds to the inner loop, it is the expectation over the distribution of all possible augmentations of a data point Commented Aug 7, 2019 at 2:33