How to Build LSTM or GRU for Sequence-to-Sequence Tasks in PyTorch vs TensorFlow

Use PyTorch for fine-grained control over hidden states and custom training loops, or TensorFlow/Keras for declarative state propagation and faster prototyping with return_state=True and initial_state parameters.

Sequence-to-sequence (seq2seq) models map an input token sequence to an output token sequence using an encoder-decoder architecture. Both the encoder and decoder can be implemented with LSTM or GRU recurrent cells. This guide walks through building these models in PyTorch and TensorFlow using the reference implementations from Microsoft's AI-For-Beginners repository.

Core Architecture: Encoder-Decoder Pattern

A seq2seq model consists of two components:

• Encoder: Processes the source sequence and compresses it into a hidden state vector
• Decoder: Unfolds the hidden state to generate the target sequence token-by-token

The encoder's final hidden state becomes the decoder's initial hidden state. During training, teacher forcing feeds ground-truth tokens to the decoder; during inference, the decoder uses its own predictions.

PyTorch Implementation: Manual State Management

In PyTorch, you instantiate torch.nn.LSTM or torch.nn.GRU directly and manage hidden states explicitly.

Cell Definition and Return Values

import torch
import torch.nn as nn

# LSTM returns (output, (h_n, c_n)) — hidden state AND cell state

lstm = nn.LSTM(input_size=256, hidden_size=512, batch_first=True)

# GRU returns (output, h_n) — hidden state only

gru = nn.GRU(input_size=256, hidden_size=512, batch_first=True)

Building the Encoder

The encoder in lessons/5-NLP/16-RNN/torchnlp.py demonstrates the LSTM pattern. For seq2seq, extract the final hidden state:

class Encoder(nn.Module):
    def __init__(self, vocab_sz, emb_dim, hid_dim, cell='LSTM'):
        super().__init__()
        self.embedding = nn.Embedding(vocab_sz, emb_dim)
        rnn_cls = nn.LSTM if cell == 'LSTM' else nn.GRU
        self.rnn = rnn_cls(emb_dim, hid_dim, batch_first=True)

    def forward(self, src):
        emb = self.embedding(src)                # (batch, src_len, emb_dim)

        outputs, hidden = self.rnn(emb)          # hidden = (h_n, c_n) for LSTM

        return hidden

Building the Decoder with Teacher Forcing

class Decoder(nn.Module):
    def __init__(self, vocab_sz, emb_dim, hid_dim, cell='LSTM'):
        super().__init__()
        self.embedding = nn.Embedding(vocab_sz, emb_dim)
        rnn_cls = nn.LSTM if cell == 'LSTM' else nn.GRU
        self.rnn = rnn_cls(emb_dim, hid_dim, batch_first=True)
        self.fc = nn.Linear(hid_dim, vocab_sz)

    def forward(self, trg, hidden):
        emb = self.embedding(trg)                # (batch, 1, emb_dim)

        output, hidden = self.rnn(emb, hidden)   # one step at a time

        pred = self.fc(output.squeeze(1))        # (batch, vocab_sz)

        return pred, hidden

Handling Variable-Length Sequences

PyTorch provides torch.nn.utils.rnn.pack_padded_sequence and pad_packed_sequence for efficient batching:

from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence

# Sort by length descending

lengths = [len(seq) for seq in sequences]
packed = pack_padded_sequence(padded, lengths, batch_first=True)
output, hidden = self.rnn(packed)
output, _ = pad_packed_sequence(output, batch_first=True)

Complete Training Loop

encoder = Encoder(VOCAB_SIZE, EMB_DIM, HID_DIM, cell='GRU')
decoder = Decoder(VOCAB_SIZE, EMB_DIM, HID_DIM, cell='LSTM')
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(list(encoder.parameters()) + list(decoder.parameters()))

for epoch in range(10):
    optimizer.zero_grad()
    src = torch.randint(0, VOCAB_SIZE, (32, 15))
    trg = torch.randint(0, VOCAB_SIZE, (32, 10))

    enc_hidden = encoder(src)
    dec_input = trg[:, 0].unsqueeze(1)
    loss = 0
    hidden = enc_hidden
    
    for t in range(1, trg.size(1)):
        pred, hidden = decoder(dec_input, hidden)
        loss += criterion(pred, trg[:, t])
        dec_input = trg[:, t].unsqueeze(1)  # teacher forcing

    
    loss.backward()
    optimizer.step()

The torchnlp.py file shows how to swap LSTM for GRU with a single line change: replace nn.LSTM with nn.GRU.

TensorFlow/Keras Implementation: Declarative State Flow

TensorFlow/Keras handles state propagation automatically through return_state=True and initial_state parameters.

Cell Definition with State Returns

import tensorflow as tf
from tensorflow.keras import layers

# With return_state=True: returns (output, h, c) for LSTM

lstm = layers.LSTM(units=512, return_state=True, name='encoder')

# With return_state=True: returns (output, h) for GRU

gru = layers.GRU(units=512, return_state=True, name='encoder')

Building the Encoder

The RNNTF.ipynb notebook demonstrates the return_state=True pattern essential for seq2seq encoders:

class Encoder(tf.keras.layers.Layer):
    def __init__(self, vocab_sz, emb_dim, hid_dim, cell='LSTM'):
        super().__init__()
        self.embedding = layers.Embedding(vocab_sz, emb_dim)
        rnn_cls = layers.LSTM if cell == 'LSTM' else layers.GRU
        self.rnn = rnn_cls(hid_dim, return_state=True, name='encoder')

    def call(self, x):
        x = self.embedding(x)
        if isinstance(self.rnn, layers.LSTM):
            output, state_h, state_c = self.rnn(x)
            return [state_h, state_c]  # list for LSTM

        else:
            output, state_h = self.rnn(x)
            return [state_h]            # list for GRU

Building the Decoder with Initial State

class Decoder(tf.keras.layers.Layer):
    def __init__(self, vocab_sz, emb_dim, hid_dim, cell='LSTM'):
        super().__init__()
        self.embedding = layers.Embedding(vocab_sz, emb_dim)
        rnn_cls = layers.LSTM if cell == 'LSTM' else layers.GRU
        self.rnn = rnn_cls(hid_dim, return_sequences=True, 
                          return_state=True, name='decoder')
        self.fc = layers.Dense(vocab_sz)

    def call(self, x, states):
        x = self.embedding(x)
        if isinstance(self.rnn, layers.LSTM):
            output, state_h, state_c = self.rnn(x, initial_state=states)
            states = [state_h, state_c]
        else:
            output, state_h = self.rnn(x, initial_state=states)
            states = [state_h]
        logits = self.fc(output)
        return logits, states

Complete Training with GradientTape

encoder = Encoder(VOCAB_SIZE, EMB_DIM, HID_DIM, cell='GRU')
decoder = Decoder(VOCAB_SIZE, EMB_DIM, HID_DIM, cell='LSTM')

loss_obj = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
optimizer = tf.keras.optimizers.Adam()

with tf.GradientTape() as tape:
    enc_states = encoder(src)
    dec_input = tf.expand_dims(trg[:, 0], 1)
    loss = 0
    states = enc_states
    
    for t in range(1, trg.shape[1]):
        logits, states = decoder(dec_input, states)
        loss += loss_obj(trg[:, t], tf.squeeze(logits, axis=1))
        dec_input = tf.expand_dims(trg[:, t], 1)  # teacher forcing

grads = tape.gradient(loss, encoder.trainable_variables + decoder.trainable_variables)
optimizer.apply_gradients(zip(grads, encoder.trainable_variables + decoder.trainable_variables))

Key Differences: PyTorch vs TensorFlow for Seq2Seq

Aspect	PyTorch	TensorFlow/Keras
State handling	Manual tuple unpack (h_n, c_n) vs h_n	Automatic via return_state and initial_state
Cell switching	Change nn.LSTM to nn.GRU — same API	Change layers.LSTM to layers.GRU — handles state count automatically
Variable lengths	pack_padded_sequence / pad_packed_sequence	Padded batches with layers.Masking
Training loop	Explicit for loop with manual loss accumulation	model.fit() or GradientTape with similar loop structure
Flexibility	Easier custom attention, beam search	Encapsulated higher-level APIs

Advanced Patterns from the Repository

Bidirectional Encoders

Both frameworks support bidirectional encoders for richer context:

PyTorch:

self.rnn = nn.LSTM(emb_dim, hid_dim, batch_first=True, bidirectional=True)

# Output is 2*hid_dim, handle concatenation in decoder

TensorFlow/Keras:

self.rnn = layers.Bidirectional(layers.LSTM(hid_dim, return_state=True))

# Returns forward and backward states; concatenate or sum as needed

The RNNTF.ipynb file shows bidirectional LSTM implementation for text classification.

Multi-Layer Stacking

The NER-TF.ipynb notebook implements a two-layer bidirectional LSTM encoder:

layers.Bidirectional(layers.LSTM(HID_DIM, return_sequences=True, return_state=True))

For PyTorch, set num_layers=2:

nn.LSTM(emb_dim, hid_dim, num_layers=2, batch_first=True, bidirectional=True)

TimeDistributed Decoders

The NER notebook uses TimeDistributed dense layers for per-token predictions:

layers.TimeDistributed(layers.Dense(num_tags, activation='softmax'))

In PyTorch, this is implicit — the linear layer operates on all timesteps when given the full sequence.

Inference: Greedy Decoding Without Teacher Forcing

PyTorch:

encoder.eval()
decoder.eval()
with torch.no_grad():
    hidden = encoder(src)
    dec_input = torch.full((batch_size, 1), START_ID)
    for _ in range(max_len):
        pred, hidden = decoder(dec_input, hidden)
        dec_input = pred.argmax(dim=-1).unsqueeze(1)

TensorFlow/Keras:

def greedy_decode(encoder, decoder, src_seq, max_len=20, start_id=1):
    enc_states = encoder(src_seq)
    dec_input = tf.expand_dims([start_id] * src_seq.shape[0], 1)
    outputs = []
    states = enc_states
    for _ in range(max_len):
        logits, states = decoder(dec_input, states)
        pred_id = tf.argmax(logits, axis=-1, output_type=tf.int32)
        outputs.append(pred_id)
        dec_input = pred_id
    return tf.concat(outputs, axis=1)

Reference Files in AI-For-Beginners

File	Description
lessons/5-NLP/16-RNN/torchnlp.py	PyTorch LSTM/GRU classifier with pack-sequence utilities
lessons/5-NLP/16-RNN/RNNTF.ipynb	TensorFlow bidirectional LSTM with return_state=True pattern
lessons/5-NLP/19-NER/NER-TF.ipynb	Full seq2seq with two-layer bidirectional LSTM encoder and TimeDistributed decoder
environment.yml	Dependency versions for TensorFlow, PyTorch, and CUDA

Summary

• PyTorch requires explicit handling of hidden-state tuples (h_n, c_n) for LSTM or h_n for GRU, offering fine-grained control for custom attention and beam search
• TensorFlow/Keras automates state propagation through return_state=True and initial_state, reducing boilerplate for standard architectures
• Both frameworks support bidirectional encoders, multi-layer stacking, and mixed cell types (LSTM encoder + GRU decoder or vice versa)
• The torchnlp.py file demonstrates single-line cell swapping; NER-TF.ipynb shows production-ready seq2seq patterns

Frequently Asked Questions

Can I mix LSTM and GRU between encoder and decoder?

Yes. In PyTorch, initialize encoder.rnn = nn.LSTM(...) and decoder.rnn = nn.GRU(...), then handle the state mismatch — GRU's single state h_n becomes both h_0 and c_0 for LSTM (or initialize c_0 to zeros). In TensorFlow, the initial_state parameter accepts lists of appropriate length automatically.

How do I handle variable-length sequences in batches?

Use PyTorch's pack_padded_sequence and pad_packed_sequence for computation efficiency. In TensorFlow, provide padded tensors and use layers.Masking or the mask parameter; Keras automatically propagates masks through compatible layers.

Which is faster for seq2seq training: PyTorch or TensorFlow?

Performance is comparable when using equivalent implementations. PyTorch's dynamic graphs allow easier debugging and modification. TensorFlow's tf.function compilation and XLA can optimize static graphs for production deployment. For research flexibility, prefer PyTorch; for production pipelines, consider TensorFlow's SavedModel format.

Where can I find working attention mechanisms for these models?

The AI-For-Beginners repository focuses on foundational RNNs. For attention, extend the decoder to compute context vectors using torch.bmm (PyTorch) or layers.Attention (TensorFlow). The torch.nn.MultiheadAttention and layers.MultiHeadAttention modules provide optimized implementations for transformer-style attention over recurrent outputs.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}