How to add two new dimensions to the input feature before the last two layers in MLP

Initially we had an MLP of multiple layers. We had an input embedding of 200 dimension. We now want to add two more dimensions to the original embedding to encode two important features. But as the original dimension is high, we fear that MLP would neglect the two new dimensions, which are quite important.

Thus, we want to add(concat) the two new dimensions before the last two layers of the MLP. I am still a new learner of ML and PyTorch, and I have searched a lot online but fail to come up with the ways to do this.

May I ask how can we achieve this using PyTorch? Thank you so much!

You could simply create two input heads. One for the embedding, which goes through its own neural network and then one for the two features. The output of both networks are then simply concatinated and passed into a final layer. Since for the one input head there is only two features (probably a vector of size two, right?)

You can combine two neural network modules simply like this:

# create a seperate network for your embedding input
class EmbeddingModel(nn.Module):
    def __init__(self):
        super(EmbeddingModel, self).__init__()
        
        self.layer1 = nn.Linear(...)
        . . .
        self.layerN = nn.Linear(...)
    
    def forward(self, x):
        x = F.activation(self.layer1(x))
        . . . 
        x = F.activation(self.layerN(x))
        return x

# create a one layer network for your "two important features"
# use the same activation function as the last layer of the "EmbeddingModel"
class FeaturesModel(nn.Module):
    def __init__(self):
        super(FeaturesModel, self).__init__()
        
        self.layer1 = nn.Linear(...)
    
    def forward(self, x):
        x = F.activation(self.layer1(x))
        return x


# finally create your main-model which combines both
class MainModel(nn.Module):
    def __init__(self):
        super(MainModel, self).__init__()
        
        self.embeddingModel = EmbeddingModel()
        self.featuresModel = FeaturesModel()

        # the input-dim to this layer has to be the output-dim of the embeddingModel + the output-dim of the featureModel
        self.outputLayer = nn.Linear(...)
    
    def forward(self, x_embeggings, x_features):
        x_embeggings = self.embeddingModel(x_embeggings)
        x_features = self.featuresModel(x_features)

        x = torch.cat((x_embeddings, x_features), -1)
        x = F.activation(self.outputLayer(x))

        return x