Linear Autoencoder Meets PCA

Train an autoencoder with no nonlinearity and see it find the PCA solution.

An autoencoder squeezes each input through a small code and tries to rebuild it. With linear layers and a squared error loss, the best possible autoencoder with a k dimensional code reconstructs exactly as well as PCA with k components. Here you build one in PyTorch and train it until it gets there.

encode(x) = (x - mean) @ W_enc          shape (n, k)
decode(z) = z @ W_dec + mean            shape (n, d)
forward(x) = decode(encode(x))
loss       = mean squared error between forward(X) and X

Write an nn.Module with encode, decode and forward, then train it with Adam on the full batch for steps steps and return it.

Function signature

def train_linear_autoencoder(X, k, steps=300):

Example

Input

X = data that lies close to a 4 dimensional plane inside 20 dimensions
model = train_linear_autoencoder(X, k=4)

Output

model.encode(X).shape  ->  torch.Size([400, 4])
mse(model(X), X)       ->  about 0.070, the same as PCA with 4 components

Reasoning

PCA keeps the k directions of largest variance and drops the rest. The error that remains is the variance in the dropped directions. A linear autoencoder cannot do better, and once trained it does no worse.

Constraints and notes

• X is a float32 tensor of shape (n, d), and its mean is far from zero. Centre the data inside the module. Store the training mean with self.register_buffer("mean", ...), so that it moves with the module and is saved in its state_dict, but is not trained.
• forward(x) must equal decode(encode(x)). No skip connections.
• Graded on the result: the reconstruction error must be within 2% of the PCA optimum, on the training data and on new data from the same source.