How does an LSTM output the correct dimensions for classes?

Question

Take the below LSTM:

input: 5x1 matrix
hidden units: 256
output size (aka classes, 1 hot vector): 10x1 matrix

It is my understanding that an LSTM of this size will do the following:

$w_x$ = weight matrix at $x$

$b_x$ = bias matrix at $x$

activation_gate = tanh($w_1$ $\cdot$ input + $w_2$ $\cdot$ prev_output + $b_1$)

input_gate = sigmoid($w_3$ $\cdot$ input + $w_4$ $\cdot$ prev_output + $b_2$)

forget_gate = sigmoid($w_5$ $\cdot$ input + $w_6$ $\cdot$ prev_output + $b_3$)

output_gate = sigmoid($w_7$ $\cdot$ input + $w_8$ $\cdot$ prev_output + $b_4$)

The size of the output of each gate should be equal to the number of hidden units, ie, 256. The problem arrises when trying to convert to the correct final output size, of 10. If the forget gate outputs 256, then it is summed with the element wise product of the activation and input gate to find the new state, this will result in a hidden state of size 256. (Also in all my research I have not found anywhere whether this addition is actually addition, or simply appending the two matrices).

So if I have a hidden state of 256, and the output gate outputs 256, doing an element wise product of these two results in, surprise surprise, 256, not 10. If I instead ensure the output gate outputs a size of 10, this no longer works with the hidden state in an element wise product.

How is this handled? I can come up with many ways of doing it myself, but I want an identical replica of the basic LSTM unit, as I have some theories I want to test, and if it is even the slightest bit different it would make the research invalid.

Answer 1

After some more research, I have found the answer.

An LSTM is comprised of 1 LSTM cell that is continuously updated by passing in new inputs, the hidden state and the previous output. All nodes inside the LSTM cell are of size hidden_units. That means the output of the activation gate, input gate, forget gate and output gate are all of size hidden_units. As an extension of this, the size of the hidden state is also equal to hidden_units.

The problem arises when the desired output is not equal to that of hidden_units. This is fixed simple and easily, by slapping a basic feed forward net onto the end of the output of the LSTM, so even if it outputs 256, you can convert it to n nodes. These final output nodes are typically put through a softmax with cross-entropy loss applied.

I came to the conclusion by looking at the code included here: https://towardsdatascience.com/lstm-by-example-using-tensorflow-feb0c1968537

In which you can see the final output has its own separate weights and bias'.