In VQ-VAE code what does this line of code signify?

Question

The VQ-VAE implimentation:https://colab.research.google.com/github/zalandoresearch/pytorch-vq-vae/blob/master/vq-vae.ipynb

quantized = inputs + (quantized - inputs).detach()

Why are we subtracting and adding input to quantized result?

Answer 1

This is to get the gradient to "skip" the quantization part.

The trick implements the red arrow in the original paper's diagram:

Simplified example: Rounding

Let's simplify this a bit and imagine we want to use rounding in our architecture:

import torch
x = torch.tensor([1.1, 2.1], requires_grad=True)
y = 2*x
z = torch.round(y)
r = z.sum()

The graph (torchviz.make_dot) looks like this:

We can look at the output:

r 
# tensor(6., grad_fn=<SumBackward0>)

All is looking good. However, when we try to compute the gradient, we get a tensor of zeroes:

r.backward()
x.grad
# tensor([0., 0.])

This makes sense: the rounding function has derivative zero almost everywhere:

However, it also means we cannot train our network.

To circumvent that, we could simply tell the gradient to skip the rounding network. To do that, we use detach, a function that tells PyTorch to detach a vector from the computational graph.

x = torch.tensor([1.1, 2.1], requires_grad=True)
y = 2*x
z = torch.round(y)
z = y + (z - y).detach() # Detach everything between z and y, including z
r = z.sum()

We still get the same answer for $r$

r 
# tensor(6., grad_fn=<SumBackward0>)

But we now also get reasonable gradients

r.backward()
x.grad
# tensor([2., 2.])

Thanks to the fact that the gradient now "skips" the rounding part:

Back to VQ-VAE

In VQ-VAE, we are replacing the output of the encoder (the input to the vector quantization layer) with the code vector. In some sense, this similar to "rounding", only this time we're "rounding" the encoder's output to the nearest code vector.

We therefore run into the same problem. Here's what the authors say about it

Note that there is no real gradient defined for [the quantization], however we approximate the gradient similar to the straight-through estimator and just copy gradients from decoder input $z_q(x)$ to encoder output $z_e(x)$

which is precisely what

quantized = inputs + (quantized - inputs).detach()

does in the notebook.