LSTM¶

Simple LSTM-based sequence-to-sequence model for neural machine translation.

Overview¶

SimpleSeq2SeqLSTM implements a classic encoder-decoder architecture using LSTM (Long Short-Term Memory) cells. The encoder reads the source sequence and compresses it into a context vector, which initializes the decoder to generate the target sequence.

Architecture¶

┌─────────────────────────────────────────────────┐
│                    ENCODER                      │
│  ┌───────────────────────────────────────────┐  │
│  │         Token Embedding                   │  │
│  │   src_vocab_size → emb_dim                │  │
│  └───────────────────────────────────────────┘  │
│                      ↓                          │
│  ┌───────────────────────────────────────────┐  │
│  │         LSTM Layers × N                   │  │
│  │   Process sequence step by step           │  │
│  │   Output: (hidden_state, cell_state)      │  │
│  └───────────────────────────────────────────┘  │
└─────────────────────────────────────────────────┘
                       ↓
           (final_hidden, final_cell)
                       ↓
┌─────────────────────────────────────────────────┐
│                    DECODER                      │
│  ┌───────────────────────────────────────────┐  │
│  │         Token Embedding                   │  │
│  │   tgt_vocab_size → emb_dim                │  │
│  └───────────────────────────────────────────┘  │
│                      ↓                          │
│  ┌───────────────────────────────────────────┐  │
│  │         LSTM Layers × N                   │  │
│  │   Initialized with encoder states         │  │
│  │   Process target sequence                 │  │
│  └───────────────────────────────────────────┘  │
│                      ↓                          │
│  ┌───────────────────────────────────────────┐  │
│  │         Linear Output                     │  │
│  │   hidden_dim → tgt_vocab_size             │  │
│  └───────────────────────────────────────────┘  │
└─────────────────────────────────────────────────┘

Quick Start¶

from torchlingo.models import SimpleSeq2SeqLSTM

model = SimpleSeq2SeqLSTM(
    src_vocab_size=10000,
    tgt_vocab_size=10000,
    emb_dim=256,
    hidden_dim=512,
    num_layers=2,
)

# Training
logits = model(src_batch, tgt_batch)  # [batch, tgt_len, vocab]

API Reference¶

SimpleSeq2SeqLSTM ¶

SimpleSeq2SeqLSTM(src_vocab_size: int, tgt_vocab_size: int, emb_dim: int = None, hidden_dim: int = None, num_layers: int = None, dropout: float = None, pad_idx: int = None, config: Config = None)

Bases: Module

Simple LSTM encoder-decoder model for sequence-to-sequence tasks.

This model encodes source sequences into a context vector using an LSTM encoder, then decodes target sequences using an LSTM decoder initialized with the encoder's final hidden and cell states. Token embeddings use padding index from config.

Parameters:

Name	Type	Description	Default
`src_vocab_size`	`int`	Size of the source vocabulary.	required
`tgt_vocab_size`	`int`	Size of the target vocabulary.	required
`emb_dim`	`int`	Embedding dimension. Falls back to config.lstm_emb_dim.	`None`
`hidden_dim`	`int`	Hidden dimension for LSTM layers. Falls back to config.lstm_hidden_dim.	`None`
`num_layers`	`int`	Number of LSTM layers in encoder and decoder. Falls back to config.lstm_num_layers.	`None`
`dropout`	`float`	Dropout rate applied between LSTM layers. Falls back to config.lstm_dropout.	`None`
`pad_idx`	`int`	Padding token index for embeddings. Falls back to config.pad_idx.	`None`
`config`	`Config`	Configuration object. Defaults to default config.	`None`

Attributes:

Name	Type	Description
`pad_idx`	`int`	Resolved padding index.
`src_embed`	`Embedding`	Source token embedding layer.
`tgt_embed`	`Embedding`	Target token embedding layer.
`encoder`	`LSTM`	LSTM encoder for source sequences.
`decoder`	`LSTM`	LSTM decoder for target sequences.
`output`	`Linear`	Linear output layer projecting decoder hidden state to target vocabulary.
`hidden_dim`	`int`	Hidden dimension size.

Source code in src/torchlingo/models/lstm_simple.py

def __init__(
    self,
    src_vocab_size: int,
    tgt_vocab_size: int,
    emb_dim: int = None,
    hidden_dim: int = None,
    num_layers: int = None,
    dropout: float = None,
    pad_idx: int = None,
    config: Config = None,
):
    super().__init__()
    cfg = config if config is not None else get_default_config()
    self.pad_idx = pad_idx if pad_idx is not None else cfg.pad_idx
    emb_dim = emb_dim if emb_dim is not None else cfg.lstm_emb_dim
    hidden_dim = hidden_dim if hidden_dim is not None else cfg.lstm_hidden_dim
    num_layers = num_layers if num_layers is not None else cfg.lstm_num_layers
    dropout = dropout if dropout is not None else cfg.lstm_dropout

    self.src_embed = nn.Embedding(src_vocab_size, emb_dim, padding_idx=self.pad_idx)
    self.tgt_embed = nn.Embedding(tgt_vocab_size, emb_dim, padding_idx=self.pad_idx)
    self.encoder = nn.LSTM(
        emb_dim,
        hidden_dim,
        num_layers=num_layers,
        batch_first=True,
        dropout=dropout,
    )
    self.decoder = nn.LSTM(
        emb_dim,
        hidden_dim,
        num_layers=num_layers,
        batch_first=True,
        dropout=dropout,
    )
    self.output = nn.Linear(hidden_dim, tgt_vocab_size)
    self.hidden_dim = hidden_dim

    # Initialize weights for better convergence
    self._init_weights()

forward ¶

forward(src: Tensor, tgt: Tensor) -> Tensor

Encode source and decode target sequences.

Encodes source tokens into context using the LSTM encoder, then passes the final hidden and cell states to the LSTM decoder to generate target token predictions. Attention masks are currently unused but accepted for API compatibility.

Parameters:

Name	Type	Description	Default
`src`	`Tensor`	Source token indices with shape (batch_size, src_len). Values should be in range [0, src_vocab_size).	required
`tgt`	`Tensor`	Target token indices with shape (batch_size, tgt_len). Values should be in range [0, tgt_vocab_size).	required

Returns:

Type	Description
`Tensor`	torch.Tensor: Logits of shape (batch_size, tgt_len, tgt_vocab_size) representing probability distributions over the target vocabulary for each position.

Source code in src/torchlingo/models/lstm_simple.py

def forward(
    self,
    src: torch.Tensor,
    tgt: torch.Tensor,
) -> torch.Tensor:
    """Encode source and decode target sequences.

    Encodes source tokens into context using the LSTM encoder, then passes
    the final hidden and cell states to the LSTM decoder to generate target
    token predictions. Attention masks are currently unused but accepted for
    API compatibility.

    Args:
        src (torch.Tensor): Source token indices with shape (batch_size, src_len).
            Values should be in range [0, src_vocab_size).
        tgt (torch.Tensor): Target token indices with shape (batch_size, tgt_len).
            Values should be in range [0, tgt_vocab_size).

    Returns:
        torch.Tensor: Logits of shape (batch_size, tgt_len, tgt_vocab_size) representing
            probability distributions over the target vocabulary for each position.
    """
    src_emb = self.src_embed(src)
    enc_out, (h, c) = self.encoder(src_emb)
    tgt_emb = self.tgt_embed(tgt)
    dec_out, _ = self.decoder(tgt_emb, (h, c))
    return self.output(dec_out)

Constructor Parameters¶

Parameter	Type	Default	Description
`src_vocab_size`	int	required	Source vocabulary size
`tgt_vocab_size`	int	required	Target vocabulary size
`emb_dim`	int	256	Embedding dimension
`hidden_dim`	int	512	LSTM hidden dimension
`num_layers`	int	2	Number of stacked LSTM layers
`dropout`	float	0.1	Dropout between LSTM layers
`pad_idx`	int	0	Padding token index
`config`	Config	None	Configuration object

Examples¶

Basic Training¶

import torch
from torchlingo.models import SimpleSeq2SeqLSTM

model = SimpleSeq2SeqLSTM(
    src_vocab_size=10000,
    tgt_vocab_size=10000,
    emb_dim=256,
    hidden_dim=512,
)

# Dummy data
src = torch.randint(0, 10000, (32, 20))  # [batch, src_len]
tgt = torch.randint(0, 10000, (32, 25))  # [batch, tgt_len]

# Forward pass
logits = model(src, tgt[:, :-1])  # [32, 24, 10000]

# Compute loss
criterion = torch.nn.CrossEntropyLoss(ignore_index=0)
loss = criterion(
    logits.reshape(-1, logits.size(-1)),
    tgt[:, 1:].reshape(-1)
)

With Config¶

from torchlingo.config import Config

config = Config(
    lstm_emb_dim=256,
    lstm_hidden_dim=512,
    lstm_num_layers=3,
    lstm_dropout=0.2,
)

model = SimpleSeq2SeqLSTM(
    src_vocab_size=10000,
    tgt_vocab_size=10000,
    config=config,
)

Inference (Greedy Decoding)¶

def greedy_decode_lstm(model, src, tgt_vocab, max_len=50, device="cpu"):
    model.eval()
    src = src.to(device).unsqueeze(0)  # Add batch dim

    # Encode
    src_emb = model.src_embed(src)
    _, (h, c) = model.encoder(src_emb)

    # Decode
    ys = [tgt_vocab.sos_idx]

    for _ in range(max_len):
        tgt_tensor = torch.tensor([ys[-1:]], device=device)
        tgt_emb = model.tgt_embed(tgt_tensor)
        output, (h, c) = model.decoder(tgt_emb, (h, c))
        logits = model.output(output)
        next_token = logits[:, -1, :].argmax().item()
        ys.append(next_token)

        if next_token == tgt_vocab.eos_idx:
            break

    return ys

How LSTMs Work¶

The Information Bottleneck¶

The encoder compresses the entire source sequence into a fixed-size vector (the final hidden state). This becomes the "context" for the decoder.

"I love cats" → [Encode] → hidden_vector → [Decode] → "Me gustan los gatos"

Limitation: Long sequences can be hard to compress into a single vector.

Hidden and Cell States¶

LSTMs maintain two types of state:

Hidden state (h): Short-term memory, used for output
Cell state (c): Long-term memory, carries information across time steps

# After encoding
# h: [num_layers, batch, hidden_dim]
# c: [num_layers, batch, hidden_dim]

Stacked Layers¶

Multiple LSTM layers create a deeper network:

Layer 3: Higher-level patterns
    ↑
Layer 2: Intermediate features  
    ↑
Layer 1: Low-level features
    ↑
Input embeddings

LSTM vs Transformer¶

Aspect	LSTM	Transformer
Processing	Sequential	Parallel
Long dependencies	Difficult	Easy (attention)
Training speed	Slower	Faster
Memory efficiency	O(n)	O(n²)
Simplicity	Simpler	More complex
Parameters	Fewer	More

When to Use LSTM¶

✅ Small datasets (< 50K examples)
✅ Limited GPU memory
✅ Learning/educational purposes
✅ Real-time inference on CPU

When to Use Transformer¶

✅ Large datasets
✅ Best translation quality
✅ GPU available for training
✅ Long sequences

Model Variants¶

Small¶

model = SimpleSeq2SeqLSTM(
    src_vocab_size, tgt_vocab_size,
    emb_dim=128,
    hidden_dim=256,
    num_layers=1,
)
# ~3M params

Medium¶

model = SimpleSeq2SeqLSTM(
    src_vocab_size, tgt_vocab_size,
    emb_dim=256,
    hidden_dim=512,
    num_layers=2,
)
# ~15M params

Large¶

model = SimpleSeq2SeqLSTM(
    src_vocab_size, tgt_vocab_size,
    emb_dim=512,
    hidden_dim=1024,
    num_layers=4,
)
# ~50M params

Limitations¶

No attention: This simple model doesn't use attention, so it struggles with long sequences
Information bottleneck: All source information must fit in the hidden state
Sequential processing: Can't parallelize across time steps

Adding Attention

For production LSTM models, consider adding Bahdanau or Luong attention to allow the decoder to look back at the encoder outputs. This is not implemented in TorchLingo's simple model.