Linear, Bias, and the MLP

nn.Linear — the most important layer

nn.Linear(in_features, out_features) applies the affine transformation y = x W^T + b. The weight matrix is (out_features, in_features), the bias is (out_features,). This single layer is where 90% of the parameters in most models live (Transformer FFN blocks, classifier heads, embedding projections — all Linear).

Initialization defaults

By default, nn.Linear initializes the weight with Kaiming uniform (good for ReLU-family activations) and the bias with a uniform distribution scaled by the input fan-in. You usually don't override these — but knowing they exist is what makes "my model isn't training" debugging tractable.

The MLP idiom

A multi-layer perceptron is just Linear → activation → Linear → activation → ... → Linear. Modern variants add Dropout for regularization, LayerNorm for stability, and skip connections for very deep networks. The pattern is universal — every Transformer FFN block is a 2-layer MLP with a non-linearity (typically GELU) in the middle.

Code

nn.Linear basics·python

import torch
import torch.nn as nn

# 20 inputs → 10 outputs
linear = nn.Linear(20, 10)
print(linear.weight.shape)   # torch.Size([10, 20])  — out_features × in_features
print(linear.bias.shape)     # torch.Size([10])

x = torch.randn(32, 20)      # batch=32, features=20
y = linear(x)
print(y.shape)               # torch.Size([32, 10])

# Without bias (rare but useful in some Transformer variants)
linear_nb = nn.Linear(20, 10, bias=False)
print(linear_nb.bias)        # None

Build an MLP — the canonical pattern·python

import torch
import torch.nn as nn

class MLP(nn.Module):
    def __init__(self, in_dim, hidden, out_dim, p_drop=0.1):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(in_dim, hidden),
            nn.GELU(),                # transformer-flavor; ReLU works too
            nn.Dropout(p_drop),
            nn.Linear(hidden, hidden),
            nn.GELU(),
            nn.Dropout(p_drop),
            nn.Linear(hidden, out_dim),
        )

    def forward(self, x):
        return self.net(x)

model = MLP(784, 256, 10)
print(f"Params: {sum(p.numel() for p in model.parameters()):,}")
# Params: 269,322

Custom init — when defaults aren't what you want·python

import torch
import torch.nn as nn
import math

class MLP(nn.Module):
    def __init__(self, in_dim, hidden, out_dim):
        super().__init__()
        self.fc1 = nn.Linear(in_dim, hidden)
        self.fc2 = nn.Linear(hidden, out_dim)
        self.act = nn.GELU()
        self._init_weights()

    def _init_weights(self):
        # Custom Xavier init — useful when you want a specific behavior
        for m in self.modules():
            if isinstance(m, nn.Linear):
                nn.init.xavier_normal_(m.weight)
                if m.bias is not None:
                    nn.init.zeros_(m.bias)

    def forward(self, x):
        return self.fc2(self.act(self.fc1(x)))

nn.Linear — the most important layer

Initialization defaults

The MLP idiom

Code

External links

Exercise

Progress

Comments 0