This is to get the gradient to "skip" the quantization part. The trick implements the red arrow in the [original paper](https://arxiv.org/abs/1711.00937)'s diagram: [![VQ-VAE architecture, with red arrow signifying gradient skipping from quantized code vector to encoder output][1]][1] ### Simplified example: Rounding Let's simplify this a bit and imagine we want to use rounding in our architecture: ```python 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: [![enter image description here][2]][2] We can look at the output: ```python r # tensor(6., grad_fn=<SumBackward0>) ``` All is looking good. However, when we try to compute the gradient, we get a tensor of zeroes: ```python r.backward() x.grad # tensor([0., 0.]) ``` This makes sense: the rounding function has derivative zero almost everywhere: [![enter image description here][3]][3] 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`](https://pytorch.org/docs/stable/generated/torch.Tensor.detach.html), a function that tells PyTorch to detach a vector from the computational graph. ```python 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$ ```python r # tensor(6., grad_fn=<SumBackward0>) ``` But we now also get reasonable gradients ```python r.backward() x.grad # tensor([2., 2.]) ``` Thanks to the fact that the gradient now "skips" the rounding part: [![enter image description here][4]][4] ### Back to VQ-VAE In VQ-VAE, we are replacing the output of the encoder (the `input` to the [vector quantization](http://www.mqasem.net/vectorquantization/vq.html) 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. [![enter image description here][5]][5] 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 ```python quantized = inputs + (quantized - inputs).detach() ``` does in the notebook. [1]: https://i.sstatic.net/HmVaG.png [2]: https://i.sstatic.net/WscOq.png [3]: https://i.sstatic.net/xyWXH.png [4]: https://i.sstatic.net/vGU4O.png [5]: https://i.sstatic.net/vcN7P.png