RNN isn't learning, unsure what I'm doing wrong

Kylotan: <p>I’m trying to make a basic RNN model to use on some torchtext datasets, initially to try and complete an assignment in the Duke University ML course but having to piece together ideas from the internet because the instruction there is very lacking.</p>
<p>The problem I have is that there doesn’t appear to be any learning happening. After each epoch, the output accuracy is equal to chance, and loss does not seem to decrease at all. (It is a classification problem with 4 possible outputs and the correct predictions are about 1 in 4.)</p>
<p>The first thing of note is that it uses a word embedding layer, and the intent is that it learns this layer alongside the RNN itself. The Duke instructors suggest this is common but most of the examples I found online suggest otherwise, though the <code>nn.Embedding</code> layer does seem to support learning weights.</p>
<p>The next thing is that I am very confused about how (and if) to use the hidden argument to the RNN. Many examples online are using an <code>init_hidden</code> member function (often copy-and-pasted from somewhere else apparently), and imply that setting the hidden input to the RNN is important. However the NLP tutorial that Pytorch provides (<a href="https://docs.pytorch.org/tutorials/intermediate/char_rnn_classification_tutorial.html" class="inline-onebox" rel="noopener nofollow ugc">NLP From Scratch: Classifying Names with a Character-Level RNN — PyTorch Tutorials 2.8.0+cu128 documentation</a>) leaves it entirely unspecified.</p>
<p>I won’t post the whole code as I suspect most of it is irrelevant, but the model and the training loop are here:</p>
<pre data-code-wrap="python"><code class="lang-python">class MyRNN(nn.Module):
    def __init__(self, vocab_size: int, embedding_size: int, hidden_dim: int, num_outputs: int):
        super().__init__()
        self._hidden_dim = hidden_dim  # remember this value in case we need it

        self.embedding = nn.Embedding(num_embeddings=vocab_size, embedding_dim=embedding_size)
        self.rnn = nn.RNN(input_size=embedding_size, hidden_size=hidden_dim, num_layers=1)
        self.hidden_to_output = nn.Linear(hidden_dim, num_outputs)
        self.softmax = nn.LogSoftmax(dim=1)

    def forward(self, x):
        embed = self.embedding(x)
        rnn_out, hidden = self.rnn(embed)  # no Hx provided, "defaults to zeros". Is this right?!
        hidden_one_layer = hidden[0] # only one layer, so extract that
        pre_output = self.hidden_to_output(hidden_one_layer)
        output = self.softmax(pre_output)
        return output


model = MyRNN(VOCAB_SIZE, EMBED_DIM, HIDDEN_DIM, NUM_OUTPUTS)
model.cuda()
model.train()  # no effect here apparently, but could matter in other contexts

criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.005)


# Iterate through training set minibatches
for epoch in range(NUM_EPOCHS):
    print(f"Starting epoch {epoch}")
    epoch_loss = 0

    model.zero_grad()  # clear the gradients each epoch

    # train_loader is an instance of torch.utils.data.DataLoader,
    # wrapping a torchtext.datasets.TextClassificationDataset
    for inputs, labels in tqdm(train_loader):
        # push cpu stuff to cuda
        inputs = inputs.to(device)
        labels = labels.to(device)

        # Forward pass, measure the result loss
        y = model(inputs)
        loss = criterion(y, labels)

        # Backward pass
        loss.backward()

        optimizer.step()
        optimizer.zero_grad()

        epoch_loss += loss.item()

    print(f'Epoch [{epoch + 1}/{NUM_EPOCHS}], Loss: {epoch_loss / len(train_loader):.4f}')
</code></pre>
<p>I’d be grateful to hear about what I am doing wrong or suboptimally, or things I could try if it’s not quite as obvious as that.</p>

WorkspaceTool: <p>That’s great to know. Many learners went through the same phase.</p>

J_Johnson: <p>Seems you really have two questions.</p>
<ol>
<li>How and when to reset the hidden layer.</li>
<li>Why is your model not learning.</li>
</ol>
<p>And you suspect 2 is tied to 1.</p>
<p>I don’t think the hidden layer is your issue here. The tutorial you link uses NLLoss, which would be appropriate to use a softmax with. But you’re using CrossEntropyLoss, which already uses softmax inside. So, essentially, you’re doubly applying softmax to your outputs. Either remove the softmax or change the loss function to NLLoss, and see if that helps.</p>

Kylotan: <p>Removing the softmax from the module made no difference, as did leaving softmax in place and switching the loss function to NLLoss.</p>

J_Johnson: <p>A couple of other things you can troubleshoot are:</p>
<ol>
<li>Lower the learning rate;</li>
<li>Try with the dataset from the tutorial, and if your model/training method learns on that, then we know the issue is with your dataset;</li>
<li>Check the class distribution of your dataset - if it’s not balanced, you may need to apply a weight vector into the loss function. See here <a href="https://discuss.pytorch.org/t/how-do-i-get-the-accuracy-of-an-unbalanced-dataset-or-segmentation-task/194579/2" class="inline-onebox">How do I get the accuracy of an unbalanced dataset or segmentation task? - #2 by J_Johnson</a></li>
</ol>

Kylotan: <p>The learning rate was 0.005 - trying it at 0.0005 had no effect, as did 0.0001. Even with 10x the epochs the accuracy stays between 24.8% and 25.2% at all times, so it is not converging at all.</p>
<p>Additionally, I would guess that this stubborn 1 in 4 result for a problem with 4 outputs implies the dataset is balanced.</p>
<p>I’m not following the tutorial linked above - I’m just trying to use it to learn how to use an RNN, and apply it here. The dataset I am using can be easily learned with a simple Embedding &gt; Linear &gt; Linear &gt; ReLU where all the inputs to the embedding are averaged before the next layer (called a ‘Simple Word Embedding Model’ elsewhere) to an accuracy of 90%. So I am confident that the data is okay - the problem is just that I don’t know how to get this RNN to work with it.</p>

Kylotan: <p>To be specific, if I swap out my RNN above with the following model, it gives me 90% accuracy on the same data set with all the same hyperparameters.</p>
<pre data-code-wrap="python"><code class="lang-python">class SWEM(nn.Module):
    def __init__(self, vocab_size, embedding_size, hidden_dim, num_outputs):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embedding_size)
        self.fc1 = nn.Linear(embedding_size, hidden_dim)
        self.fc2 = nn.Linear(hidden_dim, num_outputs)

    def forward(self, x):
        embed = self.embedding(x)
        embed_mean = torch.mean(embed, dim=0)
        h = self.fc1(embed_mean)
        h = torch.nn.functional.relu(h)
        h = self.fc2(h)
        return h
</code></pre>
<p>So this, to me, implies one of two things: either the RNN code I have has a bug, or it can’t cope with the input data somehow. The input <code>x</code> consists of padded sequences, collated in this way by the <code>DataLoader</code>:</p>
<pre data-code-wrap="python"><code class="lang-python">def collator(batch):
    labels = torch.tensor([example[0] for example in batch])
    sentences = [example[1] for example in batch]
    data = pad_sequence(sentences)
    return [data, labels]
</code></pre>
<p>The data itself is the ag_news dataset in CSV form, looking like this: <a href="https://www.kaggle.com/datasets/amananandrai/ag-news-classification-dataset" class="inline-onebox" rel="noopener nofollow ugc">AG News Classification Dataset | Kaggle</a></p>

vdw: <p>What is the shape of <code>x</code> as the input of the <code>forward()</code> method?</p>

J_Johnson: <p>A linear layer will still extract meaningful relationships simply because the words, regardless of order, have semantic meaning for the task. Where RNNs or Transformers come in is when the order of the words also carries some additional information. These models are specifically designed to extract meaning from sequential data.</p>
<p>I’m kind of curious to see more on your DataLoader and to learn what, if any, transforms you’re performing. If you’re not setting <code>batch_first = True</code> on the RNN layer, it’s quite likely that you’re inserting the sequence dim where the batch dim should be.</p>

Kylotan: <p>It is [N, 128] where the N is likely the number of values in the sequence and 128 is the batch size.</p>

Kylotan: <p>Yes, I’m aware that the RNN approach is not strictly necessary for good results on this sort of classification problem but I would expect it to be able to perform better as it can potentially spot things like negations of keywords, or avoid false positives in metaphors. (e.g. “Kill two birds with one stone” does not imply a ‘nature’ topic, even if the nouns and verbs alone might usually do so.)</p>
<p>The data loader is simply this: <code>torch.utils.data.DataLoader(agnews_train, batch_size=BATCH_SIZE, shuffle=True, collate_fn=collator)</code> with the collator above. There’s no other transformation of the data set, beyond torchtext loading in the csv.</p>
<p>I am not setting <code>batch_first=True</code> anywhere. But (as in my previous message) I believe I am providing the sequence as the first dimension and the batch as the second dimension.</p>
<p>I have seen other people use the batch_first option and I don’t fully understand why it is necessary. If I add that argument to the rnn constructor then when I run the model it confirms that the batch sizes don’t match.</p>

Kylotan: <p>Update on this: I still haven’t been able to solve it, but visualising the data in PyCharm has been quite interesting.</p>
<p>After each minibatch of input, if I visualise all the predictions as a 128x4 array (batch size of 128, 4 output categories) then the 128 rows are all almost identical. Out of 128, there are usually 1 or 2 which are quite different values, and the rest are <em>almost</em> the same value - usually 3 or 4 significant digits are identical with tiny differences at the least significant digits which make no difference to which of the 4 values is highest.Thus almost all the inputs typically yield the same prediction, meaning it’s right 1 in 4 times (given the balanced dataset).</p>
<p>I can’t understand what is happening for the system to be able to be presented 128 input tensors and to produce 126 or 127 near-identical values and 1 or 2 different ones!</p>

J_Johnson: <p>Had a chance to give your code a try and got it working. Got above 99% accuracy. Took 1000 epochs, and it didn’t really start going above 26% until after the 300th epoch. But then it started taking off shortly after passing 40% accuracy. My guess is the weight initialization is part of the issue, and the other is that embeddings used in NLP tasks are typically pre-trained(takes a long time to train them, from what I understand).</p>
<p>Here’s what I set the optimizer as:</p>
<pre><code class="lang-auto">optimizer = torch.optim.SGD(model.parameters(), lr=0.00001, momentum=True)
</code></pre>
<p>Other hyper-parameters were:</p>
<pre><code class="lang-auto">num_examples = 10000
max_seq_len = 20
VOCAB_SIZE = 500  
EMBED_DIM = 64
HIDDEN_DIM = 128
NUM_OUTPUTS = 4
BATCH_SIZE = 128
</code></pre>
<p>And I used the following code to generate examples:</p>
<pre><code class="lang-auto">def generate_learnable_dataset(num_examples=num_examples, max_seq_len=max_seq_len, vocab_size=VOCAB_SIZE):
    """
    Generate a learnable dummy dataset for AG News classification.
    Each class has distinct token patterns to make sequences learnable.

    Args:
        num_examples: Number of examples to generate (default: 100).
        max_seq_len: Maximum sequence length (default: 20).
        vocab_size: Size of the vocabulary (default: 10000).

    Returns:
        List of tuples (label, sequence), where sequence is a 1D tensor.
    """
    agnews_train = []

    # Ensure vocab_size is at least 6 (2 for start/end tokens, 1 per class)
    if vocab_size &lt; 6:
        raise ValueError("vocab_size must be at least 6 to accommodate start token (0), end token (1), and class-specific tokens.")

    # Define class-specific token ranges dynamically based on vocab_size
    # Reserve tokens 0 (start) and 1 (end)
    available_tokens = max(4, vocab_size - 2)  # At least 1 token per class
    tokens_per_class = available_tokens // 4  # Divide among 4 classes

    class_tokens = {
        0: list(range(2, 2 + tokens_per_class)),
        1: list(range(2 + tokens_per_class, 2 + 2 * tokens_per_class)),
        2: list(range(2 + 2 * tokens_per_class, 2 + 3 * tokens_per_class)),
        3: list(range(2 + 3 * tokens_per_class, min(2 + 4 * tokens_per_class, vocab_size)))
    }

    # Define range for random tokens (excluding class-specific tokens and 0, 1)
    random_token_start = min(2 + 4 * tokens_per_class, vocab_size)
    random_tokens_range = list(range(random_token_start, vocab_size)) or [2]  # Fallback to [2] if empty

    for _ in range(num_examples):
        # Randomly select a label (0–3)
        label = random.randint(0, 3)

        # Random sequence length between 3 and max_seq_len (excluding start/end tokens)
        seq_len = random.randint(3, max_seq_len)

        # Generate sequence with class-specific patterns
        sequence = []
        # Add start token
        sequence.append(0)

        # Add class-specific tokens (50% of the sequence) and random tokens (50%)
        num_class_tokens = seq_len // 2
        num_random_tokens = seq_len - num_class_tokens

        # Add class-specific tokens
        sequence.extend(random.choices(class_tokens[label], k=num_class_tokens))
        # Add random tokens
        sequence.extend(random.choices(random_tokens_range, k=num_random_tokens))

        # Shuffle the sequence (excluding start/end tokens) to avoid fixed positions
        random.shuffle(sequence[1:])
        # Add end token
        sequence.append(1)

        # Convert to tensor
        sequence_tensor = torch.tensor(sequence)
        agnews_train.append((label, sequence_tensor))

    return agnews_train
</code></pre>
<p>So the two things I’d suggest are trying some different weight initializations and/or trying a pre-trained embedding layer.</p>
<p>Cheers</p>

Kylotan: <p>Thanks for doing that - I’ll take the time to read it fully later today when I get a chance to try this for myself.</p>
<p>I understand that learning embeddings can be expensive but what I don’t understand is that a much simpler model that uses a ‘bag of words’ approach (see here: <a href="https://discuss.pytorch.org/t/rnn-isnt-learning-unsure-what-im-doing-wrong/222169/7" class="inline-onebox">RNN isn't learning, unsure what I'm doing wrong - #7 by Kylotan</a>) learns far more quickly, albeit not reaching the same accuracy level. I appreciate the RNN is a more complex model and would be expected to take longer to train, but I’m surprised that it’s 2 orders of magnitude different. Do you know if this is expected? Of course, I have no idea how well the embedding has learned the vocabulary, just that the whole model works relatively well.</p>
<p>Finally - what is the purpose of the dummy dataset? Is the idea to have a carefully controlled set of data that proves the model can work in theory, before trying it on real world data? And if the tokens are shuffled, doesn’t that work against the RNN which is specialised in understanding sequential data?</p>

J_Johnson: <p>That’s only shuffled initially. Not between epochs. The other issue is likely both the embedding and RNN weight initialization method’s that PyTorch uses may not necessarily be optimal. You can try something like this for initializing the weights:</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
import torch.nn.init as init

def init_weights(module):
    if isinstance(module, nn.Embedding):
        # Xavier uniform for embeddings
        init.xavier_uniform_(module.weight)
    elif isinstance(module, nn.RNN):
        # Custom init for RNN
        for name, param in module.named_parameters():
            if 'weight_ih' in name:
                init.xavier_uniform_(param.data)  # Input-to-hidden
            elif 'weight_hh' in name:
                init.orthogonal_(param.data)     # Hidden-to-hidden
            elif 'bias' in name:
                param.data.fill_(0)              # Biases to zero

# Example usage:
 model = YourModelWithRNNAndEmbedding()
 model.apply(init_weights)
</code></pre>
<p>I haven’t tried it, yet, so let me know how it goes, if you do.</p>
<p>You could also try CBOW or Skip-gram pre-training methods on your embedding layer.</p>
<p>Why do I prefer building dummy datasets? Because it can easily be copied and pasted by others to try and narrow down any possible issues, without requiring downloading any datasets. It may not be perfect, but it served it’s purpose in demonstrating your model works.</p>

vdw: <p>I can confirm <a class="mention" href="/u/j_johnson">@J_Johnson</a> ‘s observation. When training RNNs, I often saw the loss hardly moving for the first many iterations.</p>