I am making a NN library without any other external NN library, so I am implementing all layers, including the flatten layer, and algorithms (forward and backward pass) from scratch. I know the forward implementation of the flatten layer, but is the backward just reshaping it or not? If yes, can I just call a simple NumPy's reshape
function to reshape it?
2 Answers
Yes, a simple reshape would do the trick. A flattening layer is just a tool for reshaping data/activations to make them compatible with other layers/functions. The flattening layer doesn't change the activations themselves, so there is no special backpropagation handling needed other than changing back the shape.
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.