Linear Algebra Intuition for AI Engineers

An AI model cannot operate directly on a sentence, photograph, customer, or sound. It first represents that object with numbers. Linear algebra gives us the language for arranging those numbers, comparing them, and transforming them into increasingly useful representations.

By the end of this lesson, you will be able to:

• explain vectors, matrices, and tensors without relying on a memorized definition;
• read common AI tensor shapes and predict output shapes;
• interpret a dot product as alignment and a matrix as a transformation;
• connect y = xW + b to embeddings, neural networks, and attention;
• implement a small linear layer and semantic search calculation;
• identify common shape, broadcasting, and similarity mistakes.

From real objects to vectors#

A scalar is one number: a learning rate, loss, probability, or temperature. A vector is an ordered list of numbers. The order matters because every position participates in a shared representation.

Suppose a simple house-price model represents a property as:

x = [area_m², bedrooms, age_years]

The vector [120, 3, 8] is not merely a list. It is one point in a feature space whose axes have agreed meanings. An embedding uses the same structure, except its learned dimensions usually do not have simple labels such as “bedrooms.” Meaning is distributed across many dimensions.

"database replication" → [0.17, -0.42, 0.83, ..., 0.09]
"copying data"         → [0.14, -0.38, 0.79, ..., 0.12]
"banana bread"         → [-0.61, 0.22, -0.08, ..., 0.47]

If training produces a useful space, related objects occupy related directions or neighborhoods. This geometric idea powers embedding retrieval, recommendations, clustering, classifiers, and many parts of language models.

Coordinates are not the object

The same physical point has different coordinates under different axes. Likewise, an embedding is not the word or document itself. It is a representation produced by a particular model, version, and normalization policy.

This creates an important engineering rule: vectors from unrelated embedding models do not automatically share a meaningful coordinate system. Re-embedding only queries while leaving documents in an older space can silently destroy retrieval quality even if every vector has the same length.

Direction, magnitude, and the dot product#

For a vector x = [x_1, x_2, ..., x_n], the Euclidean or L2 norm measures length:

The dot product multiplies corresponding coordinates and adds them:

For a = [1, 2, 3] and b = [4, 5, 6]:

Geometrically, the same operation is:

The dot product is large when vectors have large magnitudes and point in similar directions. It is zero when they are perpendicular, and negative when they point in opposing directions.

Cosine similarity divides away magnitude:

“Use cosine for embeddings” is a useful starting point, not a universal law. Some embedding models are trained for dot-product retrieval; others expect normalized vectors, where dot product and cosine similarity become equal. Follow the model’s documented similarity function and store that choice with the index configuration.

A matrix is a transformation#

A matrix is a rectangular arrangement of numbers, but that description hides its most useful intuition. Treat a matrix as a function that transforms vectors.

Consider:

Multiplying x = [1, 2] by A produces [2, 1]: the horizontal coordinate doubles while the vertical coordinate halves. Other matrices can rotate, reflect, shear, project, or combine dimensions.

The diagonal example changes each coordinate independently. A dense matrix is more interesting because each output coordinate can combine every input coordinate:

This is why a learned matrix can create new features. An output dimension does not merely copy one input; it learns a weighted combination.

Shape is a contract

For matrix multiplication:

[m, n] @ [n, p] → [m, p]

The inner dimensions must match, and the outer dimensions survive. Suppose a batch contains 32 examples, each represented by 768 numbers, and a layer produces 3,072 features:

X: [32, 768]
W: [768, 3072]
Y = X @ W: [32, 3072]

Do not reach for a transpose until you can state what every axis means. A transpose that makes code run can still reverse the intended semantics.

The operation at the center of neural networks#

A linear layer computes:

The matrix W mixes and transforms input features. The bias b shifts the result. For a batch, the bias vector is broadcast across every row.

import numpy as np

rng = np.random.default_rng(7)

X = rng.normal(size=(4, 3))       # 4 examples, 3 input features
W = rng.normal(size=(3, 2))       # 3 inputs → 2 output features
b = np.zeros(2)                   # one bias per output feature

Y = X @ W + b

print("X:", X.shape)
print("W:", W.shape)
print("b:", b.shape)
print("Y:", Y.shape)

Expected shapes:

X: (4, 3)
W: (3, 2)
b: (2,)
Y: (4, 2)

Notice that a stack containing only linear layers is still one linear transformation. Neural networks insert nonlinear activation functions between layers so they can represent relationships that no single matrix can capture.

“Linear layer” is library shorthand for an affine map

A mathematical linear map f must preserve both addition and scalar multiplication:

The map f(x)=xW satisfies those rules and sends the zero vector to zero. Once we add a nonzero bias, f(x)=xW+b no longer does: f(0)=b. It is an affine map—a linear transformation followed by a translation.

Deep-learning libraries nevertheless call this module a “linear layer.” That name is conventional, but the distinction matters when reasoning about geometry. The matrix controls how directions are mixed, stretched, or projected; the bias moves the resulting coordinate system's origin.

The PyTorch equivalent

import torch
from torch import nn

torch.manual_seed(7)

X = torch.randn(4, 3)
layer = nn.Linear(in_features=3, out_features=2)
Y = layer(X)

print(layer.weight.shape)  # [2, 3]
print(layer.bias.shape)    # [2]
print(Y.shape)             # [4, 2]

PyTorch stores nn.Linear.weight as [out_features, in_features], so its implementation is conceptually X @ weight.T + bias. Library storage layout and the mathematical notation can differ; shapes reveal the truth.

From matrices to tensors#

A tensor is a multidimensional array. In practical AI engineering, “tensor” usually means an array carrying both numbers and metadata such as dtype, device, shape, strides, and gradient history.

Common language-model shapes are:

Here B is batch size, T sequence length, C hidden width, H number of heads, D head dimension, and V vocabulary size.

Matrix multiplication generalizes across leading batch dimensions. If X has shape [B, T, C] and W has shape [C, 4C], then X @ W has shape [B, T, 4C]. The same learned transformation is applied to every token in every batch item.

Linear algebra inside attention#

Self-attention begins with three learned projections:

Q = X @ Wq
K = X @ Wk
V = X @ Wv

Queries and keys are compared with dot products:

Softmax converts scores into weights, and another matrix multiplication combines the values:

token representations X
   ├── @ Wq ──> queries Q ──┐
   ├── @ Wk ──> keys K ─────┼──> QKᵀ / √d ──> softmax ──┐
   └── @ Wv ──> values V ────────────────────────────────┴──> output

The important idea is learned routing: a query uses alignment with keys to decide which value vectors to combine. Later lessons will derive the shapes, masking, stable softmax, multi-head layout, and efficient kernels.

Build a tiny semantic search system#

The following lab uses deliberately small, hand-written embeddings so every operation remains visible. Real applications obtain embeddings from a model and batch the calculation.

import numpy as np

documents = [
    "Neural networks learn useful representations",
    "A database index speeds up retrieval",
    "Gradient descent updates model parameters",
    "A cache reduces repeated database reads",
]

embeddings = np.array([
    [0.95, 0.80, 0.10],
    [0.10, 0.20, 0.95],
    [0.90, 0.85, 0.05],
    [0.05, 0.30, 0.90],
], dtype=np.float64)

query = np.array([0.92, 0.82, 0.08])

def normalize_rows(matrix):
    norms = np.linalg.norm(matrix, axis=1, keepdims=True)
    return matrix / np.maximum(norms, 1e-12)

normalized_documents = normalize_rows(embeddings)
normalized_query = query / max(np.linalg.norm(query), 1e-12)

scores = normalized_documents @ normalized_query
ranking = np.argsort(scores)[::-1]

for index in ranking:
    print(f"{scores[index]:.3f}  {documents[index]}")

The epsilon policy prevents division by zero, but a production system should decide what a zero embedding means instead of quietly accepting it. Also record the embedding model, version, dimension, distance metric, and normalization policy with the index.

Common mistakes and debugging rules#

Confusing elementwise and matrix multiplication

In NumPy and PyTorch, A * B is elementwise multiplication while A @ B is matrix multiplication. They can sometimes produce the same shape while representing entirely different computations.

Treating broadcasting as free correctness

Broadcasting aligns dimensions from the right. Adding [B, T, C] and [C] applies one feature bias everywhere, which is often intended. Adding [B, T, C] and [B, 1, 1] applies one scalar per batch item. Both run; only the problem definition says which is correct.

Comparing unnormalized vectors accidentally

Raw dot product can rank a large-magnitude but weakly aligned vector above a smaller, better-aligned one. Normalize when the model and similarity contract call for cosine similarity.

Computing an inverse to solve a system

Avoid inv(A) @ b in numerical code. Use solve(A, b) or a suitable factorization. It is usually faster and more stable.

Ignoring dtype and device

Shape-correct code can still fail because tensors use different dtypes or devices. Mixed precision can also overflow or underflow. Treat shape, dtype, device, and numerical range as one tensor contract.

Memory chart#

Object → vector representation
Matrix → learned transformation
Dot product → alignment plus magnitude
Cosine similarity → alignment without magnitude
Tensor → numbers + axes + dtype + device

[m, n] @ [n, p] → [m, p]
y = xW + b

Before an operation:
1. Name every axis.
2. Write every shape.
3. Predict the result.
4. Check dtype, device, and range.

If you can explain why a matrix transforms a representation, predict the output of a multiplication, and implement normalized similarity without copying, you have the foundation needed for the next lesson: vectors, matrices, and their operations in greater depth.

References#

• Deisenroth, Faisal, and Ong: Mathematics for Machine Learning, Chapters 2–3 — vector spaces, linear and affine mappings, analytic geometry, and projections.
• Goodfellow, Bengio, and Courville: Deep Learning, Chapter 2 — the focused linear-algebra prerequisites used in deep learning.
• NumPy: matrix multiplication — broadcasting and matrix-product behavior.
• PyTorch: Linear — weight layout and affine transformation contract.
• PyTorch: broadcasting semantics — tensor compatibility rules.
• Vaswani et al.: Attention Is All You Need — scaled dot-product attention and learned projections.
• Raschka: LLMs from Scratch, Chapter 3 code — executable Q/K/V projections, scaled attention, masks, and batched attention shapes.