GVAE prediction, concatenation of node features and condition variable

[Anneqw]

I'm trying to reconstruct node features with GVAE, where each graph has a specific condition corresponding to it. When predicting, combine random variables sampled from N(0, I) and given condition variables to generate specified node features. The loss decreases and converges during the training phase, but the results are not good when using the trained model for prediction. I think this maybe because of the inapproprite way of fusion of condition variable and node features, what I did in the model is expanding the condition variable to have the same dimension as node features and then concatenating the two variables. The dimension of node feature and condition variable are [200,2] and [num_condition,1] seperately. I transposed the condition variable and then expanded it to [200,num_condition]. The final dimension of concatenated variable is [200,2+num_condition]. The condition variable should be the condition variable of the entire graph. Is there any problem with the concatenation?

Any suggestion to would be appreciated.

Thanks a lot in advance ! :)

class model(nn.Module):
    def __init__(self,prediction):
        super(model, self).__init__()
        '''Encoder'''
        self.el = nn.Linear(args.x_dim, 128, bias=False)
        self.base_gat = GAT(args.num_layers, args.num_feature_per_layer)
        self.gat_mean = nn.Linear(args.hidden1_dim, args.hidden2_dim, bias=False)
        self.gat_logstddev = nn.Linear(args.hidden1_dim, args.hidden2_dim, bias=False)
        '''Decoder'''
        self.d_hidden1 = nn.Linear(args.decode_input_dim, args.d_hidden1_dim, bias=False)
        self.d_hidden2 = nn.Linear(args.d_hidden1_dim, args.d_hidden2_dim, bias=False)
        self.d_hidden3 = nn.Linear(args.d_hidden2_dim, args.d_hidden3_dim, bias=False)
        self.out = nn.Linear(args.d_hidden3_dim, args.output_dim, bias=False)

        self.prediction = prediction

    def encode(self, X, edge_index, c):
        new_feature= self.cat_X_c(X, c)
        input = (new_feature, edge_index)
        hidden = self.base_gat(input)
        hidden = F.leaky_relu(hidden[0])
        self.mean = self.gat_mean(hidden)
        self.logstd = self.gat_logstddev(hidden)
        gaussian_noise = torch.randn_like(self.logstd)
        sampled_z = gaussian_noise * torch.exp(self.logstd) + self.mean
        return sampled_z

    def decode(self, Z, c):
        input = self.cat_Z_c(Z, c)
        out1 = self.d_hidden1(input)
        out1 = F.leaky_relu(out1)
        out2 = self.d_hidden2(out1)
        out2 = F.leaky_relu(out2)
        out3=self.d_hidden3(out2)
        feature_pred = self.out(out3)
        return feature_pred

    def forward(self, data):
        c = (data.condition).to(torch.float32) / 100  
        X = (data.x).to(torch.float32) / 100  
        edge_index = (data.edge_index).to(torch.int64) 

        self.batch = X.shape[0] // 200

        c = c.view(-1, 1)
        Z = self.encode(X, edge_index, c)
        pred = self.decode(Z, c) 
        return pred, self.mean, self.logstd

    def cat_X_c(self, X, c):
        c = c.T
        condition = torch.empty([X.shape[0], args.c_dim]).to(device=X.device)
        for i in range(self.batch):
            c_temp = c[0, i * args.c_dim:(i + 1) * args.c_dim]
            c_temp = c_temp.repeat(200, 1)
            condition[i * 200:(i + 1) * 200, :] = c_temp
        return torch.cat((X, condition), dim=1)

    def cat_Z_c(self, Z, c):
        c = c.T
        condition = torch.empty([Z.shape[0], args.c_dim]).to(device=Z.device)
        for i in range(self.batch):
            c_temp = c[0, i * args.c_dim:(i + 1) * args.c_dim]
            c_temp = c_temp.repeat(200, 1)
            condition[i * 200: (i + 1) * 200, :] = c_temp
        return torch.cat((Z, condition), dim=1)

[rusty1s]

I think the code looks correct to me (expect for that you can probably replace cat_X_c and cat_Z_c by a single Tensor.repeat call which makes it easier to understand. If the loss decreases during training, but does not reflect on validation and test set, there is likely another issue going on or you are dealing with some serious OOO?

[Anneqw]

Thank you. I just checked my dataset, the data in the training and test sets are identically distributed. So there should be no problem with this. It doesn't seem to be the problem.
During prediction I randomly take a graph from the test dataset and generate a random vector Z whose data satisfies a normal distribution. The dimension of Z is [200,1] because there are 200 node features to predict in a graph. Is there any problem with this part?

test_dataset = MyDataset(root, filename, test=True)
test_loader = DataLoader(test_dataset, batch_size=1, shuffle=True)
for i, batch in enumerate(test_loader):
    condition = batch.condition
    X = batch.x
    break
condition=condition.T

model = model(prediction=True)
checkpoint = torch.load('model_params.pth')
model.load_state_dict(checkpoint)

prediction_number = 200

Z = torch.randn(prediction_number).view(-1, 1)
pred = model((Z, condition)) * 100

[rusty1s]

In general this should work. If it doesn't, it would mean that your variational autoencoder has not been able to properly learn to encode its latent space as a normal distribution.

[Anneqw]

It feels like a serious problem. Is there any way to tell if this is the case that you are talking about?

[rusty1s]

How big is your kl_loss? You can also generate z of training examples and compare them with your normal distribution.

[Anneqw]

Thank you for your help! I found that the proportion of KL divergence in the loss function is too large, making the distribution of z too concentrated. Finally, the problem is solved by adjusting the proportion of KL divergence in the loss function. ^_^