The Flatten
layer has no learnable parameters in itself (the operation it performs is fully defined by construction); still, it has to propagate the gradient to the previous layers.
In general, the Flatten
operation is well-posed, as whatever is the input shape you know what the output shape is.
When you backpropagate, you are supposed to do an "Unflatten", which maps a flattened tensor into a tensor of a given shape, and you know what that specific shape is from the forward pass, so it is also a well-posed operation.
More formally
Say you have Img1
in input of your Flatten
layer
$$
\begin{pmatrix}
f_{1,1}(x; w_{1,1}) & f_{1,2}(x; w_{1,2}) \\
f_{2,1}(x; w_{2,1}) & f_{2,2}(x; w_{2,2})
\end{pmatrix}
$$
So, in the output you have
$$
\begin{pmatrix}
f_{1,1}(x; w_{1,1}) & f_{1,2}(x; w_{1,2}) & f_{2,1}(x; w_{2,1}) & f_{2,2}(x; w_{2,2})
\end{pmatrix}
$$
When you compute the gradient you have
$$
\frac{df_{i,j}(x; w_{i,j})}{dw_{i,j}}
$$
and everything in the same position as in the forward pass, so the unflatten maps from the (1, 4)
tensor to the (2, 2)
tensor.