Composite Multi-fidelity neural network - doubts regarding implementation

username: philbrand
cooked: <p>Hi everyone!<br>
I am trying to replicate the examples shown in <a href="https://www-sciencedirect-com.ezproxy.biblio.polito.it/science/article/pii/S0021999119307260" rel="noopener nofollow ugc">this paper</a> but I’m having some doubts whether my code is correct.<br>
The basic idea is that we want to model some function, for which we have a few coordinates of high fidelity approximations and then some more low fidelity points.</p>
<p>To do so, the authors use a composite neural network structured as follows:</p>
<ul>
<li>A low fidelity neural network NN_L, that takes the input and outputs the low fidelity prediction. This is a non-linear neural network with Tanh as the activation function</li>
<li>A linear high fidelity neural network NN_H1, that takes as inputs the original input and the associated low fidelity prediction and outputs F_L. This should be a linear neural network with no activation function, so I’m using the torch.nn.Identity activation function</li>
<li>A non-linear high fidelity neural network NN_H2, that takes as inputs the original input and the associated low fidelity prediction and outputs F_NL. This neural network uses a Tanh activation function.<br>
The sum of F_L+F_NL is the high fidelity prediction.</li>
</ul>
<p>Here comes the problem: in the paper they do not use any hidden layer in NN_H1, which makes sense to me as it is a linear layer. Also they use L-BFGS, while I’m using Adam, but I don’t think that’s the problem here as I get similar results using L-BFGS too.<br>
If I use no hidden layers in NN_H1, I can never seem to get an accurate prediction, regardless of how long and which learning rate/weight decay I use, see the figure below:<br>
If I add 2 hidden layers with 20 neurons, I get a much better result:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/e/f/ef80bef7fad6938f3dd8ed4b9e9523550bf79d3a.png" data-download-href="https://discuss.pytorch.org/uploads/default/ef80bef7fad6938f3dd8ed4b9e9523550bf79d3a" title="On the left: prediction using no hidden layers. On the right: prediction using 2 hidden layers."><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/e/f/ef80bef7fad6938f3dd8ed4b9e9523550bf79d3a_2_690x228.png" alt="On the left: prediction using no hidden layers. On the right: prediction using 2 hidden layers." data-base62-sha1="yaJX33H3ohp66iFYeURbIB1JR7I" width="690" height="228" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/e/f/ef80bef7fad6938f3dd8ed4b9e9523550bf79d3a_2_690x228.png, https://discuss.pytorch.org/uploads/default/optimized/3X/e/f/ef80bef7fad6938f3dd8ed4b9e9523550bf79d3a_2_1035x342.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/e/f/ef80bef7fad6938f3dd8ed4b9e9523550bf79d3a.png 2x" data-dominant-color="F9F9F9"><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">On the left: prediction using no hidden layers. On the right: prediction using 2 hidden layers.</span><span class="informations">1103×365 35.1 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>The question is: am I creating the neural network correctly? Am I misunderstanding something? Maybe the problem lies elsewhere? I am creating the networks as follows:</p>
<pre><code class="lang-auto">import torch


class BasicBlock(torch.nn.Module):
    """
    Single block for a fully connected layer. Alternates a linear layer and an activation function.
    """
    def __init__(self, input_features: int, output_features: int, activation: type[torch.nn.Module]):
        super().__init__()
        self.linear_layer = torch.nn.Linear(input_features, output_features)
        self.activation = activation()

    def forward(self, x):
        x_a = self.linear_layer(x)
        out = self.activation(x_a)
        return out


class FCNN(torch.nn.Module):
    """
    Fully connected Neural Network template. Creates a Fully Connected NN with specified features and activation function
    """

    def __init__(self, in_features: int, out_features: int, mid_features: list[int], activation: type[torch.nn.Module]):
        super().__init__()
        layer_sizes = [in_features] + mid_features
        """
        Creates the inner layers by calling the BasicBlock (1 linear + 1 activation) for as many times as there are hidden layers (len(mid_features))
        """
        self.layers = torch.nn.Sequential(*[
            BasicBlock(layer_sizes[i], layer_sizes[i+1], activation)
            for i in range(len(mid_features))
        ])
        self.fc = torch.nn.Linear(layer_sizes[-1], out_features)

    def forward(self, x: torch.Tensor) -&gt; torch.Tensor:
        x_a = self.layers(x)
        out = self.fc(x_a)
        return out


class MFNN(torch.nn.Module):
    """
    A multi fidelity neural network class. This creates two subclasses: fc_low is the low fidelity part of the model
        fc_high is the high fidelity. fc_high takes two inputs (the x coordinate and the relative low-fidelity
        prediction).

    in_features: number of features given as input
    out_features: number of features computed by the COMPLETE model
    lf_layers: a list of the layer widths of the low fidelity layers
    hf1_layers: a list of the layer widths of the NNH1 layers
    hf2_layers: a list of the layer widths of the NNH2 layers
    """

    def __init__(self, in_features: int, out_features: int, lf_layers: list[int], hf1_layers: list[int],
                 hf2_layers: list[int], lf_activation: type[torch.nn.Module], hf_activation: type[torch.nn.Module]
                 ):
        super().__init__()

        # Creates the list of the layer sizes for the various fidelity
        lf_layer_sizes = [in_features] + lf_layers
        # We now create two different neural networks, for the NNlow and NNhigh
        self.fc_low = FCNN(lf_layer_sizes[0], out_features, lf_layers, lf_activation)
        self.fc_high1 = FCNN(2, out_features, hf1_layers, activation=torch.nn.Identity)
        self.fc_high2 = FCNN(2, out_features, hf2_layers, hf_activation)

    def forward(self, x: torch.Tensor):
        low_features = self.fc_low(x)
        in_h = torch.cat([x, low_features], dim=1)
        y1 = self.fc_high1(in_h)
        y2 = self.fc_high2(in_h)
        y = y1 + y2
        return y

low_fidelity_layers = [20, 20]
high_fidelity1_layers = [20, 20] # Adds 2 hidden layers
high_fidelity2_layers = [20, 20, 20, 20]
# Create the model with 2 hidden layers
mf_model = MFNN(1, 1, lf_layers=low_fidelity_layers,
                      hf1_layers=high_fidelity1_layers, hf2_layers=high_fidelity2_layers,
                      lf_activation=torch.nn.Tanh, hf_activation=torch.nn.Tanh)
# Create the model with no hidden layers
mf_model_nh = MFNN(1, 1, lf_layers=low_fidelity_layers,
                      hf1_layers=[], hf2_layers=high_fidelity2_layers,
                      lf_activation=torch.nn.Tanh, hf_activation=torch.nn.Tanh)
</code></pre>
<p>Thanks in advance!</p>