LSTM hidden_state PyTorch almost identicals which leads to a negative KLDivLoss

Viewed 175

I've been puzzling over this issue and I cannot figure out what am I doing wrong... I have trained an autoencoder (LSTM-LSTM) and I am now trying to use the encoded features for another task using the KLDivLoss. However, it turns out that the encoded features are almost always identical (see example below when setting the print precision to 10):

# last hidden state of the encoder
tensor([[[ 0.1086065620, -0.0446619801, -0.0530930459,  ...,
          -0.0573375113,  0.1083261892,  0.0037083717],
         [ 0.1086065620, -0.0446619801, -0.0530930459,  ...,
          -0.0573375151,  0.1083261892,  0.0037083712],
         [ 0.1086065620, -0.0446619801, -0.0530930459,  ...,
          -0.0573375188,  0.1083262041,  0.0037083719],
         ...,
         [ 0.1086065620, -0.0446619801, -0.0530930422,  ...,
          -0.0573375151,  0.1083262041,  0.0037083724],
         [ 0.1086065620, -0.0446619801, -0.0530930385,  ...,
          -0.0573375151,  0.1083262041,  0.0037083712],
         [ 0.1086065620, -0.0446619801, -0.0530930385,  ...,
          -0.0573375188,  0.1083261892,  0.0037083707]]],
       grad_fn=<StackBackward>)

Do you think there can be an explanation for such a behavior? I am working with high-dimensional time series and my aim is to implement an unsupervised clustering method by taking inspiration from this paper. If I feed the raw output of the encoder into the KLDivLoss I obtain a negative loss... However, if I rescale the output of the encoder (using sklearn.preprocessing.StandardScaler) I get the desired behavior: a positive loss value.

I already ensured that the first term in the KLDivLoss is log-probabilities and the second proabilities...

Code for the encoder (which contains an attention mechanism to identify the relevant driving series):

class Encoder(nn.Module):
    def __init__(self, config, input_size: int):
        super(Encoder, self).__init__()
        self.input_size = input_size
        self.hidden_size = config['hidden_size_encoder']
        self.seq_len = config['seq_len']

        self.lstm = nn.LSTM(
            input_size=self.input_size,
            hidden_size=self.hidden_size,
            num_layers=1
        )
        self.attn = nn.Linear(
            in_features=2 * self.hidden_size + self.seq_len,
            out_features=1
        )
        self.dropout = nn.Dropout(p=0.5)
        self.softmax = nn.Softmax(dim=1)

    def forward(self, input_data):
        h_t, c_t = (init_hidden(input_data, self.hidden_size), init_hidden(input_data, self.hidden_size))
        input_weighted = Variable(torch.zeros(input_data.size(0), self.seq_len, self.input_size))

        for t in range(self.seq_len):
            x = torch.cat((h_t.repeat(self.input_size, 1, 1).permute(1, 0, 2),
                           c_t.repeat(self.input_size, 1, 1).permute(1, 0, 2),
                           input_data.permute(0, 2, 1).to(device)), dim=2).to(
                device) 

            e_t = self.attn(x.view(-1, self.hidden_size * 2 + self.seq_len)) 
            a_t = self.dropout(self.softmax(e_t.view(-1, self.input_size)))

            weighted_input = torch.mul(a_t, input_data[:, t, :].to(device))
            self.lstm.flatten_parameters()
            _, (h_t, c_t) = self.lstm(weighted_input.unsqueeze(0), (h_t, c_t))

            input_weighted[:, t, :] = weighted_input

        return input_weighted[:, -1:, :], h_t, c_t

Code for clusterizing

class Clusterizer(nn.Module):
    def __init__(self, n_clusters, hidden_dim, encode, alpha=1.0):
        super(Clusterizer, self).__init__()
        self.n_clusters = n_clusters
        self.hidden_dim = hidden_dim
        self.encoder = encoder
        self.alpha = alpha
        self.centroids = None

    def init_centroids(self, encoded_x):
        '''initialize clusters center using KMeans'''
        kmeans = KMeans(n_clusters=self.n_clusters, random_state=0, n_init=10).fit(encoded_x)
        centroids = torch.tensor(kmeans.cluster_centers_, dtype=torch.float)
        self.centroids = nn.Parameter(centroids, requires_grad=True)

    def target_distribution(self, q_):
        weight = (q_ ** 2) / torch.sum(q_, 0)
        return (weight.t() / torch.sum(weight, 1)).t()

    def forward(self, encoded_x):
        num = ((1 + torch.norm(encoded_x.unsqueeze(1) - self.centroids, dim=2)) / self.alpha) ** (-(self.alpha + 1) / 2)
        den = torch.sum(num, dim=1, keepdim=True)
        return num / den

Final part of the code (where the problem occurs), here encoded_data is the stacked output of the encoder of shape (nb_observation, hidden_size).

criterion = nn.KLDivLoss(size_average=False)

...

clusterizer.init_centroids(encoded_data)
output = clusterizer(encoded_data.to(device))

target_distrib = clusterizer.target_distribution(output)
loss = criterion(output.log(), target_distrib)

I'm sorry that's a lot of code but I hope it helps to identify the source of the problem.

0 Answers
Related