LSTM loss not improving

hH1sG0n3: <p><strong>Background</strong></p>
<p>I am using an LSTM to model sequential events for a binary classification problem. The dataset is</p>
<ul>
<li>highly imbalanced (<code>ratio=0.1</code>)</li>
<li>unique tokens are <code>&lt;100</code>
</li>
<li>tokens do not strongly link with natural language (<code>token = semi natural language encoding of events</code>)</li>
<li>high number of datapoints (<code>&gt;10^6</code>)</li>
</ul>
<p>To improve performance the an additional fully connected NN port for static features (postcode, height, violations etc.) is combined with the LSTM module.</p>
<p><strong>Model architecture</strong></p>
<p>Pasted below pseudocode for model.</p>
<pre><code class="lang-auto">
class my_network(nn.Module):
    def __init__(
        self,
        vocab_size,
        output_size,
        embedding_dim,
        latent_dim,
        n_lstm,
        bidirectional=False,
        recurrent_dropout=0.3,
        dropout=0.3,
        n_aux=2,
        n_fc=1,
        n_neurons=None,
    ):
 
        super(my_network, self).__init__()

        self.output_size = output_size
        self.n_lstm = n_lstm
        self.latent_dim = latent_dim
        bidirectional = bool(bidirectional)
        self.bidirectional = bidirectional
        dir = 1 if bidirectional == 0 else 2
        self.dir = dir
        self.n_aux = n_aux
        self.n_fc = n_fc
        if n_neurons is None:
            n_neurons = (latent_dim * dir) + n_aux
        self.n_neurons = n_neurons

        self.embedding = nn.Embedding(vocab_size + 1, embedding_dim)
        self.lstm = nn.LSTM(
            embedding_dim,
            latent_dim,
            n_lstm,
            dropout=0 if n_lstm == 1 else recurrent_dropout,
            bidirectional=bidirectional,
            batch_first=True,
            bias=True,
        )
        self.batchnorm = nn.BatchNorm1d(dir * latent_dim)
        self.dropout = nn.Dropout(dropout)

        # Instantiate merger forward neural network.
        self.fc_aux = nn.Linear(n_aux, n_aux)
        self.fc = nn.ModuleList()

        for i in range(self.n_fc):
            if i == 0:
                if n_fc == 1:
                    self.fc.append(nn.Linear((latent_dim * dir) + n_aux, output_size))
                else:
                    self.fc.append(nn.Linear((latent_dim * dir) + n_aux, n_neurons))
            elif i == n_fc - 1:
                self.fc.append(nn.Linear(n_neurons, output_size))
            else:
                self.fc.append(nn.Linear(n_neurons, n_neurons))

    def forward(self, seq_x, aux_x, hidden):
 
        batch_size = seq_x.size(0)

        embedding = self.embedding(seq_x)
        lstm, (h_out, cell_out) = self.lstm(embedding, hidden)
        # Use hidden state for merger.
        lstm_out = h_out
        lstm_out = lstm_out[-1]

        lstm_out = self.dropout(lstm_out)
        lstm_out = self.batchnorm(lstm_out)
        ##-- END sequential modelling branch --##

        aux_out = self.fc_aux(aux_x)

        X = torch.cat([lstm_out, aux_out], dim=1)
        X = F.relu(X)
        # out = self.fc(X)
        for i in range(self.n_fc):
            if i == self.n_fc - 1:
                # Apply sigmoid function.
                torch.sigmoid(self.fc[i](X))
            else:
                X = F.relu(self.fc[i](X))

        return out, (h_out, cell_out)

    def init_hidden(self, batch_size):
        weight = next(self.parameters()).data
        hidden = (
                weight.new(self.dir * self.n_lstm, batch_size, self.latent_dim).zero_(),
                weight.new(self.dir * self.n_lstm, batch_size, self.latent_dim).zero_(),
            )

        return hidden
</code></pre>
<p><strong>Questions</strong></p>
<p>Validation loss plateaues at 0.3 and it seems impossible to improve average precision score no matter what I try. I have tried tuning the network and also tried to model</p>
<ul>
<li>using hidden state</li>
<li>using lstm_output</li>
<li>weighted loss using <code>pos_weight</code> in <code>nn.BCELossWithLogits</code> (amending the sigmoid in the forward)</li>
</ul>
<p>Model auditing showed that:</p>
<ul>
<li>weights are being updated in training steps</li>
<li>grads do flow to the very back of the layers and look generally healthy</li>
</ul>
<p>Potential signs of issues:</p>
<ul>
<li>Although biases do change over time, weight parameters do not update especially for the LSTM layers, e.g. like seen below:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://i.stack.imgur.com/UBZS9.png" title="" rel="noopener nofollow ugc"><img src="https://i.stack.imgur.com/UBZS9.png" alt="" width="649" height="500"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use xlink:href="#far-image"></use></svg><span class="filename"></span><span class="informations">1154×889</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use xlink:href="#discourse-expand"></use></svg></div></a></div>
</li>
</ul>
<p>Any advise would be very much appreciated.</p>