Matrix Calculus and Tensors

Introduction

Modern artificial intelligence relies on computations involving far more than vectors and two-dimensional matrices. An image may be represented by height × width × colour-channel arrays; a video by time × height × width × channels; and a batch of language-model activations by batch × sequence length × embedding dimension. These multidimensional arrays are tensors. Matrix calculus provides the language for differentiating functions involving vectors, matrices, and tensors, while the chain rule explains how gradients flow through computational graphs during backpropagation.

The central practical objective is not merely to manipulate symbols. It is to understand shapes, dependencies, and derivative conventions well enough to implement reliable machine-learning models and diagnose errors. A correct derivative with an incorrect shape is still an unusable derivative.

Learning Objectives

By the end of this chapter, you should be able to:

1. Distinguish scalars, vectors, matrices, and higher-order tensors.

2. Perform tensor indexing, reshaping, contraction, and broadcasting safely.

3. Interpret Einstein summation notation.

4. Compute gradients, Jacobians, Hessians, and selected matrix derivatives.

5. Apply the multivariable chain rule to computational graphs.

6. Explain how these ideas produce reverse-mode automatic differentiation and backpropagation.

7. Identify numerical, implementation, and governance risks associated with tensor-based AI systems.

1. Scalars, Vectors, Matrices, and Tensors

A scalar has zero dimensions and is usually denoted by x or c. A vector has one index, for example x_i, and a matrix has two indices, for example A_ij. A tensor of order or rank r has r indices:

T_{i_1 i_2 ... i_r}.

Here, “rank” means the number of axes, not matrix rank. A tensor’s shape specifies the size of each axis. For example, an RGB image of height H and width W can be represented as X ∈ R^(H×W×3). A mini-batch of B such images has shape B×H×W×3.

An axis is a dimension of an array, and an index selects an entry along an axis. The tensor order and shape must be distinguished: a tensor of shape 5×5×3 has order three, whereas a matrix of shape 25×3 has order two, even if both contain 75 numbers.

Storage layout and meaning

The mathematical tensor does not prescribe how values are stored in memory. Software may use row-major or column-major layouts, and frameworks may use channel-first (C×H×W) or channel-last (H×W×C) conventions. A transpose, reshape, or permutation is safe only when the intended index correspondence is known. Confusing H×W×C with C×H×W can produce code that runs but gives incorrect predictions.

2. Fundamental Tensor Operations

Elementwise operations

If A and B have identical shape, elementwise addition is defined by

C_{ij}=A_{ij}+B_{ij}.

Multiplication is similarly C_{ij}=A_{ij}B_{ij}; this is not matrix multiplication. In array libraries, the symbol * commonly denotes elementwise multiplication, whereas a matrix product uses a separate operation.

Reduction and contraction

A reduction removes an axis by aggregating over it. For a vector x ∈ R^n,

s = Σ_i x_i.

A tensor contraction sums over one or more pairs of matching indices. Matrix multiplication is a contraction:

C_{ik}=Σ_j A_{ij}B_{jk}.

The repeated index j is summed, while i and k remain as output indices. Contraction is the key operation behind dense layers, attention mechanisms, and many scientific models.

Transpose and permutation

For a matrix, (A^T)_{ij}=A_{ji}. For higher-order tensors, axes can be permuted. If X has shape H×W×C, a permutation may produce Y with shape C×H×W, where

Y_{c i j}=X_{i j c}.

A permutation changes the interpretation of indices, not the numerical values themselves.

Reshaping and flattening

Reshaping changes the shape while preserving a specified ordering of entries. Flattening an H×W×C image into a vector of length HWC is common before a dense layer. However, flattening is not a mathematical substitute for a permutation: if the axes are ordered incorrectly, the resulting vector connects features to weights in the wrong way.

Broadcasting

Broadcasting permits an array with fewer or singleton dimensions to participate in an elementwise operation. For example, if X has shape B×d and b has shape d, then

Y_{bi}=X_{bi}+b_i.

The vector b is conceptually copied across the batch axis. If b has shape B×1, it instead supplies one value per example and is copied across features. Broadcasting is valid when dimensions are equal or one of them is 1, under the framework’s alignment rule, usually from the trailing axes.

Broadcasting affects derivatives. For Y_{bi}=X_{bi}+b_i and scalar loss L,

∂L/∂b_i = Σ_b ∂L/∂Y_{bi}.

The gradient with respect to a broadcast parameter is reduced over every axis along which that parameter was replicated.

3. Einstein Notation

Einstein notation suppresses explicit summation signs. A repeated index appearing exactly twice in a term is summed, while free indices appear in the output.

For example,

c_i = A_{ij}b_j

means c_i=Σ_j A_{ij}b_j. Matrix multiplication is

C_{ik}=A_{ij}B_{jk}.

The index j is contracted. The expression

s=A_{ij}B_{ij}

is a Frobenius inner product, equal to Σ_iΣ_j A_{ij}B_{ij}. In contrast, A_{ij}B_{ik} leaves j and k free and produces a matrix.

Einstein notation is useful because it makes dimensions explicit. Consider a linear layer:

Y_{bi}=X_{bj}W_{ji}+b_i.

The indices b, i, and j have clear meanings: b is the batch example, j is the input feature, and i is the output feature. The index j is summed. This notation also provides a direct bridge to optimized einsum implementations.

A valid Einstein expression must obey index rules. An index should not occur more than twice in a single product term, and free indices must be consistent across terms. Violations often indicate a shape or conceptual error.

4. Derivatives in Several Variables

Scalar derivatives and gradients

For a scalar function f:R→R, the derivative f'(x) describes local change. For f:R^n→R, the gradient is conventionally a column vector:

∇_x f = [∂f/∂x_1, ..., ∂f/∂x_n]^T.

The first-order approximation is

f(x+δx) ≈ f(x)+(∇_x f)^Tδx,

for sufficiently small δx. The gradient points in the direction of steepest increase under the Euclidean norm; therefore, -∇_x f is a local descent direction.

Example: let f(x)=1/2 x^T A x, where x∈R^n and A is constant. Expanding differentials gives

df = 1/2(dx^T A x+x^T A dx).

Since x^T A dx=(A^T x)^Tdx,

∇_x f = 1/2(A+A^T)x.

If A is symmetric, this simplifies to ∇_x f=Ax. Symmetry is an assumption, not a general property.

Jacobians

For f:R^n→R^m, the Jacobian is the m×n matrix

J_{ij}=∂f_i/∂x_j.

Its rows correspond to output components and columns to input components. Under this convention, a small perturbation satisfies

df ≈ J dx.

For a scalar-valued function, the Jacobian is a 1×n row, whereas the gradient is often represented as an n×1 column. They contain the same partial derivatives but differ in orientation. Consistency is essential when multiplying Jacobians.

Example: f(x)=Ax+c, where A∈R^(m×n). Then J=A.

Hessians

For f:R^n→R, the Hessian is the n×n matrix

H_{ij}=∂²f/(∂x_i∂x_j).

If the second partial derivatives are continuous in a neighbourhood, mixed partials agree and the Hessian is symmetric. The second-order approximation is

f(x+δx)≈f(x)+(∇f)^Tδx+1/2 δx^T H δx.

The Hessian describes curvature. Positive definite curvature indicates a strict local minimum; negative definite curvature indicates a strict local maximum; an indefinite Hessian indicates saddle-like curvature. In large neural networks, explicitly forming the Hessian is often too expensive, so Hessian-vector products are preferred.

5. Matrix Derivatives and Differential Calculus

Matrix calculus has multiple conventions. This chapter uses differentials and the trace to avoid ambiguity. For a scalar L, write

dL = tr((∇_X L)^T dX).

This defines the gradient ∇_X L to have the same shape as X.

Useful identities include:

• d tr(AX)=tr(A dX), so ∇_X tr(AX)=A^T.
• d||X||_F²=2 tr(X^T dX), so ∇_X ||X||_F²=2X.
• If L=1/2||X-Y||_F², then ∇_X L=X-Y.

For L=a^T X b, with fixed vectors a and b,

dL=a^T(dX)b=tr(ba^T dX),

hence ∇_X L=ab^T.

For L=tr(X^T A X),

∇_X L=(A+A^T)X.

If A is symmetric, this becomes 2AX. These formulas require compatible dimensions and assume A is constant with respect to X.

A particularly important result concerns a linear layer. Let

Y=XW+b,

where X has shape B×d, W has shape d×p, b has shape p, and Y has shape B×p. If an upstream gradient G=∂L/∂Y has shape B×p, then

∂L/∂X=GW^T,

∂L/∂W=X^T G,

∂L/∂b=Σ_b G_{b:}.

These expressions follow from dY=dX W+X dW+db and show why a bias gradient sums across the batch.

6. The Chain Rule and Computational Graphs

Suppose z=g(x) and y=f(z), with scalar y. The chain rule states

∂y/∂x_j = Σ_i (∂y/∂z_i)(∂z_i/∂x_j).

In matrix form, with column gradients,

∇_x y = J_g(x)^T ∇_z y.

The transpose appears because the Jacobian maps input perturbations forward, while reverse-mode differentiation transports gradients backward.

For a sequence x→u→v→L, one computes the forward values and then applies local derivative rules in reverse order. This is backpropagation. Reverse mode is efficient when there are many inputs and one scalar output, as in training a neural network loss with millions of parameters.

Worked example: a two-layer scalar network

Consider

z=w x+b,

a=σ(z),

L=1/2(a-t)^2,

where x, w, b, and t are scalars and σ is differentiable. The forward values are computed first. Let δ=L’s derivative with respect to a. Then

∂L/∂a=a-t,

∂a/∂z=σ'(z),

∂z/∂w=x,

∂z/∂b=1.

Therefore,

∂L/∂w=(a-t)σ'(z)x,

∂L/∂b=(a-t)σ'(z).

For the logistic sigmoid σ(z)=1/(1+e^{-z}), σ'(z)=a(1-a). The example illustrates that each local derivative is simple; the chain rule combines them into the parameter gradient.

Vector nonlinearities

If a nonlinearity acts elementwise, such as ReLU, its Jacobian is diagonal wherever the derivative exists. For ReLU(z_i)=max(0,z_i), the derivative is 1 for z_i>0 and 0 for z_i<0; at z_i=0 it is not classically differentiable, and implementations choose a subgradient convention. Consequently, an upstream gradient is multiplied elementwise by a mask.

7. Tensor Derivatives in Neural Networks

A tensor-to-tensor derivative can be represented as a high-order Jacobian, but explicitly storing it is usually impractical. For example, a mapping from an input tensor with shape H×W×C to an output tensor with the same shape has a derivative with six indices. Automatic differentiation instead computes vector-Jacobian products or Jacobian-vector products without materializing the full object.

If y=f(x) and an upstream cotangent g_y is available, reverse mode computes

g_x = J_f(x)^T g_y.

Forward mode computes J_f(x)v for a tangent vector v. Forward mode is attractive when the number of inputs is small; reverse mode is attractive when the number of outputs is small. Neural-network training usually uses reverse mode because the loss is scalar.

Convolution as tensor contraction

A two-dimensional convolution layer can be expressed schematically as

Y_{b i j o}=Σ_{r,s,c} X_{b,i+r,j+s,c} K_{r s c o}+b_o,

where b indexes examples, i and j spatial positions, c input channels, o output channels, and r,s kernel offsets. Padding and stride alter the valid index ranges and output shape. This equation shows that convolution is a structured tensor contraction with shared parameters. The sharing is why one kernel is applied at many spatial locations.

8. Worked Tensor Example: Broadcasting and Gradients

Let X have shape 2×3 and b have shape 3:

X = [[1, 2, 3], [4, 5, 6]], b=[10,20,30].

Define Y=X+b using broadcasting. Then

Y=[[11,22,33],[14,25,36]].

Let L=Σ_{b,i}Y_{bi}. Since each element of b contributes to two entries of Y,

∂L/∂b=[2,2,2].

If instead L=Σ_{b,i}Y_{bi}², then ∂L/∂Y=2Y and

∂L/∂b_i=Σ_b 2Y_{bi}.

Thus the gradient is [50,94,138], because the column sums of Y are [25,47,69]. This example demonstrates that broadcasting in the forward pass becomes summation in the reverse pass.

9. Applied Case: Diagnosing a Shape Error in an Image Classifier

Background

Consider an illustrative case involving a small image-classification service. The model accepts batches of images and uses a convolutional layer followed by a dense classifier. Training loss decreases, but validation accuracy remains near chance level.

Problem

The data pipeline stores images in channel-last form, B×H×W×C, while the convolution implementation expects channel-first form, B×C×H×W. The program automatically reshapes rather than explicitly permuting axes. No runtime error occurs because the total number of elements is unchanged.

Analysis

The tensor has the same number of values after reshaping, but the index meaning changes. A pixel’s red, green, and blue values are no longer grouped as channels at the expected locations. Consequently, local convolutional filters receive unrelated values. Gradients are mathematically computed for the implemented function, but that function is not the intended image model.

Approach and solution

The engineering team documents a canonical layout, performs an explicit axis permutation, and checks shapes at every model boundary. They add a synthetic test image in which only one channel and one location are nonzero. The expected output verifies that the permutation is correct. Gradient checking on a small model further confirms the implementation.

Outcome

After correcting the layout, training and validation behaviour become consistent, and accuracy improves substantially. The important point is that tensor correctness includes semantic axis meaning, not only element counts.

Learning points

1. Shape annotations should accompany every layer.

2. Reshape and transpose are different operations.

3. Unit tests should test index meaning with interpretable synthetic data.

4. Automatic differentiation cannot detect a logically incorrect tensor layout.

5. Reproducible pipelines require documented conventions for channels, batches, and sequence axes.

10. Numerical and Implementation Considerations

Gradient checking

For a scalar function f, a central finite-difference approximation is

∂f/∂x_i ≈ [f(x+h e_i)-f(x-h e_i)]/(2h),

where e_i is the i-th coordinate vector and h is small. This is useful for testing derivatives on small systems. It is not a replacement for automatic differentiation in production: finite differences are expensive and subject to cancellation and rounding errors. Relative and absolute tolerances should be selected according to scale and numerical precision.

Stability

Overflow, underflow, poorly scaled features, and saturated nonlinearities can produce unreliable gradients. Stable implementations of softmax and log-sum-exp subtract the maximum input before exponentiation. Normalization, suitable initialization, gradient clipping, and mixed-precision loss scaling can improve training, but each introduces assumptions and should be validated empirically.

Sparse and structured tensors

Many tensors contain repeated or zero structure. Sparse representations can reduce memory and computation, but their operations and gradients may have limited framework support. Structured derivatives, such as convolutional gradients, exploit parameter sharing rather than constructing dense Jacobians.

11. Limitations, Governance, and Professional Practice

Matrix calculus guarantees derivatives of the implemented mathematical function; it does not guarantee that the function is appropriate, fair, or safe. Tensor pipelines may encode sensitive personal data, and gradients can sometimes reveal information about training examples. Access control, data minimization, secure model interfaces, and privacy assessments are therefore relevant in professional deployments.

Engineers should document tensor schemas, units, axis conventions, derivative assumptions, precision, and treatment of nondifferentiable points. Model cards and data documentation can record intended use, limitations, and evaluation populations. Reproducibility requires fixed preprocessing, versioned libraries, deterministic tests where feasible, and independent validation of critical numerical routines.

Contemporary automatic-differentiation systems increasingly support compiler transformations, accelerator execution, sparse computation, distributed tensors, and higher-order derivatives. These developments improve scale but can make debugging more difficult. A sound practice is to validate a small transparent implementation before optimizing it with fusion, compilation, quantization, or distributed execution.

Best Practices

• Annotate every tensor with shape, axis meaning, and data type.
• Distinguish elementwise multiplication, matrix multiplication, and contraction.
• Use explicit permutations when changing axis order.
• Derive gradients symbolically for small modules before coding them.
• Check broadcast reductions in backward computations.
• Prefer Jacobian-vector or vector-Jacobian products to explicit Jacobian construction.
• Test gradients numerically on small, well-scaled examples.
• Treat nondifferentiable points and subgradient conventions explicitly.
• Monitor numerical precision and overflow in training pipelines.
• Validate both mathematical correctness and semantic correctness of data layouts.

Chapter Summary

Tensors generalize vectors and matrices to multidimensional arrays whose axes carry explicit meanings. Their essential operations include elementwise computation, broadcasting, permutation, reshaping, reduction, and contraction. Einstein notation provides a compact way to express these operations and exposes which indices are summed and which remain free.

For scalar functions, gradients describe first-order change and Hessians describe curvature. Jacobians represent derivatives of vector-valued functions. Matrix derivatives are most safely handled through differentials and trace identities, with gradients defined to match the shape of the differentiated variable. The multivariable chain rule connects local derivatives through a computational graph. Reverse-mode automatic differentiation applies transposed Jacobian actions backward and forms the basis of neural-network backpropagation.

In practice, tensor shape and axis semantics are as important as algebra. Broadcasting creates summations in backward passes, convolutions are structured tensor contractions, and explicit Jacobians are usually avoided in favour of Jacobian-vector or vector-Jacobian products. Reliable AI engineering additionally requires numerical testing, documented conventions, privacy awareness, and validation that the implemented tensor function is the intended one.

Key Terms

Scalar: A zero-order quantity with one value.

Vector: An ordered one-dimensional array.

Matrix: A two-dimensional array with row and column indices.

Tensor: A multidimensional array with one or more indexed axes.

Axis: A dimension of a tensor.

Shape: The ordered sizes of a tensor’s axes.

Broadcasting: Expanding singleton or missing dimensions for compatible elementwise operations.

Contraction: Summation over one or more pairs of tensor indices.

Einstein notation: Index notation in which repeated indices imply summation.

Gradient: The vector of partial derivatives of a scalar function.

Jacobian: The matrix of first derivatives of a vector-valued function.

Hessian: The matrix of second derivatives of a scalar function.

Differential: A first-order linear description of change.

Chain rule: A rule for differentiating compositions of functions.

Backpropagation: Reverse application of the chain rule through a computational graph.

Automatic differentiation: Programmatic evaluation of exact derivatives using elementary operations.

Vector-Jacobian product: A product of an upstream vector with a Jacobian, usually used in reverse mode.

Jacobian-vector product: A Jacobian multiplied by a tangent vector, usually used in forward mode.

Frobenius norm: The square root of the sum of squared tensor or matrix entries.

Axis permutation: Reordering the dimensions of a tensor.

Review Questions

Short Answer Questions

1. What is the difference between tensor order and tensor shape?

2. Why does a broadcast bias gradient require summation over the batch axis?

3. Define the Jacobian under the row-output, column-input convention.

4. When is the Hessian of a scalar function symmetric?

5. Why are explicit Jacobians usually avoided in neural-network implementations?

Long Answer Questions

1. Explain tensor contraction using Einstein notation and relate matrix multiplication to contraction.

2. Derive the gradients with respect to X, W, and b for Y=XW+b given an upstream gradient G.

3. Compare gradients, Jacobians, and Hessians, including their shapes and practical uses.

4. Explain reverse-mode automatic differentiation and why it is efficient for scalar loss functions.

5. Discuss why reshape and axis permutation must not be treated as interchangeable operations.

Analytical / Application Questions

1. Let X have shape 4×6 and b have shape 6, with Y=X+b and L=ΣY². Express ∂L/∂b in index notation and state its shape.

2. A model receives data in B×H×W×C format but expects B×C×H×W. Design two tests that verify the axis conversion semantically rather than merely checking the number of elements.

3. For f(x)=1/2 x^T A x, derive the gradient without assuming A is symmetric, then state the simplification when A is symmetric.

Multiple Choice Questions with Answers

Question 1

For C_{ik}=Σ_j A_{ij}B_{jk}, which index is contracted?

A. i

B. j

C. k

D. Both i and k

Correct Answer: B. j

Question 2

If X has shape B×d and b has shape d, what is the shape of X+b under valid broadcasting?

A. B×d

B. B×1

C. d×B

D. d×d

Correct Answer: A. B×d

Question 3

Under the convention J_{ij}=∂f_i/∂x_j, what is the shape of the Jacobian for f:R^n→R^m?

A. n×m

B. m×n

C. n×n

D. m×m

Correct Answer: B. m×n

Question 4

For L=1/2||X-Y||_F², what is ∇_X L?

A. X+Y

B. Y-X

C. X-Y

D. 2X

Correct Answer: C. X-Y

Question 5

If Y=XW+b and the upstream gradient G has shape B×p, what is ∂L/∂W?

A. GW^T

B. X^T G

C. G X^T

D. W^T G

Correct Answer: B. X^T G

Question 6

Which derivative is most appropriate for a scalar loss with millions of input parameters?

A. Reverse-mode differentiation

B. Forward-mode differentiation only

C. Explicit storage of the full Hessian

D. Finite differences for every parameter during training

Correct Answer: A. Reverse-mode differentiation

Question 7

For f(x)=1/2 x^T A x with arbitrary constant A, the gradient is:

A. Ax

B. A^T x

C. 1/2(A+A^T)x

D. (A-A^T)x

Correct Answer: C. 1/2(A+A^T)x

Question 8

In a broadcast addition Y_{bi}=X_{bi}+b_i, how is ∂L/∂b_i obtained?

A. By selecting only the first batch element

B. By summing ∂L/∂Y_{bi} over b

C. By multiplying by W

D. By transposing the batch dimension

Correct Answer: B. By summing ∂L/∂Y_{bi} over b

Question 9

What does a repeated index in Einstein notation generally indicate when it appears twice in a term?

A. It is retained as an output index

B. It is summed or contracted

C. It must equal zero

D. It denotes a batch dimension only

Correct Answer: B. It is summed or contracted

Question 10

Why can an incorrect image axis order produce poor predictions without a runtime shape error?

A. The tensor contains no numerical values

B. The total element count may remain unchanged while index meanings change

C. Gradients cannot be computed for images

D. Broadcasting always removes channels

Correct Answer: B. The total element count may remain unchanged while index meanings change