LSTM vs Transformer: The Evolution of Sequence Learning in AI

LSTM vs Transformer: The Evolution of Sequence Learning in AI

The transformer did not replace the LSTM because it had a smarter idea. It replaced it because it solved a bottleneck no amount of clever LSTM engineering could touch: recurrent computation cannot be parallelized, and GPUs are built for parallel work.

LSTMs were the best available tool for sequence modeling from 1997 to roughly 2018 — not a naive stepping stone. Understanding why they work, and precisely where they break, is what makes the transformer's design choices make sense.

The Fundamental Problem of Sequences

Text, audio, time series, and code are sequences. Their meaning depends on order. "The cat sat on the mat" means something different from "the mat sat on the cat."

Fully-connected networks are permutation-invariant — shuffle the inputs and you get the same output. CNNs handle local order through convolution but cannot capture long-range dependencies across hundreds of tokens. Sequences need architectures that explicitly model temporal or positional relationships.

The key challenges:

1. Variable length inputs: sentences have different numbers of words
2. Long-range dependencies: the pronoun "it" in a sentence might refer to a noun from 20 words earlier
3. Sequential processing vs. parallelism: processing token by token is sequential; training data efficiently requires parallelism

LSTMs solve problems 1 and 2 brilliantly. They fail at problem 3, and that failure is what transformers exploit.

How LSTMs Actually Work

An LSTM (Long Short-Term Memory network) is a recurrent network with a dedicated memory cell that information can be written to, read from, or erased from — controlled by learned gates, the way a notebook lets you decide what to keep, what to cross out, and what to write down next. Hochreiter and Schmidhuber published the architecture in 1997, solving the vanishing-gradient problem that crippled plain RNNs.

At each timestep t, an LSTM receives the current input x_t and the previous hidden state h_{t-1}, then computes:

f_t = σ(W_f · [h_{t-1}, x_t] + b_f)    # forget gate: what to erase
i_t = σ(W_i · [h_{t-1}, x_t] + b_i)    # input gate: what to write
g_t = tanh(W_g · [h_{t-1}, x_t] + b_g) # candidate values
o_t = σ(W_o · [h_{t-1}, x_t] + b_o)    # output gate: what to expose

c_t = f_t ⊙ c_{t-1} + i_t ⊙ g_t       # update cell state
h_t = o_t ⊙ tanh(c_t)                  # compute hidden state

Where σ is sigmoid (outputs near 0 or 1) and ⊙ is element-wise multiplication.

The genius is the cell state c_t. When the forget gate is near 1 and the input gate is near 0, the cell state flows unchanged: c_t ≈ c_{t-1}. Information can persist across hundreds of timesteps without vanishing. When the forget gate is near 0, the cell erases old information. The gates themselves are learned — the network figures out what to remember.

import torch
import torch.nn as nn

class LSTMClassifier(nn.Module):
    """
    LSTM for text classification.
    Input: tokenized sequences → Output: class probabilities
    """
    def __init__(self, vocab_size, embed_dim, hidden_dim, num_classes, num_layers=2):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embed_dim, padding_idx=0)
        
        # LSTM processes the sequence token by token
        self.lstm = nn.LSTM(
            input_size=embed_dim,
            hidden_size=hidden_dim,
            num_layers=num_layers,
            batch_first=True,      # input shape: [batch, seq_len, features]
            dropout=0.3,
            bidirectional=True     # process forward and backward
        )
        
        # bidirectional doubles the hidden size
        self.classifier = nn.Sequential(
            nn.Dropout(0.3),
            nn.Linear(hidden_dim * 2, num_classes)
        )
    
    def forward(self, x):
        # x shape: [batch, seq_len]
        embedded = self.embedding(x)            # [batch, seq_len, embed_dim]
        
        # output: [batch, seq_len, hidden*2]
        # h_n:    [num_layers*2, batch, hidden] — final hidden states
        output, (h_n, c_n) = self.lstm(embedded)
        
        # Use the last hidden state from both directions
        # h_n[-2]: last forward layer, h_n[-1]: last backward layer
        hidden = torch.cat([h_n[-2], h_n[-1]], dim=1)
        
        return self.classifier(hidden)

The Bottleneck: Sequential Computation

LSTMs must process a sequence strictly left to right — computing h_t requires h_{t-1}, which requires h_{t-2}, all the way back to the start. It's an assembly line where station 10 can't start until station 9 finishes.

On GPU hardware — which is built for parallel computation — this sequential dependency is devastating. A sequence of 512 tokens requires 512 sequential LSTM steps. All 512 cannot execute in parallel because each depends on the previous.

The practical consequence: training a large LSTM on a large dataset is painfully slow. The parallelism that makes CNNs train fast on image data is not available for recurrent models on sequence data.

There is a second problem: the information bottleneck. Everything the network needs to remember must be compressed into the hidden state h_t. For a long document, by the time the LSTM reaches the end, information from the beginning must have been compressed through hundreds of sequential transformations. The forget gate cannot perfectly preserve everything — information fades.

Attention: The Key Idea

Attention lets a model look directly at every earlier position in a sequence, instead of relying on one compressed summary vector. Bahdanau et al. (2015) first bolted this onto encoder-decoder LSTMs for machine translation, letting the decoder consult all encoder hidden states rather than a single bottleneck vector.

# Bahdanau-style attention
def attention(decoder_hidden, encoder_outputs):
    # decoder_hidden: [1, hidden]
    # encoder_outputs: [seq_len, hidden]
    
    # Compute alignment scores
    scores = torch.tanh(
        linear_decoder(decoder_hidden) + 
        linear_encoder(encoder_outputs)
    )
    scores = linear_v(scores).squeeze(-1)   # [seq_len]
    
    # Normalize to weights summing to 1
    weights = torch.softmax(scores, dim=0)  # [seq_len]
    
    # Weighted sum of encoder outputs
    context = (weights.unsqueeze(-1) * encoder_outputs).sum(0)
    
    return context, weights

This improved translation quality substantially, especially for long sentences. The attention weights are also interpretable — you can visualize which source words the decoder attends to when generating each target word.

The natural question: if attention lets you access any position directly, why keep the LSTM at all?

"Attention Is All You Need"

Vaswani et al. (2017) proved that attention alone — with no recurrence at all — could match or beat the LSTM-plus-attention state of the art on machine translation. More importantly, that attention-only design could be fully parallelized on a GPU.

Scaled dot-product self-attention is the mechanism that makes it work: every token computes a compatibility score against every other token, then blends their values by that score.

import torch
import torch.nn.functional as F
import math

def scaled_dot_product_attention(Q, K, V, mask=None):
    """
    Q: [batch, heads, seq_len, d_k]  — queries
    K: [batch, heads, seq_len, d_k]  — keys  
    V: [batch, heads, seq_len, d_v]  — values
    """
    d_k = Q.size(-1)
    
    # Compute attention scores: how compatible is each query with each key?
    scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(d_k)
    # scores shape: [batch, heads, seq_len, seq_len]
    
    # For causal (decoder) attention: mask future positions
    if mask is not None:
        scores = scores.masked_fill(mask == 0, -1e9)
    
    # Convert to probabilities
    weights = F.softmax(scores, dim=-1)
    
    # Weighted sum of values
    output = torch.matmul(weights, V)
    # output shape: [batch, heads, seq_len, d_v]
    
    return output, weights

In self-attention, the same sequence produces queries, keys, and values. Every token can attend to every other token in one shot — no sequential dependency. The entire operation is a batch of matrix multiplications, which GPUs execute in parallel.

Multi-Head Attention

A single attention head learns one kind of relationship between tokens; multi-head attention runs several heads in parallel, each with its own projections. It's the difference between one editor reading a draft and a team of specialist editors — grammar, tone, facts — reading it simultaneously.

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, num_heads):
        super().__init__()
        assert d_model % num_heads == 0
        
        self.d_k = d_model // num_heads
        self.num_heads = num_heads
        
        # Single projection matrices for all heads
        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, d_model)
        self.W_v = nn.Linear(d_model, d_model)
        self.W_o = nn.Linear(d_model, d_model)
    
    def split_heads(self, x, batch_size):
        # [batch, seq, d_model] → [batch, heads, seq, d_k]
        x = x.view(batch_size, -1, self.num_heads, self.d_k)
        return x.transpose(1, 2)
    
    def forward(self, x, mask=None):
        batch_size, seq_len, _ = x.shape
        
        Q = self.split_heads(self.W_q(x), batch_size)
        K = self.split_heads(self.W_k(x), batch_size)
        V = self.split_heads(self.W_v(x), batch_size)
        
        attn_output, _ = scaled_dot_product_attention(Q, K, V, mask)
        
        # Recombine heads: [batch, heads, seq, d_k] → [batch, seq, d_model]
        attn_output = attn_output.transpose(1, 2).contiguous()
        attn_output = attn_output.view(batch_size, seq_len, -1)
        
        return self.W_o(attn_output)

Different heads learn different relationships. In a trained model, one head might capture syntactic dependencies (subject-verb agreement), another semantic similarity, another coreference (pronouns pointing to nouns). The heads operate independently and their outputs are concatenated.

Positional Encoding: Giving Transformers a Sense of Order

Self-attention on its own has no sense of word order — shuffle the input tokens and the output shuffles right along with them. Positional encoding injects that missing order information back in, the way page numbers restore sequence to a stack of shuffled pages.

class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=5000, dropout=0.1):
        super().__init__()
        self.dropout = nn.Dropout(dropout)
        
        # Precompute sinusoidal encodings
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len).unsqueeze(1).float()
        
        div_term = torch.exp(
            torch.arange(0, d_model, 2).float() * 
            (-math.log(10000.0) / d_model)
        )
        
        pe[:, 0::2] = torch.sin(position * div_term)  # even dims: sin
        pe[:, 1::2] = torch.cos(position * div_term)  # odd dims: cos
        
        self.register_buffer('pe', pe.unsqueeze(0))   # [1, max_len, d_model]
    
    def forward(self, x):
        # x: [batch, seq_len, d_model]
        x = x + self.pe[:, :x.size(1)]
        return self.dropout(x)

The sinusoidal encoding is elegant: the relative position of two tokens can be expressed as a linear function of the encodings, which helps the model generalize to sequence lengths not seen during training. Modern models (GPT-NeoX, LLaMA) use learned relative position encodings like RoPE instead.

Architecture Comparison

LSTMs propagate a hidden state step by step through gates; transformers connect every position to every other position through attention, computed all at once. The diagram below lines up both data flows side by side.

Head-to-Head Comparison

The two architectures trade off in almost every dimension that matters — sequential versus parallel, fixed memory versus quadratic memory, streaming speed versus training speed.

NLP Benchmark Evolution

GLUE (General Language Understanding Evaluation) scores show exactly how fast transformers overtook LSTMs — within a single year, not a decade.

The jump from ELMo (the best LSTM-based model) to BERT happened in the same year, 2018. Within 12 months, BERT had exceeded the LSTM frontier by a margin that would have taken years to close incrementally.

When to Use Each Architecture Today

Use an LSTM when:

• Streaming inference is required — processing one token at a time as it arrives, with no ability to look ahead.
• Hardware has strict memory limits — microcontrollers and edge devices where transformer memory demands are prohibitive.
• Sequences are extremely long — where O(n²) attention becomes computationally infeasible.
• The dataset is small and compute is limited — too little data or budget to justify fine-tuning a large transformer.

Use a Transformer when:

• Building any modern NLP application — classification, generation, translation, QA all default to transformers today.
• Fine-tuning a pretrained model — BERT, T5, or the GPT family, all built on this architecture.
• Working with multimodal data — vision-language models are transformer-based almost without exception.
• Scale matters — transformers scale further with both data and compute than LSTMs ever did.

For practical NLP projects, the advice is simple: start with a pretrained transformer — BERT for classification, GPT-2 for generation. LSTMs are worth understanding for the theory; in most production systems, they've been replaced.

State Space Models: The Next Chapter

State space models (SSMs), most notably Mamba (Gu & Dao, 2023), are a 2023-2024 development that brings back LSTM-like sequential structure but designed for fast parallel training. They show performance competitive with transformers at lower computational cost on long sequences.

The field keeps moving. Understanding both LSTMs and transformers gives you the foundation to follow wherever it goes next.

For deeper context on how transformers power modern language models, see the transformer architecture notes and LLM concepts. The LLM Learning section covers the application of transformers in large language models.

Test your understanding with the Deep Learning Quiz, and the ML Basics Quiz covers the foundations that both architectures build on.

The Machine Learning course has hands-on sequence modeling projects, and the Embeddings and Vector Database notes explain how transformer representations are used in retrieval systems.