Linear Head Prediction

Transfer Learning: Linear Head

In feature extraction (the simplest form of transfer learning), a pretrained network is used as a fixed feature extractor. We freeze its weights and train only a small linear head on top.

The pipeline:

1. Pass data through the frozen backbone → get embeddings $Z$ (shape: m × d)

2. Apply a linear head: $\hat{y}=ZW$ (shape: m × 1, for regression)

3. Compute MSE loss on the head's predictions

Only the head's weights $W$ are updated during fine-tuning.

Your task:

Implement linear_head_predict(Z, W) that returns the predictions $\hat{y}=ZW$ and linear_head_mse(Z, W, y_true) that returns the MSE loss. Implement them as two separate functions.