Neural networks, deep learning, NLP · insight-0439

Watch the flatten-to-dense layer, not the convolutions, for parameter blow-up

Situation
A small CNN for shelf images or X-ray triage has far more parameters than expected and overfits quickly.
Action
Count parameters per layer. Replace flatten-plus-dense with global average pooling or more downsampling, so the head sees channels, not pixels.
Reason
A convolution's weights depend on kernel size and channel counts only, while a dense layer's weights grow with every spatial position fed into it.
Boundary
CNNs built from scratch or with a custom head on a pretrained backbone.

Symptom or risk

Flattening a large feature map into a dense layer creates a weight matrix that dwarfs all convolutions combined.

Tasks needing precise spatial layout may keep spatial features, typically through convolutional heads.

Global pooling removes position information; check that the task does not need it.

Increasing input resolution later and silently multiplying the head's size.

A 3 x 3 convolution from 64 to 128 channels has 73,856 parameters; flattening a 128 x 56 x 56 map into 256 units needs about 102.8 million.

PyTorch Conv2d weights have shape (out_channels, in_channels/groups, kernel height, kernel width), independent of the input's spatial size.
Conv2d (PyTorch documentation)
LeCun et al. describe convolutional networks with shared weights over local receptive fields.
Gradient-based learning applied to document recognition