Fast Quadratic Functions of Large Datasets

username: arnavs
cooked: <p>Hi,</p>
<p>I’m working on a problem which involves learning taking a vector <code>X = [1, c, x_1, ..., x_N]</code> (currently <code>N ~ 100</code>, but much more in the future) and learning a specific quadratic form, <code>Y'PY + b</code>, where <code>Y = [1, c, X_bar]</code> and <code>X_bar</code> is the average of the <code>x</code>'s.</p>
<p>The current way I’ve been doing this is by using a bilinear layer with a bias term, e.g. something like:</p>
<pre><code class="lang-auto">class vf_bilin_3(nn.Module):
    def __init__(self):
        super().__init__()
        self.Pb = nn.Bilinear(3, 3, 1, bias=True) # Y -&gt; Y'PY + b

    def forward(self, x):
        # forward pass
        M = x.shape[0]
        X_bar = torch.mean(x[:, 2:], dim=1, keepdims=True) 
        Y = torch.cat((x[:, 0:2].reshape([M, 2]), X_bar), dim=1) # [1, c, X_bar]
        output = self.Pb(Y, Y)
        return output
</code></pre>
<p>But I’m concerned that this won’t scale well when we have large numbers of observations per batch, or huge <code>N</code>, or on the GPU, etc.</p>
<p>I’d be very grateful if someone could point me in the right direction, or provide any optimization tips for this network.</p>
<p>– Arnav</p>

username: KFrank
cooked: <p>Hello Arnav!</p>
<aside class="quote no-group" data-username="arnavs" data-post="1" data-topic="93362" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/arnavs/48/27973_2.png" class="avatar"> arnavs:</div>
<blockquote>
<p>I’m concerned that this won’t scale well when we have large numbers of observations per batch, or huge <code>N</code></p>
</blockquote>
</aside>
<p>The short answer is to pre-compute <code>X_bar</code>.</p>
<aside class="quote no-group">
<blockquote>
<pre><code class="lang-auto">    def forward(self, x):
        # forward pass
        M = x.shape[0]
        X_bar = torch.mean(x[:, 2:], dim=1, keepdims=True) 
        Y = torch.cat((x[:, 0:2].reshape([M, 2]), X_bar), dim=1) # [1, c, X_bar]
        output = self.Pb(Y, Y)
        return output
</code></pre>
</blockquote>
</aside>
<p><code>output</code> depends on <code>X_bar</code> but not on the individual values of<br>
the components of <code>X</code>.  So if you call <code>.forward()</code> repeatedly (for<br>
example in your training loop) you will be needlessly computing<br>
<code>X_bar</code> multiple times.</p>
<p>Move the computation of <code>X_bar</code> out of <code>.forward()</code> – and thus<br>
out of the loop – to avoid this unnecessary computation.</p>
<p>(In a similar vein, but less important in terms of computation, I<br>
would also move the construction of <code>Y</code> out of <code>.forward()</code>.)</p>
<p>Best.</p>
<p>K. Frank</p>

username: arnavs
cooked: <p>Thanks so much <a class="mention" href="/u/kfrank">@KFrank</a>, that sounds like the right advice.</p>
<p>One question, though. I understand that the weight updating step should be different from the precomputing inputs step, but I thought <code>forward</code> was the model entry point. How can I separate the two in practice?</p>
<p>— Arnav</p>

username: KFrank
cooked: <p>Hello Arnav!</p>
<aside class="quote no-group" data-username="arnavs" data-post="3" data-topic="93362" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/arnavs/48/27973_2.png" class="avatar"> arnavs:</div>
<blockquote>
<p>How can I separate the two in practice?</p>
</blockquote>
</aside>
<p>The idea would be to preprocess your data before you feed it to your<br>
model.</p>
<p>Depending on the details of your use case, you could, for example,<br>
read your original data off of disk (independent of pytorch), construct<br>
your <code>Y</code> vectors, and then write this preprocessed data back to disk.<br>
When you run your pytorch model, just used the preprocessed data<br>
as your input data.</p>
<p>Note, you will only get (significant) savings if you are going to run<br>
(a lot) more than one <em>epoch</em> in your training loop.  (By “epoch” I<br>
mean that you run each sample in your dataset through your model<br>
once.)  If you are only going to run any given data sample through<br>
your model once (or just a few times), it doesn’t matter whether you<br>
perform the <code>X_bar</code> computation ahead of time or in during training.<br>
So in such a case it might be more convenient to leave the <code>X_bar</code><br>
computation in your <code>.forward()</code>.</p>
<p>Best.</p>
<p>K. Frank</p>

username: arnavs
cooked: <p>Hi <a class="mention" href="/u/kfrank">@KFrank</a>, thanks again for the quick and detailed reply. Clarity on terms like “epoch” is super helpful to neophytes to the ML world.</p>
<p>And that makes a lot of sense. The model’s the thing in the training loop (or broadcast), so it should only do things we’re comfortable repeating.</p>