Feature distortion, in two dimensions
Companion to the notes on Kumar, Raghunathan, Jones, Ma & Liang, Fine-Tuning can Distort Pretrained Features and Underperform Out-of-Distribution (ICLR 2022). Back to the notes.
The model is f(x) = v · (B·x) with a single feature B = (bx, by) and a scalar head v. The truth is w* = 2·(cos 45°, sin 45°). The only training example is x = (1, 0), on the horizontal axis, so training can only see w's horizontal coordinate. A test point on the vertical axis is the "out-of-distribution" direction: its error is the gap between the predictor's vertical coordinate and the truth's, |w*y − v·by|.
Fine-tuning runs gradient flow on both
v and B. The gradient with respect to B is (2 v r, 0), where r is the training residual. It has no vertical component, so by never changes, while bx and v do. That is the whole distortion.1. Drag the pretrained feature and the head
2. What to look for
- Set v0 equal to 2, the true head size. The fine-tuning error nearly vanishes: when the head is already right it does not have to drag the feature to fit the training point (Theorem 3.3's lower bound scales with the alignment error
|v0² − v*²|). - Move B0 far from 45°. Linear probing now overshoots, because it only matches the horizontal coordinate, and fine-tuning can win. The theorem is about good features (small ε); the toy shows the other side too.
- Put B0 near 46° and v0 far from 2: fine-tuning is far worse than linear probing, which is the paper's regime.
- Scrub time:
by(the vertical coordinate of the feature) never moves, onlybxdoes.