How can I do memory efficient attention with square mask

Hyeonuk_Woo: <p>Hi, I wan to do attention calculation. For my task, from <em>priori</em> , I can block lots of query-key pair. (like Figure 1 and Figure 2) In Figure 1,2 , I paint gray for blocked query-key pair.</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/6/c/6c51ad95b7d29365f35f1af59774ae79371b658f.png" data-download-href="https://discuss.pytorch.org/uploads/default/6c51ad95b7d29365f35f1af59774ae79371b658f" title="figure1"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/6/c/6c51ad95b7d29365f35f1af59774ae79371b658f_2_690x286.png" alt="figure1" data-base62-sha1="fseyNtkYR8XBMXeuw8cxBDZPCN1" width="690" height="286" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/6/c/6c51ad95b7d29365f35f1af59774ae79371b658f_2_690x286.png, https://discuss.pytorch.org/uploads/default/original/3X/6/c/6c51ad95b7d29365f35f1af59774ae79371b658f.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/6/c/6c51ad95b7d29365f35f1af59774ae79371b658f.png 2x" data-dominant-color="DCDCDC"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">figure1</span><span class="informations">944×392 25.6 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p><strong>Figure 1. example of query-key pair tensor</strong></p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/e/e/ee1378dd83ec58a746412dcb6b629b3d1ef8f2c4.png" data-download-href="https://discuss.pytorch.org/uploads/default/ee1378dd83ec58a746412dcb6b629b3d1ef8f2c4" title="tmp1"><img src="https://discuss.pytorch.org/uploads/default/original/3X/e/e/ee1378dd83ec58a746412dcb6b629b3d1ef8f2c4.png" alt="tmp1" data-base62-sha1="xY7m7Ays9Pr8p5e6UShqtiCihrC" width="690" height="286" data-dominant-color="D4D4D4"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">tmp1</span><span class="informations">1475×613 10.9 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div><br>
<strong>Figure 2. example of query-key pair square mask</strong></p>
<p>I use <code>mask</code> for attention calculation as below.</p>
<pre><code class="lang-auto">square_mask= (-1*) square_mask
square_mask= inf*square_mask
attention_logit += square_mask 
attention_prob = nn.functional.softmax(attention_logit)
</code></pre>
<p>I think that , computing in this way, even only a fraction of query-key pair what I should be really interested in, the amount of memory required to perform the calculation is still the same as using whole query-key pair.</p>
<p>As shown below(Figure 3&amp;4), if we extract only the parts we need from the entire matrix and proceed with the calculation, I think that it can be an efficient calculation in terms of memory. How can we do this effectively?<br>
Is there a better way other than this one?</p>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/3/b/3b30ada5cefbbcecc24e3b4cb18a7d1b2b4d1050.png" alt="figure2" data-base62-sha1="8rCsJUIoRnJV1zncmjH4Q1N8UA8" width="400" height="390"></p>
<p><strong>Figure 3. extracted query-key tensor</strong></p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/b/4/b479200556635933621235439c8b7a5f42105182.png" data-download-href="https://discuss.pytorch.org/uploads/default/b479200556635933621235439c8b7a5f42105182" title="tmp4"><img src="https://discuss.pytorch.org/uploads/default/original/3X/b/4/b479200556635933621235439c8b7a5f42105182.png" alt="tmp4" data-base62-sha1="pKxs93ZH58ypcgQ82eZ7xkHSS1I" width="463" height="500" data-dominant-color="E0E0E0"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">tmp4</span><span class="informations">565×609 6.21 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div><br>
<strong>Figure 3. extracted query-key sqaure mask</strong></p>