Hung-yi Lee ML 2025 Lecture 4: Beyond Transformers - Linear Attention, Mamba, and the RNN Renaissance

Notes from Professor Hung-yi Lee’s (李宏毅) course “Machine Learning in the Era of Generative AI (2025)” at National Taiwan University.

Why Every Architecture Exists

Professor Lee starts with a philosophical point: every neural network architecture exists for a reason, and those reasons differ.

CNN: A special case of fully-connected networks with receptive fields and parameter sharing. Exists to reduce unnecessary parameters for image processing, preventing overfitting.

Residual Connections: Exist to make optimization easier. Without them, deeper networks are harder to train - not because of overfitting, but because the loss landscape becomes rugged with many local minima. Residual connections smooth the error surface.

So when you see a new architecture, ask: what problem was it designed to solve?

The RNN Family

Before Self-Attention dominated, RNN-style architectures ruled sequence processing. The core idea: maintain a hidden state that accumulates information from all previous inputs.

The general RNN formula:

• $h_t = f_a(h_{t-1}) + f_b(x_t)$ (update hidden state)
• $y_t = f_c(h_t)$ (produce output)

Where:

• $h$ is the hidden state (can be a vector or matrix)
• $f_a$ is the “reflection” function (processes previous memory)
• $f_b$ is the “write” function (incorporates new input)
• $f_c$ is the “read” function (extracts output)

This mirrors the AI agent memory system from Lecture 2: write module, reflection module, read module.

LSTM and GRU are RNNs where $f_a$, $f_b$, $f_c$ vary with time (controlled by gates). This allows the model to decide what to remember and what to forget.

Why Self-Attention Won

The real advantage of Self-Attention isn’t “infinite memory” - that’s an illusion. The real advantage is parallelization during training.

RNN’s problem: To compute $h_6$, you must first compute $h_1, h_2, h_3, h_4, h_5$ sequentially. GPUs hate waiting.

Self-Attention’s advantage: Given a complete input sequence, you can compute all outputs $y_1$ through $y_T$ in parallel. Everything is matrix multiplication, which GPUs love.

During training, you have the complete target sentence. Self-Attention can process it all at once. RNN must process token by token.

The Inference Problem

But at inference time, the tables turn:

	Self-Attention	RNN
Computation	Grows with sequence length	Fixed per step
Memory	Grows with sequence length	Fixed (just hidden state)
Training	Highly parallelizable	Sequential, slow

As context windows grow (GPT-3.5: one Harry Potter chapter → Gemini 1.5: 2 million tokens), Self-Attention’s quadratic scaling becomes painful.

Linear Attention: The Bridge

Here’s the mind-blowing revelation: Linear Attention is just RNN without the reflection function.

Start with RNN, remove $f_a$ (reflection): $h_t = h_{t-1} + f_b(x_t)$

Expand this:

• $h_1 = f_b(x_1)$
• $h_2 = f_b(x_1) + f_b(x_2)$
• $h_T = f_b(x_1) + f_b(x_2) + ... + f_b(x_T)$

Now assume:

• $h$ is a $D \times D$ matrix
• $f_b(x_t) = v_t \cdot k_t^T$ (outer product of two vectors)
• Output: $y_t = h_t \cdot q_t$

Work through the math, and you get: $y_t = \sum_{i=1}^{t} (k_i^T \cdot q_t) \cdot v_i$

This is exactly Self-Attention… minus the softmax!

Linear Attention = Self-Attention without softmax = RNN without reflection

This was known since 2020 (paper: “Transformers are RNNs”). The implication: Linear Attention can be trained like a Transformer (parallel) but run inference like an RNN (efficient).

Why Softmax Matters

If Linear Attention is so similar, why does it underperform Self-Attention?

The memory limit myth: People think RNN has limited memory while Attention has infinite memory. But both have limits! In a D-dimensional space, you can only have D orthogonal vectors. With sequences longer than D, keys will inevitably have non-zero inner products with unintended queries.

The real difference: Softmax enables relative importance.

Without softmax: An attention weight of 1.0 always means “important.”

With softmax: Importance is relative. If a new item has attention weight 2.0, the old item with weight 1.0 becomes less important (0.45 → 0.17 after softmax).

Professor Lee’s analogy: In “Solo Leveling,” Iron seemed strong until Igris appeared. Igris seemed strong until Beru appeared. Importance is relative.

Linear Attention’s fatal flaw: Memory never changes. Once information enters, it stays forever. No forgetting mechanism.

Adding Forgetting: RetNet and Gated Retention

RetNet: Multiply $h_{t-1}$ by a constant $\gamma$ (between 0 and 1): $h_t = \gamma \cdot h_{t-1} + v_t \cdot k_t^T$

This causes gradual forgetting. Still parallelizable during training.

Gated Retention: Make $\gamma$ time-dependent: $\gamma_t = \sigma(x_t \cdot W_\gamma)$

Now the model learns when to forget. See a paragraph break? Set $\gamma_t \approx 0$ to clear memory.

More complex forgetting: Use a matrix $G_t = \mathbf{1} \cdot s_t^T$ for element-wise multiplication with $h_{t-1}$. This controls forgetting per memory column.

Mamba: The Famous Challenger

Mamba (late 2023) was the first Linear Attention variant to genuinely beat Transformers. It uses a complex state-space formulation, but Mamba2 simplified to essentially Gated Retention.

Key results:

• Slightly outperforms Transformer++ across model sizes
• Massive inference speedup (much higher tokens/second)

The famous meme: A giant stone statue (Mamba) towering over Transformers. Professor Lee notes this is actually from “Solo Leveling” - and realizes Meruem (the Chimera Ant King) is basically Beru with a different name.

Delta Net: Memory as Gradient Descent

Delta Net introduces a clever idea: before writing new information, first erase what’s already in that memory location.

The update rule can be rewritten as gradient descent: $h_t = h_{t-1} - \beta_k (h_{t-1} k_t - v_t) k_t^T$

This is gradient descent on the loss: $L = \frac{1}{2}\|h \cdot k_t - v_t\|^2$

The interpretation: Update memory so that querying with $k_t$ retrieves $v_t$ as accurately as possible.

This led to Titan (January 2025): “Learning to Memorize at Test Time” - treating memory updates as learned gradient descent.

Current State of the Field

Large-scale validation:

• JAMBA: Mamba-based, up to 52B parameters
• MiniMax-01: Linear Attention, 400B+ parameters

Image generation: SANA uses Linear Attention for fast inference.

But not everywhere: “MambaOut” paper shows Mamba doesn’t help image classification - you don’t need long-range attention for that task.

Practical approach: Instead of training from scratch, fine-tune existing models (like LLaMA) by swapping Self-Attention for Mamba layers.

The Bet

There’s an actual bet on whether “Attention is All You Need” will still hold in 2027:

• For: Jonathan Frankle (Harvard, Mosaic ML)
• Against: Sasha Rush (Cornell, HuggingFace)

656 days until we find out.

Key Takeaways

1. Every architecture exists for a reason - understand the “why”
2. Self-Attention won because of training parallelization, not infinite memory
3. Linear Attention = Self-Attention - softmax = RNN - reflection
4. Softmax enables relative importance and implicit forgetting
5. Mamba and variants add learnable forgetting to Linear Attention
6. The field is converging: train like Transformer, infer like RNN