How To Zero Grad for selective inputs (Compressive Transformer Memory)

a2c: <p>Hi, I’m implementing the memory mechanism from the Compressive Transformer paper (<a href="https://arxiv.org/abs/1911.05507" rel="nofollow noopener">https://arxiv.org/abs/1911.05507</a>). There are two loss functions: the main cross-entropy of the language model and a reconstruction loss (MSE) for the difference between the attention on the real hidden states versus the compressed hidden states (see page 4). The latter is used to optimize how the compressed memory is constructed: it enforces a sort of scale-invariance between attention on original uncompressed memory and the smaller, compressed memory.</p>
<p>I’m constructing the reconstruction loss as follows:</p>
<pre><code class="lang-auto">compressor_attn, _ = checkpoint(self.attn, q=h, k=original_memory, v=original_memory)
compressed_attn, _ = checkpoint(self.attn, q=h, k=compressed_memory, v=compressed_memory)
</code></pre>
<p>Where <code>self.attn</code> is an nn.Module containing all the attention logic + parameters. And where the <code>compressed_memory</code> is created using an nn.Module that downsamples the size of the the original context memory -&gt; e.g. from length 1000 to length 100, see below for an example:</p>
<pre><code class="lang-auto">class MemoryCompressor(nn.Module):
    
    def __init__(self, input_dim, latent_dim):
        super(MemoryCompressor, self).__init__()
        self.input_dim = input_dim
        self.latent_dim = latent_dim
        self.encoder = nn.Linear(input_dim, latent_dim)
        self.apply(self.init_weights)
        
    def init_weights(self, module):
        if isinstance(module, (nn.Linear)):
            module.weight.data.normal_(mean=0.0, std=0.1 / np.sqrt(self.input_dim))
            
    def forward(self, input):
        return self.encoder(input).transpose(0, 2)

# Example of memory being compressed:
original_memory= ...
compressed_memory = MemoryCompressor(1000, 100)(original_memory)
</code></pre>
<p>Then the reconstruction loss is (at the end):</p>
<p><code>recons_loss = F.mse_loss(compressed_attn, compressor_attn)</code></p>
<p><strong>My problem</strong> is that right now the backwards pass goes from <code>recons_loss</code> -&gt; <code>compressor_attn</code> -&gt; the parameters of <code>self.attn</code>. The desired behaviour I want is that <em>only</em> the parameters from <code>MemoryCompressor</code> that are used to build <code>compressed_memory</code> have gradients that affect <code>recons_loss</code>; and vica versa, the actual attention parameters only have gradients that affect the main crossentropy loss of the language model.</p>
<p><strong>Why?</strong> Because I want the attention part of the model to focus purely on getting good prediction; I don’t want it distracted by trying to prioritize reconstruction loss. So the only bit I want focused on the reconstruction loss is the <code>MemoryCompressor</code>, so I need to selectively construct the gradients somehow so only the gradients from <code>MemoryCompressor</code> flow through to <code>recons_loss</code> (even if we use attention parameters as an intermediate step).</p>
<p><strong>Current State</strong>: I’ve figured out that I need to somehow zero-out the gradients of the bits I don’t want, but I’m unsure how to do it with this particular control flow. Am going to try and figure this out myself today - think it’s just my inexperience! - but any pointers would be appreciated (happy to clarify if anything is unclear).</p>

tom: <p>I think the key here is the “stop_gradient” pseudocode in Algorithm 2. This is just <code>h = hidden.detach()</code> and <code>old_mem = old_mem.detach()</code> or somesuch.</p>
<p>Best regards</p>
<p>Thomas</p>