Jonathan Lalou's Blog

Modern artificial intelligence workloads—particularly those associated with deep learning—have reshaped the way computation is structured and executed. While CPUs remain indispensable for general-purpose tasks, GPUs have become the de facto standard for training and running machine learning models.

This shift is not incidental. It is driven by a deep alignment between the mathematical structure of AI and the architectural characteristics of GPUs. In this article, we examine this alignment and illustrate it with representative code commonly found in real-world AI systems.

The Computational Nature of AI

At its core, modern machine learning is an exercise in large-scale numerical optimization. Whether training a convolutional network or a transformer, the dominant operations are:

• Matrix multiplications
• Tensor contractions
• Element-wise transformations
• Non-linear activations

These operations are instances of linear algebra applied at scale. Crucially, they exhibit a high degree of data parallelism: the same operation is applied repeatedly across large datasets.

From Mathematical Abstraction to Code

To understand why GPUs excel, it is instructive to look at how AI code is written in practice.

Example 1: A Simple Neural Network Layer (PyTorch)

import torch
import torch.nn as nn

# Define a simple linear layer
layer = nn.Linear(in_features=1024, out_features=512)

# Simulated batch of input data
x = torch.randn(64, 1024)  # batch size = 64

# Forward pass
y = layer(x)

The operation above is fundamentally a matrix multiplication followed by a bias addition. Each output element is computed independently, making the workload inherently parallel.

Example 2: Training Step in a Neural Network

import torch
import torch.nn as nn
import torch.optim as optim

model = nn.Sequential(
    nn.Linear(1024, 512),
    nn.ReLU(),
    nn.Linear(512, 10)
)

optimizer = optim.Adam(model.parameters())
criterion = nn.CrossEntropyLoss()

# Dummy input and labels
inputs = torch.randn(64, 1024)
targets = torch.randint(0, 10, (64,))

# Forward pass
outputs = model(inputs)

# Compute loss
loss = criterion(outputs, targets)

# Backward pass
loss.backward()

# Update weights
optimizer.step()
optimizer.zero_grad()

Both the forward and backward passes are dominated by tensor operations applied across entire batches, reinforcing the highly parallel nature of AI workloads.

Example 3: Convolutional Operation (Core of CNNs)

import torch
import torch.nn as nn

conv = nn.Conv2d(in_channels=3, out_channels=64, kernel_size=3, padding=1)

# Batch of images: (batch_size, channels, height, width)
images = torch.randn(32, 3, 224, 224)

# Apply convolution
features = conv(images)

Convolutions apply the same kernel across spatial dimensions, resulting in a massive number of independent computations—ideal for parallel execution.

Example 4: Attention Mechanism (Transformer Core)

import torch
import torch.nn.functional as F

def attention(Q, K, V):
    scores = Q @ K.transpose(-2, -1)
    scores = scores / (Q.size(-1) ** 0.5)
    weights = F.softmax(scores, dim=-1)
    return weights @ V

# Simulated query, key, value tensors
Q = torch.randn(32, 8, 128, 64)  # batch, heads, seq_len, dim
K = torch.randn(32, 8, 128, 64)
V = torch.randn(32, 8, 128, 64)

output = attention(Q, K, V)

This pattern—matrix multiplication followed by normalization and weighted aggregation—is central to modern transformer architectures and exemplifies the computational intensity of AI workloads.

Architectural Alignment

A clear pattern emerges from these examples:

• Uniform operations applied across large tensors
• Minimal branching or complex control flow
• Heavy reliance on linear algebra primitives

These characteristics align closely with GPU design, which emphasizes throughput and parallel execution.

Memory Throughput and Data Movement

AI workloads are not only compute-intensive but also data-intensive. Large tensors must be moved efficiently between memory and compute units. GPUs provide significantly higher memory bandwidth than CPUs, enabling sustained performance for such operations.

The Role of Frameworks

Modern frameworks abstract away hardware complexity while exposing high-level primitives such as tensor operations and automatic differentiation. This allows developers to write expressive code while leveraging specialized hardware.

Conclusion

The preference for GPUs in AI is a consequence of structural compatibility between workload and architecture. AI code is inherently parallel, tensor-centric, and dominated by linear algebra operations.

GPUs are designed precisely to execute such workloads efficiently at scale. For software developers, understanding this alignment is essential to building performant and scalable machine learning systems.

Further Exploration

• Computational graphs and automatic differentiation
• Transformer architectures
• Mixed-precision training
• GPU execution models