Training neural networks to characterize turbulence

Preamble: Run the cells below to import the necessary Python packages

Training a model¶

• Given a model $f(\boldsymbol{x}; \boldsymbol{\theta})$ and a dataset $\{(\boldsymbol{x}_i, y_i)\}$, we want to find the parameters $\boldsymbol{\theta}$ that best describe the data.

• Linear regression had an analytic equation for the weight matrix $\boldsymbol{\theta}$, turns out to minimize the least-squares cost function $J(\boldsymbol{\theta}) = \frac{1}{2} \sum_i (y_i - f(\boldsymbol{x}_i; \boldsymbol{\theta}))^2$.

• Usually we aren’t so lucky as to have an exactly solvable convex optimization problem

• The next best thing is to have an optimization problem (possibly non-convex) where the loss is a smooth, differentiable function of the weights. In this case, we can use gradient descent and its variants

An aside: What if our problem is non-differentiable?¶

• Brute force (try random parameters or grid search and see what gives the best result).

• Genetic algorithms, reinforcement learning

• Can attempt to use gradient descent, if we can construct a differentiable surrogate function for the true loss

Predicting the Reynolds number of a turbulent flow¶

Our dataset will consist of snapshots of the velocity field of a turbulent flow. Each snapshot is labelled with the corresponding Reynolds number, which is a dimensionless number that characterizes the flow regime of a fluid. The higher the Reynolds number, the more turbulent the flow. Examples of low Reynolds number flows are the flow of honey around a spoon, or the motion of water around a swimming bacterium. Examples of high Reynolds number flows are the motion of air around an aircraft wing, or the flow created by the flapping of a bird’s wing.

Because we are going to use large models for this problem, we will downsample the dataset substantially in order to ensure that we can train the model in a reasonable amount of time.

We will first implement a data loader that will fetch the dataset from a remote repository, perform standardization and downsampling, and then split the data into training and testing sets to help prevent overfitting.

We can see that each of our data points is a flattened snapshot of the vorticity field of a turbulent flow, $\mathbf{x} \in \mathbb{R}^{D}$ with $D = 64^2 = 4096$ pixels. Each data point is labelled with the corresponding Reynolds number, $y \in \mathbb{R}$. We can first inspect our dataset qualitatively by plotting a few examples of each Reynolds number.

Our problem consists of a supervised regression problem: given a snapshot of the velocity field, predict the Reynolds number of the flow.

where $\mathbf{x} \in \mathbb{R}^{D}$ is the flattened vorticity field, and $\hat{y} \in \mathbb{R}$ is the predicted Reynolds number.

Before fitting a model, check the data¶

We first perform an unsupervised embedding. We ignore the labels $y$ and inspect only the data $\mathbf{x}$. We can start by visualizing the full training dataset in a lower-dimensional space using principal component analysis (PCA).

We see that the data is not linearly separable, so this is likely a non-trivial learning problem. We can also try using a non-linear embedding, such as t-SNE or UMAP, to see if either one better separates the data.

We can see that the low Reynolds number dataset ($Re = 300$) cleanly separates from the higher Reynolds number datasets. This tracks with our qualitative inspection of the dataset above, where the vortex street is clear at low Reynolds numbers, but quickly becomes turbulent at higher Reynolds numbers.

We next check the balance of the dataset by inspecting the training labels $y$, and ignoring the data $\mathbf{x}$.

Generate baseline predictions¶

Since this is a regression problem, we will generate baseline predictions using a linear regression model. We suspect that the lack of linear separability of the data will lead to poor performance, but this helps us gain some intuition for what a “good” score should look like.

We will use the sklearn implementation of ridge regression. Instead of manually tuning the regularization parameter, we will use cross-validation to find the best value. sklearn has a built-in model that will perform this automatically, RidgeCV

Let’s try a neural network¶

Recall that a multilayer perceptron extends linear regression by adding non-linear activation functions between the linear transformation of the input and the output.

where $\sigma$ is an elementwise function (such as the logistic function), $\boldsymbol{\theta}_1$ and $\boldsymbol{\theta}_2$ are trainable weight matrices, and $\mathbf{x}$ is a single input datapoint.

For our problem, $\mathbf{x} \in \mathbb{R}^{D}$, $D = 64^2 = 4096$ is the flattened velocity field, and $\hat{\mathbf{y}} \in \mathbb{R}$ is the predicted Reynolds number. While the output is a scalar here, this isn’t a requirement of multilayer perceptron. The shapes of the weight matrices are $\boldsymbol{\theta}_1 \in \mathbb{R}^{M \times D}$ and $\boldsymbol{\theta}_2 \in \mathbb{R}^{1 \times M}$. While the input and output shapes are fixed by the problem, the number of neurons $M$ is a hyperparameter. This is sometimes called the “width” of the network, or the “number of hidden units” of the network.

We will use scikit-learn’s own implementation of a multi-layer perceptron (MLP). This will handle fitting the values of the weight matrices $\boldsymbol{\theta}_1$ and $\boldsymbol{\theta}_2$ that minimize the model’s prediction across the training set. Behind the scenes, it will handle the forward and backward passes, and the updates to the weight matrices. The objective function that this model minimizes is the mean-squared error between the predicted and true labels across the training set:

where $N$ is the number of datapoints in the training set, $\hat{y}_i$ is the model’s predicted Reynolds number for the $i^{th}$ datapoint (which depends on the current values of the weight matrices $\boldsymbol{\theta}_1$ and $\boldsymbol{\theta}_2$), and $y_i$ is the true Reynolds number.

We can see that the constructor specifies the network architecture and several hyperparameters that affect the training process:

• The “Adam” optimizer is a variant of gradient descent that uses adaptive learning rates and momentum. This is used to determine the updates to the weight matrices.

• The “relu” activation function is a non-linear function $\sigma(x) = \max(0, x)$ that is more robust to vanishing gradients than the logistic function

• The “batch size” is the number of samples used to compute the gradient at each step. This is a compromise between using the full dataset (which is slow) and using a single sample (which is noisy).

• The “epoch” is the number of times the entire dataset is used to compute the gradient.

What about other supervised learning models?¶

There are many supervised learning models besides variants of linear regression and neural networks. For example, random forests and their generalization, gradient boosted forests, use a different set of models based on ensembling to perform classification and regression.

How do we train a multilayer perceptron?¶

Given a model

and an objective function

where the sum indexes into the $N$ training datapoints. Each training point $\mathbf{x}$ has a known true label $y$. We can update the weights with gradient descent. For our two-layer network, our relevant gradients are

where $\eta$ is the learning rate. Gradient descent operates elementwise on the weights. Fortunately, the gradient of a scalar objective with respect to a tensor of a given shape is a tensor of the same shape. Given a training set with $N$ examples, we usually take the average of the gradient computed over the $N$ examples.

Taking gradients defines a computational graph¶

Let’s start by revisiting how we optimize in low dimensions. Suppose we wanted to take the gradient of the double well potential

We can find the gradient analytically by taking the derivative

We implicitly performed several subcomputations to do this. Let’s write out these functions in terms of basic computational steps.

Given an $x$, computing $V(x)$ (the forward evaluation, or forward pass) can be broken up into

Now, given $V$, we want to compute $\frac{\partial V}{\partial x}$ (the backward evaluation, or backward pass). We can break this up into

Backpropagation through a multilayer perceptron¶

Forward pass¶

When a trained neural network is given an input $\mathbf{x}$, it computes an output $\hat{\mathbf{y}}$ by passing the input through a series of layers. Each layer is a linear transformation followed by a non-linear function. For example, a two-layer network is given by

for concreteness, we will use the sigmoid nonlinearity $\sigma(x) = \frac{1}{1 + e^{-x}}$. This model takes an input $\mathbf{x} \in \mathbb{R}^D$ and performs a matrix multiplication with $\boldsymbol{\theta}_1 \in \mathbb{R}^{H_1 \times D}$ to produce a hidden representation $\mathbb{R}^{H_1}$. The nonlinearity is then applied elementwise to produce a hidden representation with the same shape $\mathbf{h}_1 \in \mathbb{R}^{H_1}$. A matrix multiplication with $\boldsymbol{\theta}_2 \in \mathbb{R}^{1 \times H_1}$ is then used to contract the hidden representation into a single output, followed by a sigmoid nonlinearity to produce the final prediction $\hat{y} \in \mathbb{R}$.

We start by breaking this into single steps

Backward pass¶

Suppose we pass an input $\mathbf{x} \in \mathbb{R}^D$ and get an output $\hat{y} \in \mathbb{R}$, and we know the true output is $y \in \mathbb{R}$.

If we are using the mean-squared error loss function, then for a single training example $\mathbf{y}$, the loss $\mathcal{L} \in \mathbb{R}$ is

We want to update all the parameters $\boldsymbol{\theta}_1$ and $\boldsymbol{\theta}_2$ to minimize this loss. To calculate $\frac{\partial \mathcal{L}}{\partial \boldsymbol{\theta}_2} \in \mathbb{R}^{M \times H_2}$, we need to use the chain rule,

In a reverse-mode chain rule evaluation, we can solve this sequence from left-to-right. We start with the scalar term. Since the loss is a scalar, and both the prediction and the true label are scalars, the gradient is a scalar.

Next, we need to compute the gradient of the sigmoid nonlinearity with respect to its input. Since the sigmoid nonlinearity is applied elementwise to the scalar input $h_2^{u}$, this is a scalar as well. A useful identity for the sigmoid nonlinearity is $\sigma'(x) = \sigma(x) (1 - \sigma(x))$.

Finally, we need to compute the gradient of the linear transformation with respect to the weight matrix. Since the input to the linear transformation is a vector, the gradient is a row vector.

Putting this all together,

Notice how we actually traversed the network backwards, starting with the gradient of the last evaluation we performed during the forward pass. In order to update $\boldsymbol{\theta}_2$, we need to know both the output of the forward pass $\hat{y}$ and the intermediate hidden representation $\mathbf{h}_1$. Fortunately, we cached these values during the forward pass.

Backpropagation is the chain rule on a computation graph¶

Diagrams from UW CSE599W lecture notes

We can now compute the gradient for the first layer.

$$ \frac{\partial \mathcal{L}}{\partial \boldsymbol{\theta}_1} =

\frac{\partial \mathcal{L}}{\partial \hat{y}} \frac{\partial \hat{y}}{\partial h_2^{u}} \frac{\partial h_2^{u}}{\partial \mathbf{h}_1} \frac{\partial \mathbf{h}_1}{\partial \mathbf{h}_1^{u}} \frac{\partial \mathbf{h}_1^{u}}{\partial \boldsymbol{\theta}_1} $$

We can find the first few terms as above. For the term $\frac{\partial \mathbf{h}_1}{\partial \mathbf{h}_1^{u}}$, we are taking the gradient of an elementwise operation over a vector, and so this returns a vector.

where the product is taken elementwise.

Putting this all together,

where $\odot$ denotes the elementwise product. The parenthetical term corresponds to a column vector of shape $H_1$. Since $\mathbf{x}$ appears as row vector of shape $D$, the gradient is a matrix of shape $H_1 \times D$.

Notice that the gradient of $\boldsymbol{\theta}_1$ is the same shape as the parameter itself, and so gradient descent can be applied elementwise.

Some things to note when performing the backwards pass¶

• Several of the “intermediate” values we computed during the forward pass re-appear when we perform the backwards pass. Because we cached these values, then we don’t have to recompute them during the backward pass.

• Notice that when we compute the “forward” pass we start with the innermost portion of the function composition, and then work our way out. For the backwards pass, we start with the outermost portion of the function composition, and work our way in.

• Notice that the sigmoid function’s derivative $\sigma'(x) = \sigma(x) (1 - \sigma(x))$ can be written as a function purely of its output $\sigma(x)$. This is a common pattern in autodiff.

• In order to keep track of indicies and transposes, it helps to think about the dimensions of the input/output of each primitive operation.

• It’s useful to remember the some rules for matrix derivatives:

This improves the model, but it’s still not very good. Why does the prediction $\hat{y}$ plateau at 1.0? Let’s try implementing our model with a different activation function, the RELU.

This looks good, but how well does our trained model perform on unseen data?

Let’s modify our training loop to exploit that we have many datapoints. At each iteration, we pass multiple datapoints through the network, and compute the average gradient of the loss with respect to the weights. This is called mini-batch gradient descent. The size of the batch is a hyperparameter. Because we select a random batch of datapoints, the dynamics in weight space are stochastic.

A few technical details: we did not vectorize our code, and so we had to write a for loop for the forward and backwards passes.

A reminder regarding indices...¶

For a single datapoint (vector) $\mathbf{x} \in \mathbb{R}^{d}$, the forward pass of a multilayer perceptron is

We normally aggregate many datapoints into a matrix $\mathbf{X} \in \mathbb{R}^{N \times d}$, where $N$ is the number of datapoints. In this case, the forward pass returns $\hat{\mathbf{Y}} \in \mathbb{R}^{N \times 1}$, where each row is the prediction for a single datapoint. For some regression problems (such as the fluid forecasting problem on the homework), $\hat{\mathbf{Y}} \in \mathbb{R}^{N \times M}$

The first index, the data or “batch” index, is usually vectorized: it passes right through the model, since every data point has an output. The first index usually recieves special treatment in high-level packages like PyTorch and TensorFlow, since it is often used to aggregate gradients from multiple datapoints.

The loss function is usually a function of the entire dataset, including the batch index. For example, the mean squared loss is

For a multivariate prediction problem, the loss function is

where $\hat{\mathbf{y}}_i \in \mathbb{R}^{M}$ and $\mathbf{y}_i \in \mathbb{R}^{M}$.

Importantly, when we compute gradients of the loss, we treat the entire dataset $X$, the true labels $Y$, and even the current predicted labels $\hat{Y}$ as a fixed set of constant parameters

Vector-valued functions¶

We can see that we were able to simplify the calculation by breaking the forward pass into discrete steps that got combined, and that the backward pass was able to reuse some of the intermediate calculations. What if we go to higher dimensions? Recall that our neural network has a forward pass that looks like

where $\sigma$ is the non-linear activation function applied elementwise to the output of matrix multiplication. We combine a given prediction on a single datapoint $\hat{\mathbf{y}}$ with the true label $\mathbf{y}$ to compute the mean-squared error loss

In order to perform gradient descent, we need to compute the gradient of the loss with respect to the weights $\boldsymbol{\theta}_1$ and $\boldsymbol{\theta}_2$. While the loss is a scalar, the gradient is a vector.

Some useful identities¶

Linear model¶

We saw that linear regression and ridge regression have an analytic solution for the optimal weights. This is not the case for neural networks. Instead, we need to use gradient descent. It turns out that we can also frame linear regression as a neural network, and use the same gradient descent algorithm to train it. This is useful for large datasets, where pseudo-inverse methods are too slow.

A linear model is given by

where $\mathbf{x} \in \mathbb{R}^D$ is a datapoint, $\boldsymbol{\theta} \in \mathbb{R}^D$ is a vector of weights, and $\hat{y} \in \mathbb{R}$ is the predicted output. Writing out the model in terms of the components of $\mathbf{x}$ and $\boldsymbol{\theta}$ gives

For linear regression, we typically use the mean-squared loss function,

where $i$ indexes into different data points, and $j$ indexes the features.

Taking the gradient of the loss with respect to each component of $\boldsymbol{\theta}$ gives, for a single datapoint $\mathbf{x}_i$,

where $i$ indexes the data points and $j$ indexes the features.

Logistic model:¶

The model is

Instead of mean-squared error, logistic models are often trained using cross-entropy loss:

where $i$ indexes the data points and $j = 1, 2, ..., M$ indexes the features. $\frac{\partial \mathcal{L}}{\partial \theta_j} \in \mathbb{R}^M$ is the gradient of the loss with respect to the weights.

Multivariate Linear model¶

A multivariate linear model is given by

where $\mathbf{x}$ is a $N \times D$ matrix with $N$ data points and $D$ features, the weight matrix $\boldsymbol{\theta} \in \mathbb{R}^{M \times D}$ , and the prediction vector is $\hat{\mathbf{y}} \in \mathbb{R}^M$. Usually, we work with the data matrix $X = \{ \mathbf{x}_i \} \in \mathbb{R}^{N \times D}$.

The mean-squared loss function is

The gradient is

Notice that the gradient is a matrix $\frac{\partial \mathcal{L}}{\partial \boldsymbol{\theta}} \in \mathbb{R}^{M \times D}$, where each row is the gradient of the loss with respect to the weights for a single output. This means that we have update values separately for each element of the weight matrix. Also, note that the right hand side is an outer product, while the forward pass is an inner product.

Harder cases: branching paths, loops, and shared weights¶

• For more sophisticated architectures, we might have multiple paths through the network. For example, we might have a residual connection, where the output of the first layer is added to the output of the second layer.

• Backpropagation is more complicated, but as long as everything is deterministic and continuous, it should still be possible to implement all gradients at nodes.

Diagram from Szegedy et al. (2014)

Math is hard: Automatic Differentiation¶

• Computing symbolic gradients of vector functions is exact but laborious

• Finite-difference gradients are approximate, easy, but expensive in high dimensions:

• • If the input to our network is $D$-dimensional, then we need to compute $2\,M$ gradients. For a 1 Megapixel image, this is a million gradient operations per gradient descent step

Automatic differentiation is a technique that can compute the exact gradient of a function in a computationally-efficient manner.

Backpropagation is just reverse-mode autodiff applied to train deep neural networks.

Let the computer do the work¶

• Most existing deep learning frameworks are built on top of autodiff

• Tensorflow, PyTorch, JAX, and many others allow you to specify the neural network as a function composition graph, and can compute the gradient automatically as long as everything is differentiable

• Implictly, these networks build a computation “graph” that is later traversed backwards to compute the gradient

• Caching forward pass values can speed up the backwards pass, since many derivatives depend on forward pass values

Some optimization tricks¶

• Instead of “full batch” gradient descent, where we compute the average of the gradient $\nabla_\theta \mathcal{L}(\theta)$ over all training examples $x_i$, we can use stochastic gradient descent, where we randomly sample a subset of the training data during each gradient descent epoch. Smaller batch sizes are more noisy, while larger batch sizes are more accurate to the global geometry of the loss function.

• Instead of just gradient descent, we can use momentum and many other tricks we saw in optimization:

  • Stochastic

  • Momentum

  • Second-order methods (e.g. Newton’s method)

  • Adaptive learning rates

  • Constrained Optimization (e.g. L1/L2 regularization)

• Currently, many practioners use the Adam optimizer, which is a combination of momentum and weight-based learning rate adjustment. Second-order methods are less commonly used for large problems, due to the expense of finding the full Hessian, but approximates methods based on its eigenvalues are viable.

How do different hyperparameters affect training?¶

Two common hyperparameters that we’ll encounter when training large networks are the learning rate $\eta$ and the batch size $B$.

• Batch size $B$: The number of training examples used to compute the gradient. Larger batch sizes are better estimators of the overall gradient of the loss landscape, but more expensive to compute. Smaller batch sizes are more noisy, but can be faster to compute. The stochasticity of small batch sizes can also help the network avoid local minima.

• Learning rate: The size of the step taken in the direction of the gradient. If the learning rate is too small, then the network will take a long time to converge. If the learning rate is too large, then the network will oscillate around the minimum, or even diverge.

• We’ll demonstrate these hyperparameters using scikit-learn’s own MLPRegressor class, which is a multilayer perceptron for regression problems. We’ll use the adam iterative optimization algorithm, which is a combination of momentum and adaptive learning rates.

We can see how changing the batch size affects the training process.