Understanding output of LSTM for regression

Question

I am working with embeddings and wanted to see how feasible it is to predict some scores attached to some sequences of words. The details of the scores are not important.

Input (tokenized sentence): ('the', 'dog', 'ate', 'the', 'apple')
Output (float): 0.25

I have been following this tutorial which tries to predict part-of-speech tags of such input. In such case, the output of the system is a distribution of all possible tags for all tokens in the sequence, e.g. for three possible POS classes {'DET': 0, 'NN': 1, 'V': 2}, the output for ('the', 'dog', 'ate', 'the', 'apple') could be

tensor([[-0.0858, -2.9355, -3.5374],
        [-5.2313, -0.0234, -4.0314],
        [-3.9098, -4.1279, -0.0368],
        [-0.0187, -4.7809, -4.5960],
        [-5.8170, -0.0183, -4.1879]])

Each row is a token, the index of the highest value in a token is the best predicted POS tag.

I understand this example relatively well, so I wanted to adapt it to a regression problem. The full code is below, but I am trying to make sense of the output.

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

torch.manual_seed(1)


class LSTMRegressor(nn.Module):
    def __init__(self, embedding_dim, hidden_dim, vocab_size):
        super(LSTMRegressor, self).__init__()
        self.hidden_dim = hidden_dim

        self.word_embeddings = nn.Embedding(vocab_size, embedding_dim)

        # The LSTM takes word embeddings as inputs, and outputs hidden states
        # with dimensionality hidden_dim.
        self.lstm = nn.LSTM(embedding_dim, hidden_dim)

        # The linear layer that maps from hidden state space to a single output
        self.linear = nn.Linear(hidden_dim, 1)
        self.hidden = self.init_hidden()

    def init_hidden(self):
        # Before we've done anything, we dont have any hidden state.
        # Refer to the Pytorch documentation to see exactly
        # why they have this dimensionality.
        # The axes semantics are (num_layers, minibatch_size, hidden_dim)
        return (torch.zeros(1, 1, self.hidden_dim),
                torch.zeros(1, 1, self.hidden_dim))

    def forward(self, sentence):
        embeds = self.word_embeddings(sentence)

        lstm_out, self.hidden = self.lstm(embeds.view(len(sentence), 1, -1), self.hidden)
        regression = F.relu(self.linear(lstm_out.view(len(sentence), -1)))

        return regression


def prepare_sequence(seq, to_ix):
    idxs = [to_ix[w] for w in seq]
    return torch.tensor(idxs, dtype=torch.long)

# ================================================

training_data = [
    ("the dog ate the apple".split(), 0.25),
    ("everybody read that book".split(), 0.78)
]

word_to_ix = {}
for sent, tags in training_data:
    for word in sent:
        if word not in word_to_ix:
            word_to_ix[word] = len(word_to_ix)

tag_to_ix = {"DET": 0, "NN": 1, "V": 2}

# ================================================

EMBEDDING_DIM = 6
HIDDEN_DIM = 6

model = LSTMRegressor(EMBEDDING_DIM, HIDDEN_DIM, len(word_to_ix))
loss_function = nn.MSELoss()
optimizer = optim.Adam(filter(lambda p: p.requires_grad, model.parameters()))

# See what the results are before training
with torch.no_grad():
    inputs = prepare_sequence(training_data[0][0], word_to_ix)
    regr = model(inputs)

    print(regr)

for epoch in range(100):  # again, normally you would NOT do 300 epochs, it is toy data
    for sentence, target in training_data:
        # Step 1. Remember that Pytorch accumulates gradients.
        # We need to clear them out before each instance
        model.zero_grad()

        # Also, we need to clear out the hidden state of the LSTM,
        # detaching it from its history on the last instance.
        model.hidden = model.init_hidden()

        # Step 2. Get our inputs ready for the network, that is, turn them into
        # Tensors of word indices.
        sentence_in = prepare_sequence(sentence, word_to_ix)
        target = torch.tensor(target, dtype=torch.float)

        # Step 3. Run our forward pass.
        score = model(sentence_in)

        # Step 4. Compute the loss, gradients, and update the parameters by
        #  calling optimizer.step()
        loss = loss_function(score, target)
        loss.backward()
        optimizer.step()

# See what the results are after training
with torch.no_grad():
    inputs = prepare_sequence(training_data[0][0], word_to_ix)
    regr = model(inputs)

    print(regr)

The output is:

# Before training
tensor([[0.0000],
        [0.0752],
        [0.1033],
        [0.0088],
        [0.1178]])
# After training
tensor([[0.6181],
        [0.4987],
        [0.3784],
        [0.4052],
        [0.4311]])

But I don't understand why. I was expecting a single output. The size of the tensor is the same as the number of tokens of the input. I would, then, guess that for each step in the input, the hidden state is given. Is that correct? Does that mean that the last item in the tensor (tensor[-1], or is it the first tensor[0]?) is the final prediction? Why are all outputs given? Or lies my misunderstanding earlier in the forward-pass? Perhaps I should only feed the last item of the LSTM layer to the linear layer?

I am also interested to know how this extrapolates to bidirectional LSTMs and multilayer LSTMs, and even how this would work with GRUs (bidirectional or not).

The bounty will be given to the person who can explain why we would use the last output or the last hidden state or what the difference means from a goal-directed perspective. In addition, some information about multilayer architectures and bidirectional RNNs is welcome. For instance, is it common practice to sum or concatenate the output and hidden state of bidirectional LSTM/GRU to get your data into sensible shape? If so, how do you do it?

Answer 1

Part of the confusion could be because POS tagging is not a regression problem. POS tagging is a classification problem. Classification is predicting the most likely category from a set of discrete categories. Regression is predicting a numeric output.

POS tagging is multi-class classification (e.g., DET, NN, V, ...). The last layer in a neural network for a multi-class classification should be a softmax function. A softmax function will convert the activation of the nodes for each category into a probability. The largest probability will be the predicted category, in this case a POS tag.

Answer 2

I think what you are missing here is a clear understanding of LSTMs. I will answer in parts.

"why we would use the last output or the last hidden state": The last hidden state and the last output are not the same. The last output is generated from the last hidden state by passing it through a linear layer, such as softmax.

"what the difference means from a goal-directed perspective": The last hidden state is simply a set of weights, while the last output is a prediction based on those weights.

"is it common practice to sum or concatenate the output and hidden state of bidirectional LSTM/GRU to get your data into sensible shape": It is not impossible to concatenate the output and the hidden state, because, well, they are tensors. But, it is inadvisable to do so. Also, it makes no practical sense. The output is generated by using the hidden state as one of the inputs in LSTMs.

Regarding the main question, the final tensor you are getting is simply logits. If you want a single value out of it, like the other answer states, you should stack a softmax layer on top of your existing architecture.

Answer 3

Very nice question; In the field of time-series prediction, LSTMs use sequence prediction. The LSTM networks can learn to predict which value is next to a given sequence of inputs that are related on time.

Remembering that input of a LSTM network should be normalized previously, so having a single vector of y1, y2, y3, ..., yn, the LSTM will try to learn, memorize and forget the relations within the instances in the vector; the memorizing and forget rate are configurable in LSTM model and by having layers of neurons, LSTM can create complex structure of relations in order to predict the next value minimizing the error of sequence prediction.

Now if you have exogenous variables in your scenario, LSTM can learn more by looking at the extra information, i.e. provided vectors of variables. The other point is that if you have not a good number of historical points of your time-series corresponding a large sequence of values, normally the LSTMs have inferior results to the mathematical models.

Here, I provide the formal description from the paper [1]: A LSTM network is consist of a chain of cells while each LSTM cell is configured mainly by four gates: input gate, input modulation gate, forget gate and output gate. Input gate takes a new input point from outside and process newly coming data. Memory cell input gate takes input from the output of the LSTM cell in the last iteration. Forget gate decides when to forget the output results and thus selects the optimal time lag for the input sequence. Output gate takes all results calculated and generate output. For a time-series input as $X= (x_1, x_2, ..., x_T)$, hidden state cells as $H = (h_1, h_2, ..., h_T)$ and output sequence as $Y = (y_1, y_2, ..., y_T)$.

For $t = 1, ..., T$ LSTM computes:

\begin{equation} \label{eq:lstm1} h_t = H(W_{hy}x_t + W_{hh}h_{t-1} + b_h) \end{equation} \begin{equation} \label{eq:lstm2} y_t = W_{hy}h_t + b_y \end{equation}

[1]: Long Short-Term Memory Sepp Hochreiter and Jürgen Schmidhuber Posted Online March 13, 2006, https://www.mitpressjournals.org/doi/abs/10.1162/neco.1997.9.8.1735

Answer 4

I went ahead, and testing a lot of stuff out and I came up with this network which seems to work alright as far as I have tested it

def __init__(self, hidden_dim, ms_dim, embeddings):
        super(LSTMRegressor, self).__init__()
        self.hidden_dim = hidden_dim

        # load pretrained embeddings, freeze them
        self.word_embeddings = nn.Embedding.from_pretrained(embeddings)
        embed_size = embeddings.shape[1]
        self.word_embeddings.weight.requires_grad = False

        self.w2v_lstm = nn.LSTM(embed_size, hidden_dim, bidirectional=True)

        self.ms_lstm = nn.LSTM(ms_dim, hidden_dim, bidirectional=True)

        self.linear = nn.Linear(hidden_dim, 1)
        self.relu = nn.LeakyReLU()

    def forward(self, batch_size, sentence_input, ms_input):
        # 1. Embeddings
        embeds = self.word_embeddings(sentence_input)
        w2v_out, _ = self.w2v_lstm(embeds.view(-1, batch_size, embeds.size(2)))

        # separate bidirectional output into first/last, then sum them
        w2v_first_bi = w2v_out[:, :, :self.hidden_dim]
        w2v_last_bi = w2v_out[:, :, self.hidden_dim:]
        w2v_sum_bi = w2v_first_bi + w2v_last_bi

        # 2. Other features
        ms_out, _ = self.ms_lstm(ms_input.view(-1, batch_size, ms_input.size(1)))

        ms_first_bi = ms_out[:, :, :self.hidden_dim]
        ms_last_bi = ms_out[:, :, self.hidden_dim:]
        ms_sum_bi = ms_first_bi + ms_last_bi

        # 3. Concatenate LSTM outputs
        summed = torch.cat((w2v_sum_bi, ms_sum_bi))

        # 4. Only use the last item of the sequence's output
        summed = summed[-1, :, :]

        # 5. Send output to linear layer, then ReLU
        regression = self.linear(summed)
        regression = self.relu(regression)

        return regression