Welcome to The Nonlinear Library, where we use Text-to-Speech software to convert the best writing from the Rationalist and EA communities into audio. This is: What’s the backward-forward FLOP ratio for Neural Networks?, published by Marius Hobbhahn on December 13, 2021 on The AI Alignment Forum.
Summary:
Classic settings, i.e. deep networks with convolutional layers and large batch sizes, almost always have backward-forward FLOP ratios close to 2:1.
Depending on the following criteria we can encounter ratios between 1:1 and 3:1
Type of layer: Passes through linear layers have as many FLOP as they use to do weight updates. Convolutional layers have many more FLOP for passes than for weight updates. Therefore, in CNNs, FLOP for weight updates basically play no role.
Batch size: Weights are updated after the gradients of the batch have been aggregated. Thus, FLOP for passes increase with batch size but stay constant for weight updates.
Depth: The first layer has a backward-forward ratio of 1:1 while all others have 2:1. Therefore, the overall ratio is influenced by the fraction of FLOP in first vs. FLOP in other layers.
We assume the network is being optimized by stochastic gradient descent (w += ɑ⋅dw) and count the weight update as part of the backward pass. Other optimizers would imply different FLOP counts and could create ratios even larger than 3:1 for niche settings (see appendix B). However, the ratio of 2:1 in the classic setting (see point 1) should still hold even when you use momentum or Adam.
Compute-intensity of the weight update
Most compute-intensive layers
Backward-forward ratio
Large batch size OR compute-intensive convolutional layer
First layer
1:1
Other layers
2:1
Small batch size AND no compute-intensive convolutional layers
First layer
Other layers
3:1
Introduction:
As part of the Parameter, Compute and Data Trends in Machine Learning project we want to understand how many floating-point operations (FLOP) backward pass takes relative to a forward pass is for typical NN applications and which factors it depends on. We call this the backward-forward FLOP ratio.
In this post, we first provide a theoretical analysis of the ratio, and we then corroborate our findings empirically.
Theory:
To understand where the differences in ratios come from, we need to look at the classical equations of backpropagation.
Let’s start with a simple example---a neural network with 2 hidden layers.
In this example, we have the following computations for forward and backward pass assuming linear layers with ReLU activations. The “@”-symbols denote matrix multiplications.
Operation
Computation
FLOP forward
Computation
FLOP backward
δ1R=dL/dA2
=W2@δ2
dL/dW2
=δ2@A1R
We separate the weight update from the individual layers since the update is done after aggregation, i.e. we first add all gradients coming from different batches and then multiply with the learning rate.
From this table we see
ReLUs and the loss function contribute a negligible amount of FLOP compared to layers.
For the first layer, the backward-forward FLOP ratio is 1:1
For all other layers, the backward-forward FLOP ratio is 2:1 (ignoring ReLUs)
In equation form, the formula for the backward-forward FLOP ratio is:
backward / forward =
(FIRST LAYER FORWARD FLOP + 2OTHER LAYERS FORWARD FLOP + WEIGHT UPDATE) / (FIRST LAYER FORWARD FLOP + OTHER LAYERS FORWARD FLOP)
There are two considerations to see which terms dominate in this equation:
How much of the computation happens in the first layer?
How many operations does the weight update take compared to the computation in the layers? If the batch size is large or many parameters are shared, this term can be dismissed. Otherwise, it can be approximated as WEIGHT UPDATE ≈ FIRST LAYER FORWARD FLOP + OTHER LAYERS FORWARD FLOP.
This leads us to four possible cases:
Big weight update
Small weight update
First layer dominant
2FIRST LAYER FORWARD FLOP / FIRST LAYER FORWARD FLOP = 2:1
FIRST LAYER FORWARD FLOP / FIRST LAYER FORWARD ...