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 = None, hidden_dim: int | None = None, num_layers: int | None = None, dropout: float | None = None, pad_idx: int | None = None, attention: bool | None = None, attn_type: str | None = None, config: Config | None = 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.

Without attention, the only channel from encoder to decoder is that final hidden state -- a fixed-size summary of the whole source sentence. With attention=True the decoder additionally attends over every encoder output, removing that bottleneck.

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`
`attention`	`bool`	Whether the decoder attends over encoder outputs. Falls back to config.lstm_attention (default `False`).	`None`
`attn_type`	`str`	Scoring function, `"dot"` (Luong) or `"additive"` (Bahdanau). Ignored when `attention` is False. Falls back to config.lstm_attn_type.	`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.
`attention`	`Module \| None`	Attention module, or `None` when disabled.
`attn_combine`	`Linear \| None`	Merges context with the decoder state, or `None` when attention is disabled.
`output`	`Linear`	Linear output layer projecting decoder hidden state to target vocabulary.
`hidden_dim`	`int`	Hidden dimension size.

Note

Enabling attention adds parameters, so a checkpoint saved with attention=False will not load into a model built with attention=True (and vice versa). Build the model the same way you trained it.

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

def __init__(
    self,
    src_vocab_size: int,
    tgt_vocab_size: int,
    emb_dim: int | None = None,
    hidden_dim: int | None = None,
    num_layers: int | None = None,
    dropout: float | None = None,
    pad_idx: int | None = None,
    attention: bool | None = None,
    attn_type: str | None = None,
    config: Config | None = 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
    attention = attention if attention is not None else cfg.lstm_attention
    attn_type = attn_type if attn_type is not None else cfg.lstm_attn_type

    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,
    )
    if attention:
        self.attention = build_attention(attn_type, hidden_dim)
        # Luong's "attentional hidden state": fold the context back into the
        # decoder state before predicting, so the context actually reaches
        # the output layer.
        self.attn_combine = nn.Linear(hidden_dim * 2, hidden_dim, bias=False)
    else:
        self.attention = None
        self.attn_combine = None

    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, return_attention: bool = False) -> Tensor | tuple[Tensor, Tensor | None]

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. When the model was built with attention=True, the decoder states are additionally blended with a weighted average of the encoder outputs before the output projection.

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
`return_attention`	`bool`	Also return the attention weights. Defaults to False, which preserves the plain-tensor return type.	`False`

Returns:

Type	Description
`Tensor \| tuple[Tensor, Tensor \| None]`	torch.Tensor: Logits of shape (batch_size, tgt_len, tgt_vocab_size) representing probability distributions over the target vocabulary for each position.
`Tensor \| tuple[Tensor, Tensor \| None]`	If `return_attention` is True, returns a `(logits, weights)`
`Tensor \| tuple[Tensor, Tensor \| None]`	tuple instead, where `weights` has shape
`Tensor \| tuple[Tensor, Tensor \| None]`	(batch_size, tgt_len, src_len) -- or is `None` when the model has
`Tensor \| tuple[Tensor, Tensor \| None]`	no attention.

Examples:

>>> model = SimpleSeq2SeqLSTM(50, 50, emb_dim=8, hidden_dim=8, attention=True)
>>> src, tgt = torch.randint(1, 50, (2, 5)), torch.randint(1, 50, (2, 3))
>>> logits, weights = model(src, tgt, return_attention=True)
>>> weights.shape
torch.Size([2, 3, 5])

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

def forward(
    self,
    src: torch.Tensor,
    tgt: torch.Tensor,
    return_attention: bool = False,
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor | None]:
    """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. When the model was built with ``attention=True``, the
    decoder states are additionally blended with a weighted average of the
    encoder outputs before the output projection.

    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).
        return_attention (bool, optional): Also return the attention weights.
            Defaults to False, which preserves the plain-tensor return type.

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

        If ``return_attention`` is True, returns a ``(logits, weights)``
        tuple instead, where ``weights`` has shape
        (batch_size, tgt_len, src_len) -- or is ``None`` when the model has
        no attention.

    Examples:
        >>> model = SimpleSeq2SeqLSTM(50, 50, emb_dim=8, hidden_dim=8, attention=True)
        >>> src, tgt = torch.randint(1, 50, (2, 5)), torch.randint(1, 50, (2, 3))
        >>> logits, weights = model(src, tgt, return_attention=True)
        >>> weights.shape
        torch.Size([2, 3, 5])
    """
    enc_out, hidden, src_pad_mask = self.encode_source(src)
    logits, _hidden, weights = self.decode_prefix(
        tgt, hidden, enc_out, src_pad_mask
    )
    return (logits, weights) if return_attention else logits

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
`attention`	bool	False	Let the decoder attend over encoder outputs
`attn_type`	str	`"dot"`	Scorer: `"dot"` (Luong) or `"additive"` (Bahdanau)
`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 and Beam Search)¶

Use the library's decoders rather than writing the loop yourself. Both accept an LSTM model:

from torchlingo.inference import beam_search_decode, greedy_decode

decoded = greedy_decode(model, src_batch, max_len=50)      # list[list[int]]
tokens = beam_search_decode(model, src, beam_size=5)       # one sentence

The search is the same code that decodes a Transformer — see Decoding. Three methods make that possible:

Method	Purpose
`encode_source(src)`	Returns `(enc_out, hidden, src_pad_mask)` — everything the decoder may need
`decode_prefix(tgt, hidden, enc_out, mask)`	Scores a whole target prefix; the counterpart to a Transformer's `decode(tgt, memory)`
`decode_step(token, hidden, enc_out, mask)`	Advances one token, carrying state forward

If you do want to step manually — to inspect attention at each step, say — use encode_source and decode_step so the encoder outputs actually reach the decoder:

enc_out, hidden, src_pad_mask = model.encode_source(src)
ys = [tgt_vocab.sos_idx]
alignments = []

for _ in range(max_len):
    last = torch.tensor([ys[-1:]], device=src.device)
    logits, hidden, weights = model.decode_step(last, hidden, enc_out, src_pad_mask)
    alignments.append(weights)          # None when attention is disabled
    next_token = logits[:, -1, :].argmax().item()
    ys.append(next_token)
    if next_token == tgt_vocab.eos_idx:
        break

Don't re-implement the decoder loop

Driving model.decoder directly from the encoder's final (h, c) — as earlier versions of this page showed — throws away the per-token encoder outputs. On an attention model that silently disables attention, and the model will appear to work while producing worse translations.

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.

With attention=True, the decoder additionally reads every encoder output, weighted per step, so the sentence no longer has to survive that squeeze. Run python examples/attention_alignment.py to see the difference measured on a task with a known correct alignment.

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; easier with `attention=True`	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¶

Information bottleneck (default): with attention=False, all source information must fit in the final hidden state
Sequential processing: can't parallelize across time steps

Adding Attention

Pass attention=True to let the decoder look back at every encoder output instead of relying on the final hidden state alone. Both classic scorers are implemented — see Attention. Attention is off by default, so the bottlenecked model remains the baseline you compare against.