Generate NN weights using another NN

Neoxion: <p>Hello,</p>
<p>I have the following setup. There is a small NN <code>NetA</code> that predicts values <code>y</code> from a set of inputs <code>x1</code>:</p>
<pre><code class="lang-auto">class NetA(nn.Module):
    def __init__(self, nf = 4):
        super(NetA, self).__init__()
        self.main = nn.Sequential(
            nn.Linear(2, nf, bias=False),
            nn.BatchNorm1d(nf),
            nn.Tanh(),

            nn.Linear(nf, nf, bias=False),
            nn.BatchNorm1d(nf),
            nn.Tanh(),

            nn.Linear(nf, nf, bias=False),
            nn.BatchNorm1d(nf),
            nn.Tanh(),

            nn.Linear(nf, 1, bias=False),
        )
    def forward(self, x1):
        x = self.main(x1).view(-1)
        return x
</code></pre>
<p>There is also a set of inputs <code>x2</code> and every <code>x2[i]</code> produces a different set of <code>y[i]</code> for the entire span of <code>x1</code>, used in <code>NetA</code>. Thus, I pre-train <code>NetA</code> for <code>x2[0]</code> to get some kind of base weights and transform NN into a functional form <code>fNetA</code>:</p>
<pre><code class="lang-auto">class fNetA(nn.Module):
    def __init__(self):
        super(fNetA, self).__init__()

    def forward(self, x1, bn_param_, param_, param6_, out_param_):
        x = F.tanh(F.batch_norm(F.linear(x1, weight = param_[:, :2], bias = None), running_mean=bn_param_[:,0], running_var=bn_param_[:,1], 
                    weight=bn_param_[:,2], bias=bn_param_[:,3]))
        x = F.tanh(F.batch_norm(F.linear(x, weight = param_[:, 2:], bias = None), running_mean=bn_param_[:,4], running_var=bn_param_[:,5], 
                    weight=bn_param_[:,6], bias=bn_param_[:,7]))
        x = F.tanh(F.batch_norm(F.linear(x, weight = param6_, bias = None), running_mean=bn_param_[:,8], running_var=bn_param_[:,9], 
                    weight=bn_param_[:,10], bias=bn_param_[:,11]))
        x = F.linear(x, weight = out_param_, bias = None)

        return x.view(-1)
</code></pre>
<p>I do this because I would like to generate <code>out_param_</code> with the help of a different NN <code>NetB</code>:</p>
<pre><code class="lang-auto">class NetB(nn.Module):
    def __init__(self, nf, gnf):
        super(NetB, self).__init__()
        self.nf = nf
        self.embedding = nn.Embedding(5, gnf)

        self.main = nn.Sequential(
            nn.Linear(gnf, gnf, bias=False),
            nn.BatchNorm1d(gnf),
            nn.Tanh(),

            nn.Linear(gnf, gnf, bias=False),
            nn.BatchNorm1d(gnf),
            nn.Tanh(),

            nn.Linear(gnf, nf, bias=True), # output layer weights out_param_
            nn.Tanh(), # limit delta because y values are very close to baseline
        )

    def forward(self, x2):
        x = self.embedding(x2) # x2 represents specific class
        x = self.main(x).mean(0).view(-1, self.nf)
        return x
</code></pre>
<p>The training loop for <code>NetB</code> is:</p>
<pre><code class="lang-auto">fNetA.train()
NetB.train()
for i, (x1, x2, y) in enumerate(train_dataloader):
    netb_optimizer.zero_grad()
    new_out_param_ = NetB(x2) + out_param_ # NetB output + baseline weights
    out = fNetA(x1, bn_param_, param_, param6_, new_out_param_) # bn_param_, param_, param6_ are just saved tensors
    g_loss = criterion(out, y)
    loss.backward()
    netb_optimizer.step()
</code></pre>
<p>I am aware that I calculate mean over the batch in <code>NetB</code>. I structured batches in a way, so only values for a specific <code>x2[i]</code> can be in one batch and they all have the same <code>new_out_param_</code>.<br>
This setup works if I want to tune the baseline weights of any specific <code>x2[i]</code>, meaning that <code>NetB</code> can learn the delta for <code>new_out_param_</code> for a single case, but when I start using all <code>x2[i]</code> in the train data-loader, <code>NetB</code> cannot converge at all. I tried gradient accumulation over multiple batches to ensure that weights are updated based on multiple <code>x2[i]</code>. I also tried to mix various <code>x2[i]</code> within one batch, produce batched weights by removing mean from <code>NetB</code> and iterate over them accumulating gradients to update the weights. Nothing really works. Removing <code>nnTanh()</code> in the output of <code>NetB</code> doesn’t change anything either.<br>
I noticed that gradients of <code>NetB</code> weights are very small (around 1e-9), but they are that small in a single case too, but NN can still learn. Technically, I can learn all required <code>new_out_param_</code> independently and then train the generator NN <code>NetB</code>, but I wanted to do it in a proper way, that is scalable. So, I have two questions:</p>
<ol>
<li>What could be the issue, which is preventing <code>NetB</code> learning in this case?</li>
<li>Is it possible to do it differently? For example, back-propagate loss into the beginning of <code>NetA</code> and use its pure value to train <code>NetB</code> (only using graph from <code>NetB</code>)?</li>
</ol>
<p>Thanks in advance</p>