LSTM from scratch in tensorflow 2

Question

I'm trying to make LSTM in tensorflow 2.1 from scratch, without using the one already supplied with keras (tf.keras.layers.LSTM), just to learn and code something. To do so, I've defined a class "Model" that when called (like with model(input)) it computes the matrix multiplications of the LSTM. I'm pasting here part of my code, the other parts are on github (link)

class Model(object):
    [...]

    def __call__(self, inputs):
        assert inputs.shape == (vocab_size, T_steps)
        outputs = []
        for time_step in range(T_steps):
            x = inputs[:,time_step]
            x = tf.expand_dims(x,axis=1)

            z = tf.concat([self.h_prev,x],axis=0)

            f = tf.matmul(self.W_f, z) + self.b_f
            f = tf.sigmoid(f)

            i = tf.matmul(self.W_i, z) + self.b_i
            i = tf.sigmoid(i)

            o = tf.matmul(self.W_o, z) + self.b_o
            o = tf.sigmoid(o)

            C_bar = tf.matmul(self.W_C, z) + self.b_C
            C_bar = tf.tanh(C_bar)

            C = (f * self.C_prev) + (i * C_bar)

            h = o * tf.tanh(C)

            v = tf.matmul(self.W_v, h) + self.b_v
            v = tf.sigmoid(v)

            y = tf.math.softmax(v, axis=0)

            self.h_prev = h
            self.C_prev = C

            outputs.append(y)
        outputs = tf.squeeze(tf.stack(outputs,axis=1))
        return outputs

But this neural netoworks has three problems:

1) it is way slow during training. In comparison a model that uses tf.keras.layers.LSTM() is trained more than 10 times faster. Why is this? Maybe because I didn't use a minibatch training, but a stochastic one?

2) the NN seems to not learn anything at all. After just some (very few!) training examples, the loss seems to settle down and it won't decrease anymore, but rather it oscillates around the reached value. After training, I tested the NN making it generate some text, but it just outputs non-sense gibberish. Why isn't learning anything?

3) the loss function outputs very high values. I've coded a categorical cross-entropy loss function but, with 100 characters long sequence, the value of the function is over 370 per training example. Shouldn't it be way lower than this?

I've wrote the loss function like this:

def compute_loss(predictions, desired_outputs):
    l = 0
    for i in range(T_steps):
        l -= tf.math.log(predictions[desired_outputs[i], i])
    return l

I know they're open questions, but unfortunately I can't make it works. So any answer, even a short answer that help me to make myself solve the problem, is fine :)

LSTM from scratch in tensorflow 2

0 Answers