A Diagram That Shows How LSTM Modules Work

LSTM (long short-term memory) software modules are fascinating. An LSTM component accepts a sequence of values that represent a sequence of words, and produce an abstract output that represents the entire sequence. The output representation can be used to predict things such as a classification label, an equivalent sequence in a different language, or a predicted next word.

I recently updated my personal LSTM diagram that I use before I start on an LSTM project.



There is a lot of information contained in the diagram. The lower case t means current time (or sequence number in the input). The x is the current word representation (a word embedding). The h is the output vector.

The internal c is the cell state — the LSTM memory.

If you follow the arrows, you can see that the output h(t) depends on the current word representation x(t), and on the previous output h(t-1), and on the previous cell state c(t-1).

LSTM modules save all the sequence outputs. Therefore, the current output h(t) is the same as the last output.

Here is some sample PyTorch code:

class LSTM_Net(T.nn.Module):
  def __init__(self):
    # vocab_size = 129892
    super(LSTM_Net, self).__init__()
    self.embed = T.nn.Embedding(129892, 32)
    self.lstm = T.nn.LSTM(32, 100)
    self.fc1 = T.nn.Linear(100, 2)  # 0=neg, 1=pos
 
  def forward(self, x):
    # x = review/sentence. length = fixed w/ padding
    z = self.embed(x)  # x can be arbitrary shape - not
    z = z.reshape(50, -1, 32)  # seq bat embed
    lstm_oupt, (h_n, c_n) = self.lstm(z)
    z = lstm_oupt[-1]  # or use h_n. [1,100]
    z = T.log_softmax(self.fc1(z), dim=1)  # CrossEntropy
    return z 

The numeric values of 12892, 32, 100, 2, -1 and 50, all have specific meanings. For example, the 100 is the size of the internal cell state c(t).

All of the internal mechanisms are a separate topic. A deep understanding of the internal workings is useful but isn’t completely necessary to create prediction systems that use LSTMs.

In some ways, you can think of an LSTM module as a tiny computer that input, memory, and output.

Interesting stuff!



There are some parallels between the incredible development of aviation and the equally incredible development of machine learning. Left: The British S.E.5a (1918) had a top speed of about 130 mph. Center: Just 20 years later, the U.S. Curtiss P-40 (1938) had a top speed of about 350 mph. Right: And 20 years later, the U.S. Lockheed F-104 (1958) could fly at 1,680 mph.


This entry was posted in Machine Learning. Bookmark the permalink.