Training models interactively in Pytorch

Question

I need to train two models in parallel. Each model has a different activation function with trainable parameters. I want to train model one and model two in the way that the parameters of the activation function from model one (e.g., alpha1) is separated from the parameters in model two (e.g., alpha2) by a gap of 2; i.e., |alpha_1 - alpha_2| > 2. I wonder how I could include it into the loss function for training.

Answer 1

Example module definition

I will use torch.nn.PReLU as parametric activation you talk about. get_weight created for convenience.

import torch


class Module(torch.nn.Module):
    def __init__(self, in_features, out_features):
        super().__init__()
        self.input = torch.nn.Linear(in_features, 2 * in_features)
        self.activation = torch.nn.PReLU()
        self.output = torch.nn.Linear(2 * in_features, out_features)

    def get_weight(self):
        return self.activation.weight

    def forward(self, inputs):
        return self.output(self.activation(self.inputs(inputs)))

Modules and setup

Here I'm using one optimizer to optimize parameters of both modules you talk about. criterion can be mean squared error, cross entropy or any other thing you need.

module1 = Module(20, 1)
module2 = Module(20, 1)

optimizer = torch.optim.Adam(
    itertools.chain(module1.parameters(), module2.parameters())
)
critertion = ...

Training

Here is a single step, you should pack it in a for-loop over your data as is usually done, hopefully it's enough for you to get the idea:

inputs = ...
targets = ...

output1 = module1(inputs)
output2 = module2(inputs)

loss1 = criterion(output1, targets)
loss2 = criterion(output2, targets)

total_loss = loss1 + loss2
total_loss += torch.nn.functional.relu(
    2 - torch.abs(module1.get_weight() - module2.get_weight()).sum()
)
total_loss.backward()

optimizer.step()

This line is what you are after in this case:

total_loss += torch.nn.functional.relu(
    2 - torch.abs(module1.get_weight() - module2.get_weight()).sum()
)

relu is used so the network won't reap infinite benefit solely from creating divergent weights. If there wasn't one, loss would become negative the greater the difference between weights would be. In this case the bigger the difference the better, but it makes no difference after the gap is greater or equal to 2.

You may have to increase 2 to 2.1 or something if you have to pass the threshold of 2 as the incentive to optimize the value when it's close to 2.0 would be small.

Edit

Without explicitly given threshold it might be hard, but maybe something like this would work:

total_loss = (
    (torch.abs(module1) + torch.abs(module2)).sum()
    + (1 / torch.abs(module1) + 1 / torch.abs(module2)).sum()
    - torch.abs(module1 - module2).sum()
)

It's kinda hackish for the network, but might be worth a try (if you apply additional L2 regularization).

In essence, this loss will have optimum at -inf, +inf pairs of weights in the corresponding positions and never will be smaller than zero.

For those weights

weights_a = torch.tensor([-1000.0, 1000, -1000, 1000, -1000])
weights_b = torch.tensor([1000.0, -1000, 1000, -1000, 1000])

Loss for each part will be:

(torch.abs(module1) + torch.abs(module2)).sum() # 10000
(1 / torch.abs(module1) + 1 / torch.abs(module2)).sum() # 0.0100
torch.abs(module1 - module2).sum() # 10000

In this case network can reap easy benefits just by making the weights greater with opposite signs in both modules and disregard what you want to optimize (large L2 on weights of both modules might help and I think optimum value would be 1/-1 in case L2's alpha is equal to 1) and I suspect the network might be highly unstable.

With this loss function if the network gets a sign of large weight wrong it will be heavily penalized.

In this case you would be left with L2 alpha parameter to tune to make it work, which is not that strict, but still requires a hyperparameter choice.