How does PyTorch initialize nn.Linear weights by default, and when do you override it?

## Layers initialize themselves A newly constructed `nn.Linear(64, 10)` is not empty. Its `__init__` allocates the `weight` and `bias` tensors and then calls `self.reset_parameters()`, which fills them. This is why you can build a model and immediately run a forward pass and get finite, reasonably scaled activations. For `nn.Linear` the current implementation uses `init.kaiming_uniform_(self.weight, a=math.sqrt(5))` for the weight, and initializes the bias uniformly in `(-bound, bound)` where `bound = 1 / sqrt(fan_in)`. Convolution layers use the same scheme. The `a=sqrt(5)` choice makes the effective scale equivalent to an older uniform heuristic; it is historical, not a claim that your network uses leaky ReLU with that slope. ## Why initialization scale matters at all Each layer multiplies its input by a weight matrix. If the weights are systematically too large, activation magnitudes grow layer by layer, and in a deep stack you get overflow or a loss that is `nan` on the first step. Too small, and activations shrink toward zero, gradients shrink with them, and the early layers barely learn — the classic vanishing-signal picture. The standard schemes all answer the same question: what variance should the weights have so that the variance of the activations is roughly preserved from layer to layer? Two families dominate: - **Xavier / Glorot** (`nn.init.xavier_uniform_`, `nn.init.xavier_normal_`) balances fan-in and fan-out, derived assuming a symmetric activation with unit derivative near zero — tanh, sigmoid, linear. - **Kaiming / He** (`nn.init.kaiming_uniform_`, `nn.init.kaiming_normal_`) accounts for ReLU zeroing roughly half the activations, so it scales up by a compensating factor. The `nonlinearity="relu"` argument selects the right gain; `mode="fan_in"` (the default) preserves forward-pass variance, `mode="fan_out"` preserves backward-pass variance. `nn.init.calculate_gain(nonlinearity)` exposes the multiplier the library uses for a given activation, which is handy when you write a custom scheme. ## Applying a scheme with apply() `nn.Module.apply(fn)` recursively applies `fn` to every submodule and then to the module itself. That makes it the standard place to put initialization: ``` def init_weights(m): if isinstance(m, nn.Linear): nn.init.kaiming_normal_(m.weight, nonlinearity="relu") if m.bias is not None: nn.init.zeros_(m.bias) model.apply(init_weights) ``` Note the `m.bias is not None` guard: layers built with `bias=False` have `bias` set to `None`, and forgetting the check is a common crash. Note also that all `nn.init` functions end in an underscore — they are in-place, and they are wrapped so they do not record on the autograd graph. ## When the default is fine and when it is not For a shallow MLP or any network with normalization layers between the linear layers, the default is almost always fine — normalization re-centres and re-scales activations at every layer, which washes out moderate initialization differences. The cases where overriding earns its keep: - **Deep nets without normalization.** Here the scale compounds and the difference between Xavier and Kaiming is the difference between training and not training. - **Residual stacks.** A common trick is to zero-initialize the last layer of each residual branch (`nn.init.zeros_`) so that at step zero every block is the identity and the signal passes through unchanged. Training then grows the branches gradually. - **Output layers with known priors.** Initializing a classifier's bias to the log-odds of the class prior makes the first few steps far more stable for very imbalanced data. - **Embedding layers.** `nn.Embedding` defaults to `N(0, 1)`, which is a much larger scale than a fan-in-scaled linear layer; many architectures re-initialize embeddings to a smaller standard deviation with `nn.init.normal_(emb.weight, std=0.02)`. ## Practical cautions Do the initialization **before** you construct the optimizer only if the optimizer holds references — in fact `torch.optim` holds references to the parameter tensors themselves, and `nn.init` functions mutate in place, so the order does not matter. What *does* matter is doing it before loading a checkpoint: `load_state_dict` overwrites everything, so any custom initialization applied afterwards silently destroys the loaded weights. And initialization is a training-stability tool, not a substitute for the rest of the recipe. If a model diverges, initialization is one candidate among several, and the cheapest way to test it is to print the standard deviation of activations at a few depths on one forward pass and see where the scale drifts.