Understanding gumbel-softmax backpropagation in Wav2Vec papers - Artificial Intelligence Stack Exchange most recent 30 from ai.stackexchange.com 2022-01-27T23:34:43Z https://ai.stackexchange.com/feeds/question/32109 https://creativecommons.org/licenses/by-sa/4.0/rdf https://ai.stackexchange.com/q/32109 0 Understanding gumbel-softmax backpropagation in Wav2Vec papers Peter Franek https://ai.stackexchange.com/users/9092 2021-10-19T11:37:56Z 2021-10-24T21:11:59Z <p>I'm studying the series of Wav2Vec papers, in particular, the <a href="https://arxiv.org/pdf/1910.05453.pdf" rel="nofollow noreferrer">vq-wav2vec</a> and <a href="https://arxiv.org/pdf/2006.11477.pdf" rel="nofollow noreferrer">wav2vec 2.0</a>, and have a problem understanding some details about the quantization procedure.</p> <p>The broader context is this: they use raw audio and first convert it to &quot;features&quot; <span class="math-container">$z$</span> via a convolutional network. Then they project any feature <span class="math-container">$z$</span> to a &quot;quantized&quot; element <span class="math-container">$\hat{z}$</span> from a given finite codebook (or concatenation of finitely many finite codebooks). To find <span class="math-container">$\hat{z}$</span>, they compute scores <span class="math-container">$l_j$</span> for each codebook entry <span class="math-container">$v_j$</span>, convert these scores to Gumbel-Softmax probabilities <span class="math-container">$p_j$</span> (using a formula which is not deterministic, the formula involves random choices of some numbers from some distribution) and then use these probabilities <span class="math-container">$p_j$</span> to choose <span class="math-container">$\hat{z}$</span>. Further stages of the pre-training pipeline are trained to predict <span class="math-container">$\hat{z}'s$</span> by either predicting &quot;future&quot; from the &quot;past&quot;, or &quot;reconstructing masked segments&quot;.</p> <p>My question is this is about this sentence:</p> <blockquote> <p>During the forward pass, <span class="math-container">$i = \text{argmax}_j p_j$</span> and in the backward pass, the true gradient of the Gumbel-Softmax outputs is used.</p> </blockquote> <ul> <li>I have trouble seeing what exactly is happening in the loss function and back-propagation. Could someone please help me to break this down into details?</li> </ul> <hr /> <p>My mental attempts to make sense out of it (I'm using the notation <span class="math-container">$\hat{z}$</span> for quantized vectors, in the second paper they use <span class="math-container">$q$</span>)</p> <p>(1) I would say that during the forward pass, in the Gumbel-Softmax, random variables from the Gumbel-distribution <span class="math-container">$n_j$</span> are sampled every time (for every training example) to compute the Gumbel-softmax probabilities <span class="math-container">$p_j$</span>.</p> <p>(1a) In the back-propagation, these <span class="math-container">$n_j$</span>'s are kept constant, and <span class="math-container">$p_j$</span> is treated as a function of <span class="math-container">$l_j's$</span> only.</p> <p>(2) The loss function has 2 parts here, Contrastive loss and Diversity loss.</p> <p>(2a) Based on the description, I would say that in the contrastive loss, the &quot;sampled&quot; vectors <span class="math-container">$\hat{z}_j$</span> are used, and probabilities never appear (even not in back-propagation of this part of the loss).</p> <p>(2b) I would believe that in the gradient of the Diversity loss, which only uses probabilities <span class="math-container">$p_{g,v}$</span>, that here the gradient or the loss actually <strong>is</strong> used, as this is responsible for maximizing the entropy. This part of the gradient probably does not use the sampled values <span class="math-container">$\hat{z}_j$</span>.</p> <p>Is this approximately correct?</p> <p>If yes, then I still fail to understand what exactly is happening in the vq-wav2vec paper. The sentence</p> <blockquote> <p>During the forward pass, <span class="math-container">$i = \text{argmax}_j p_j$</span> and in the backward pass, the true gradient of the Gumbel-Softmax outputs is used.</p> </blockquote> <p>is there as well, but I cannot see any part of the loss function (in this paper) where the probabilities are explicitly used (such as the diversity loss).</p>