Computational Analysis of Transformers

Chapter Overview

Understanding computational requirements is crucial for deploying transformers. This chapter analyzes time and space complexity, memory footprints, and inference costs. We derive exact FLOP counts, memory requirements, and scaling laws for transformers of different sizes.

Learning Objectives

1. Calculate FLOPs for transformer forward and backward passes
2. Analyze memory requirements for training and inference
3. Understand scaling laws for model size, data, and compute
4. Optimize inference through batching and caching
5. Estimate training time and costs for large models

Computational Complexity

Understanding the computational complexity of transformers is essential for making informed decisions about model architecture, hardware requirements, and deployment strategies. The transformer's computational profile differs fundamentally from recurrent architectures, trading sequential dependencies for quadratic memory scaling—a trade-off that profoundly impacts both training and inference.

Self-Attention Complexity

Self-attention is the defining operation of transformers, and its computational characteristics determine much of the model's behavior. For a sequence of length $n$ with model dimension $d_{\text{model}}$, we analyze each component of the attention mechanism in detail.

QKV Projections: The first step projects the input $\mX \in \R^{n \times d_{\text{model}}}$ into query, key, and value spaces. Each projection is a matrix multiplication:

Each matrix multiplication $\mX \mW$ requires $2nd_{\text{model}}d_k$ floating-point operations (FLOPs): for each of $n \times d_k$ output elements, we perform $d_{\text{model}}$ multiply-add operations. With three projections:

For the common case where $d_k = d_{\text{model}}$ (single-head or considering all heads together):

Why this matters for hardware: These are dense matrix multiplications, which achieve high utilization on modern GPUs. NVIDIA A100 GPUs can perform up to 312 TFLOPS (FP16 with Tensor Cores), meaning these projections are typically compute-bound rather than memory-bound. However, for small batch sizes or short sequences, the operations may become memory-bandwidth limited, achieving only 10-20\% of peak FLOPS.

Attention Score Computation: Computing $\mS = \mQ \mK\transpose$ involves multiplying $\mQ \in \R^{n \times d_k}$ by $\mK\transpose \in \R^{d_k \times n}$, producing $\mS \in \R^{n \times n}$:

The attention matrix $\mS$ has $n^2$ elements, and computing each requires $d_k$ multiply-add operations. This quadratic scaling in sequence length is the fundamental bottleneck for long-context transformers.

Dimension tracking example: For BERT-base with $n=512$, $d_k=64$ (per head), and $h=12$ heads:

Across 12 heads, we compute 12 separate $512 \times 512$ attention matrices, requiring:

Hardware implications: The attention matrix requires $n^2$ memory per head. For $n=512$ and 12 heads with FP32:

This seems modest, but for $n=2048$ (GPT-2): $12 \times 2048^2 \times 4 = 201\text{ MB}$ per sequence. With batch size 32: $6.4\text{ GB}$ just for attention matrices! This is why long-context models require substantial GPU memory.

Softmax and Scaling: Applying softmax to each row of $\mS$ requires $O(n^2)$ operations (exponentials and normalization), which is negligible compared to the matrix multiplications but can become significant for very long sequences due to memory access patterns.

Attention Output: Computing $\mO = \mA \mV$ multiplies the attention weights $\mA \in \R^{n \times n}$ by values $\mV \in \R^{n \times d_v}$:

Again, this scales quadratically with sequence length. For $d_v = d_k$:

Output Projection: Finally, concatenated head outputs are projected back to model dimension:

Total Self-Attention FLOPs:

For typical configurations where $d_k = d_v = d_{\text{model}}/h$:

Complexity regime analysis: The relative importance of the two terms depends on the ratio $n/d_{\text{model}}$:

• Short sequences ($n \ll d_{\text{model}}$): The $8nd_{\text{model}}^2$ term dominates. For BERT-base with $n=128$, $d=768$: $8 \times 128 \times 768^2 \approx 603\text{M}$ vs $4 \times 128^2 \times 768 \approx 50\text{M}$. The projections dominate.
• Long sequences ($n \gg d_{\text{model}}$): The $4n^2d_{\text{model}}$ term dominates. For $n=8192$, $d=768$: $8 \times 8192 \times 768^2 \approx 38.7\text{G}$ vs $4 \times 8192^2 \times 768 \approx 206\text{G}$. The attention computation dominates.
• Crossover point: When $8nd_{\text{model}}^2 \approx 4n^2d_{\text{model}}$, solving gives $n \approx 2d_{\text{model}}$. For $d=768$, this occurs around $n \approx 1536$.

This analysis explains why efficient attention mechanisms (Chapter 16) focus on reducing the $O(n^2)$ term for long-context applications.

Feed-Forward Network Complexity

The position-wise feed-forward network (FFN) in each transformer layer typically expands the representation to a higher dimension before projecting back. This two-layer network with GELU or ReLU activation is applied independently to each position in the sequence.

Architecture: For input $\mX \in \R^{n \times d_{\text{model}}}$:

The intermediate dimension $d_{ff}$ is typically $4d_{\text{model}}$ in standard transformers (BERT, GPT), though some models use different ratios. This expansion allows the network to learn complex non-linear transformations.

First Projection FLOPs: Computing $\mX \mW_1$ requires:

For $d_{ff} = 4d_{\text{model}}$:

Second Projection FLOPs: Computing $\mH \mW_2$ requires:

Total FFN FLOPs:

Activation function: GELU requires additional operations (exponentials, multiplications) but these are $O(nd_{ff})$, negligible compared to the matrix multiplications.

Why FFN dominates computation: Comparing FFN to self-attention:

For typical sequence lengths where $n < 2d_{\text{model}}$, the FFN requires roughly twice the FLOPs of attention! This is why some efficient transformer variants (e.g., mixture-of-experts) focus on making the FFN more efficient.

Memory and bandwidth considerations: The FFN intermediate activations $\mH \in \R^{n \times d_{ff}}$ must be stored for backpropagation. For BERT-base with $n=512$, $d_{ff}=3072$:

With 12 layers and batch size 32: $6 \times 12 \times 32 = 2.3\text{ GB}$ just for FFN intermediate activations. This is a significant portion of training memory.

Hardware utilization: FFN matrix multiplications are highly regular and achieve excellent GPU utilization (often 70-90\% of peak FLOPS on modern GPUs). The operations are:

• Compute-bound for reasonable batch sizes and sequence lengths
• Well-suited for Tensor Cores on NVIDIA GPUs (FP16/BF16 operations)
• Easily parallelizable across the sequence dimension

On an NVIDIA A100 GPU (312 TFLOPS FP16), computing the FFN for BERT-base with batch size 32 and $n=512$:

Per-Layer Total Complexity

Combining self-attention and FFN, a complete transformer layer requires:

Additional operations: Layer normalization, residual connections, and dropout add $O(nd_{\text{model}})$ operations, which are negligible compared to the matrix multiplications.

Breakdown by component:

• FFN: $16nd_{\text{model}}^2$ (typically 60-70\% of layer FLOPs for short sequences)
• Attention projections: $8nd_{\text{model}}^2$ (typically 25-35\%)
• Attention computation: $4n^2d_{\text{model}}$ (grows with sequence length)

This breakdown is crucial for optimization: for short sequences, optimizing FFN yields the largest gains; for long sequences, efficient attention mechanisms become critical.

Complexity Analysis

Comparison with RNN:

• RNN: $O(Lnd^2)$ time, $O(Ld^2)$ space
• Transformer: Quadratic in $n$ but parallel; RNN sequential

Bottleneck regimes:

• Short sequences $(n < d)$: FFN dominates, $O(Lnd^2)$
• Long sequences $(n > d)$: Attention dominates, $O(Ln^2d)$

Memory Requirements

Memory is often the limiting factor in training and deploying large transformer models. Understanding memory requirements at a granular level enables informed decisions about model architecture, batch sizes, and hardware selection. We analyze memory consumption across four categories: model parameters, gradients, optimizer states, and activations.

Model Parameters

Model parameters must be stored in GPU memory during both training and inference. The memory footprint depends on the numerical precision used.

Precision options:

• FP32 (float32): 4 bytes per parameter, standard precision
• FP16 (float16): 2 bytes per parameter, half precision
• BF16 (bfloat16): 2 bytes per parameter, better range than FP16
• INT8: 1 byte per parameter, quantized inference

For BERT-base with 110 million parameters:

Parameter breakdown for BERT-base:

Per transformer layer:

12 layers:

Total BERT-base:

In FP32: $110\text{M} \times 4 = 440\text{ MB}$

Larger models scale dramatically:

GPT-3 in FP32 requires 700 GB just for parameters—far exceeding single GPU memory (A100 has 80 GB). This necessitates:

• Model parallelism: Split model across multiple GPUs
• Mixed precision: Use FP16/BF16 (350 GB for GPT-3)
• Quantization: INT8 inference (175 GB for GPT-3)

Activation Memory

During training, intermediate activations must be stored for backpropagation. Activation memory scales with batch size and sequence length, often dominating memory consumption.

Activations per transformer layer:

• Input to layer: $B \times n \times d_{\text{model}}$
• Query, Key, Value: $3 \times B \times n \times d_{\text{model}}$
• Attention scores: $B \times h \times n \times n$ (quadratic in sequence length!)
• Attention output: $B \times n \times d_{\text{model}}$
• FFN intermediate: $B \times n \times d_{ff}$
• Layer norm activations: $2 \times B \times n \times d_{\text{model}}$

Total activation memory per layer (approximate):

For BERT-base ($B=32$, $n=512$, $d_{\text{model}}=768$, $h=12$, $d_{ff}=3072$):

For 12 layers: $12 \times 1\text{ GB} = 12\text{ GB}$ just for activations!

Impact of sequence length: The attention matrix term $B \times h \times n^2$ grows quadratically. For $n=2048$ (4× longer):

For 12 layers: $77\text{ GB}$ just for attention matrices—nearly filling an A100 GPU!

This quadratic scaling is why:

• Long-context models require gradient checkpointing (recompute activations during backward pass)
• Efficient attention mechanisms (sparse, linear) are crucial for long sequences
• Batch sizes must be reduced for longer sequences

Gradient checkpointing trade-off: Recomputing activations during backward pass:

• Memory savings: Reduce activation memory by $\sim$80\%
• Compute cost: Increase training time by $\sim$20-30\%
• When to use: When memory-constrained, especially for long sequences

Training Memory Budget

Training requires memory for parameters, gradients, optimizer states, and activations. Understanding this breakdown is essential for selecting hardware and configuring training.

Total training memory:

For AdamW optimizer (most common for transformers):

• Model parameters: $P$ parameters $\times$ 4 bytes (FP32) = $4P$ bytes
• Gradients: $P$ parameters $\times$ 4 bytes = $4P$ bytes
• First moment (momentum): $P$ parameters $\times$ 4 bytes = $4P$ bytes
• Second moment (variance): $P$ parameters $\times$ 4 bytes = $4P$ bytes
• Activations: $A$ bytes (depends on batch size, sequence length, model depth)

Total: $16P + A$ bytes

Mixed precision training (FP16/BF16 with FP32 master weights):

• FP16 parameters (forward/backward): $2P$ bytes
• FP32 master parameters: $4P$ bytes
• FP32 gradients: $4P$ bytes
• FP32 optimizer states: $8P$ bytes
• FP16 activations: $A/2$ bytes

Total: $18P + A/2$ bytes

Surprisingly, mixed precision uses slightly MORE memory for parameters/optimizer (18P vs 16P) but saves significantly on activations ($A/2$ vs $A$). Since activations often dominate, mixed precision typically reduces total memory.

Hardware Selection Guide

GPU memory requirements by model size (mixed precision + gradient checkpointing):

Inference memory requirements (FP16):

• Parameters only: $2P$ bytes
• KV cache (autoregressive): $2 \times L \times h \times n_{\max} \times d_k \times B$ bytes
• Activations (single forward pass): Minimal compared to training

For GPT-2 (117M params) inference:

GPT-2 inference easily fits on consumer GPUs or even CPUs!

Inference Optimization

Inference optimization is critical for deploying transformers in production. Unlike training, which prioritizes throughput (tokens/second across large batches), inference prioritizes latency (time to generate a single response) while maintaining reasonable throughput. We analyze key optimization techniques and their trade-offs.

KV Caching for Autoregressive Decoding

Autoregressive generation (used in GPT, decoder-only models) generates tokens sequentially, where each new token attends to all previous tokens. Naive implementation recomputes attention for all previous positions at each step—highly inefficient.

Problem analysis: Generating sequence of length $T$ tokens:

• Step 1: Compute attention for position 1 (attends to position 1)
• Step 2: Compute attention for position 2 (attends to positions 1-2)
• Step 3: Compute attention for position 3 (attends to positions 1-3)
• Step $T$: Compute attention for position $T$ (attends to positions 1-$T$)

Total attention computations: $\sum_{t=1}^{T} t = \frac{T(T+1)}{2} \approx \frac{T^2}{2}$

For $T=1000$ tokens: $\approx 500{,}000$ attention computations!

KV Caching solution: Key and value projections depend only on input tokens, not on the query position. Cache $\mK$ and $\mV$ from previous steps:

Computational savings: With caching, each step computes attention once (not recomputing previous positions):

Speedup: For $T=1000$: $\frac{500{,}000}{1{,}000} = 500\times$ faster!

Memory cost: Store keys and values for all positions and layers:

For GPT-2 ($L=12$, $h=12$, $d_k=64$, FP16):

Memory scaling with sequence length:

For GPT-3, the cache reaches 9.7~GB per sequence (see Chapter~[ref] for the derivation), nearly filling an A100 for batch size 8 ($77.6$~GB). Practitioners manage this through smaller batch sizes for long contexts, dynamic batching, and INT8 quantization to reduce cache size by 2--4$\times$.

Batched Inference

Processing multiple sequences simultaneously increases GPU utilization and throughput.

Single sequence inference: For GPT-2 generating 100 tokens:

• Compute: $\approx 100 \times 8\text{ GFLOPs} = 800\text{ GFLOPs}$
• Time on A100: $\frac{800\text{ GFLOPS}}{312\text{ TFLOPS} \times 0.3} \approx 8.5\text{ ms}$
• GPU utilization: $\approx 30\%$ (memory-bound, not compute-bound)

Batched inference (batch size 32):

• Compute: $32 \times 800\text{ GFLOPs} = 25{,}600\text{ GFLOPs}$
• Time on A100: $\frac{25{,}600\text{ GFLOPS}}{312\text{ TFLOPS} \times 0.7} \approx 117\text{ ms}$
• GPU utilization: $\approx 70\%$ (much better!)
• Throughput: $\frac{32 \times 100}{117\text{ ms}} \approx 27{,}350\text{ tokens/sec}$

Latency vs. throughput trade-off:

• Batch size 1: Latency = 8.5 ms, throughput = 11,765 tokens/sec
• Batch size 32: Latency = 117 ms (13.8× worse), throughput = 27,350 tokens/sec (2.3× better)

Padding challenge: Sequences in a batch must have the same length. Shorter sequences are padded, wasting computation:

• Sequence lengths: [512, 256, 128, 64]
• Padded to: [512, 512, 512, 512]
• Wasted computation: $(512-256) + (512-128) + (512-64) = 1024$ positions (50\%!)

Solutions:

• Dynamic batching: Group sequences of similar lengths
• Bucket batching: Pre-defined length buckets (128, 256, 512, 1024)
• Packed sequences: Concatenate sequences without padding (requires careful attention masking)

Quantization for Inference

Quantization reduces memory and increases throughput by using lower-precision arithmetic.

Precision options:

• FP32: 4 bytes, full precision
• FP16/BF16: 2 bytes, half precision (1.5-2× speedup)
• INT8: 1 byte, 8-bit integer (2-4× speedup, 4× memory reduction)
• INT4: 0.5 bytes, 4-bit integer (4-8× speedup, 8× memory reduction)

INT8 quantization: Map FP32 weights $w \in [-w_{\max}, w_{\max}]$ to INT8 $w_q \in [-128, 127]$:

Dequantize during computation:

Quantization impact on GPT-2:

Accuracy trade-offs:

• FP16/BF16: Negligible accuracy loss (<0.1\% perplexity increase)
• INT8: Small accuracy loss (0.5-2\% perplexity increase) with calibration
• INT4: Moderate accuracy loss (2-5\% perplexity increase), requires careful quantization

Hardware support:

• NVIDIA Tensor Cores: Accelerate FP16/BF16 (up to 2× speedup)
• NVIDIA INT8 Tensor Cores: Accelerate INT8 (up to 4× speedup)
• CPU AVX-512 VNNI: Accelerate INT8 on CPUs

Model Distillation

Train smaller "student" model to mimic larger "teacher" model:

• DistilBERT: 66M params (vs. BERT-base 110M), 97\% performance, 2× faster
• TinyBERT: 14M params, 96\% performance, 7× faster

Distillation enables deployment on resource-constrained devices (mobile, edge).

Inference Optimization Summary

Combined optimizations: KV caching + FP16 + batching + INT8 can achieve 1000× speedup with minimal accuracy loss!

Scaling Laws

Performance scales as power laws with model size $N$, dataset size $D$, and compute budget $C$. The Kaplan scaling laws established that larger models are more sample-efficient, while the Chinchilla scaling laws (Hoffmann et al., 2022) refined the optimal allocation: $N_{\text{opt}} \propto C^{0.5}$ and $D_{\text{opt}} \propto C^{0.5}$, implying that many large models are over-parameterized and under-trained. See Section~[ref] for the full treatment including GPT-3 vs Chinchilla analysis and practical implications for training strategy.

Exercises

Solutions

Full solutions for all exercises are available at \url{https://deeplearning.hofkensvermeulen.be}.

chapters/chapter12_computational_analysis.texin{solution} Exercise 3: KV Caching Implementation

import torch
import torch.nn as nn
import time
import matplotlib.pyplot as plt

class GPT2WithCache(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.transformer = nn.ModuleList([
            GPT2Block(config) for _ in range(config.n_layer)
        ])
        self.ln_f = nn.LayerNorm(config.n_embd)
        self.lm_head = nn.Linear(config.n_embd, config.vocab_size)
    
    def forward(self, input_ids, past_key_values=None, use_cache=False):
        hidden_states = self.wte(input_ids) + self.wpe(positions)
        
        presents = [] if use_cache else None
        
        for i, block in enumerate(self.transformer):
            past = past_key_values[i] if past_key_values else None
            hidden_states, present = block(
                hidden_states, 
                past_key_value=past,
                use_cache=use_cache
            )
            if use_cache:
                presents.append(present)
        
        hidden_states = self.ln_f(hidden_states)
        logits = self.lm_head(hidden_states)
        
        return logits, presents

class GPT2Block(nn.Module):
    def forward(self, x, past_key_value=None, use_cache=False):
        # Self-attention with optional KV cache
        attn_output, present = self.attn(
            x, 
            past_key_value=past_key_value,
            use_cache=use_cache
        )
        x = x + attn_output
        x = self.ln_1(x)
        
        # Feed-forward
        x = x + self.mlp(x)
        x = self.ln_2(x)
        
        return x, present

class GPT2Attention(nn.Module):
    def forward(self, x, past_key_value=None, use_cache=False):
        B, T, C = x.shape
        
        # Compute Q, K, V
        q, k, v = self.c_attn(x).split(self.n_embd, dim=2)
        
        # Reshape for multi-head attention
        q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
        k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
        v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
        
        # Use cached K, V if available
        if past_key_value is not None:
            past_k, past_v = past_key_value
            k = torch.cat([past_k, k], dim=2)
            v = torch.cat([past_v, v], dim=2)
        
        # Store K, V for next iteration
        present = (k, v) if use_cache else None
        
        # Attention computation
        attn = (q @ k.transpose(-2, -1)) / math.sqrt(k.size(-1))
        attn = F.softmax(attn, dim=-1)
        out = attn @ v
        
        # Reshape and project
        out = out.transpose(1, 2).contiguous().view(B, T, C)
        out = self.c_proj(out)
        
        return out, present

Benchmarking Code:

def generate_without_cache(model, prompt, max_length=256):
    """Generate tokens without KV caching"""
    input_ids = prompt
    times = []
    
    for _ in range(max_length):
        start = time.time()
        logits, _ = model(input_ids, use_cache=False)
        times.append(time.time() - start)
        
        next_token = logits[:, -1, :].argmax(dim=-1, keepdim=True)
        input_ids = torch.cat([input_ids, next_token], dim=1)
    
    return input_ids, times

def generate_with_cache(model, prompt, max_length=256):
    """Generate tokens with KV caching"""
    input_ids = prompt
    past_key_values = None
    times = []
    
    for i in range(max_length):
        start = time.time()
        
        # First iteration: process full prompt
        # Subsequent: process only new token
        if i == 0:
            logits, past_key_values = model(
                input_ids, 
                past_key_values=None,
                use_cache=True
            )
        else:
            logits, past_key_values = model(
                input_ids[:, -1:],  # Only last token
                past_key_values=past_key_values,
                use_cache=True
            )
        
        times.append(time.time() - start)
        
        next_token = logits[:, -1, :].argmax(dim=-1, keepdim=True)
        input_ids = torch.cat([input_ids, next_token], dim=1)
    
    return input_ids, times

# Run benchmark
prompt = torch.randint(0, 50257, (1, 50))  # 50 token prompt

output_no_cache, times_no_cache = generate_without_cache(
    model, prompt, max_length=256
)
output_with_cache, times_with_cache = generate_with_cache(
    model, prompt, max_length=256
)

print(f"Without cache: {sum(times_no_cache):.2f}s")
print(f"With cache: {sum(times_with_cache):.2f}s")
print(f"Speedup: {sum(times_no_cache)/sum(times_with_cache):.2f}x")

Experimental Results:

For GPT-2 small (124M parameters), generating 256 tokens:

Generation Time vs Sequence Length:

# Benchmark different sequence lengths
seq_lengths = [32, 64, 128, 256, 512]
times_no_cache = []
times_with_cache = []

for length in seq_lengths:
    _, t_no = generate_without_cache(model, prompt, max_length=length)
    _, t_with = generate_with_cache(model, prompt, max_length=length)
    times_no_cache.append(sum(t_no))
    times_with_cache.append(sum(t_with))

# Plot results
plt.figure(figsize=(10, 6))
plt.plot(seq_lengths, times_no_cache, 'o-', label='Without cache', linewidth=2)
plt.plot(seq_lengths, times_with_cache, 's-', label='With cache', linewidth=2)
plt.xlabel('Sequence Length')
plt.ylabel('Generation Time (seconds)')
plt.title('KV Caching Impact on Generation Speed')
plt.legend()
plt.grid(True)
plt.savefig('kv_cache_speedup.png')

Analysis:

Without caching, time complexity: $O(n^2)$ where $n$ is sequence length $$T_{\text{no cache}} = \sum_{i=1}^n c \cdot i = c \cdot \frac{n(n+1)}{2} \approx O(n^2)$$

With caching, time complexity: $O(n)$ $$T_{\text{cache}} = c \cdot n$$

Speedup grows with sequence length: $$\text{Speedup}(n) = \frac{n(n+1)/2}{n} = \frac{n+1}{2} \approx O(n)$$

For $n=256$: Speedup $\approx 128/2 = 64\times$ (theoretical)

Actual speedup (16.2$\times$) is lower due to:

• Memory bandwidth bottleneck (loading cached K, V)
• Overhead of cache management
• Other non-attention computations (FFN, embeddings)

Memory Cost:

KV cache size: $2 \times L \times B \times n \times d = 2 \times 12 \times 1 \times 256 \times 768 = 4.7$MB

Small memory cost for massive speedup makes KV caching essential for inference.