Where Fine-Tuning Changes DistilBERT's Representations

In the previous post, a linear probe on frozen DistilBERT reached 85.9% accuracy on IMDB sentiment. Centered Kernel Alignment (CKA) makes it possible to compare the pretrained and fine-tuned representations layer by layer. The lower layers remain similar; the largest differences appear in the final two transformer blocks.

What CKA compares

At a fixed layer, each checkpoint produces one activation vector for every review in the same input set. CKA compares the geometry of those two sets of vectors. If reviews that are similar in one representation are also similar in the other, CKA is high. A value near 1 indicates closely aligned geometry under the chosen kernel; lower values indicate greater divergence.

Linear CKA ignores rotations and uniform rescaling but still distinguishes more general changes to a representation. This balance makes it useful for comparing neural networks with wide hidden layers[1]. It does not identify which computation changed or whether the difference matters for the task.

Experiment

The comparison uses pretrained distilbert-base-uncased and an IMDB-fine-tuned checkpoint[2]. DistilBERT[3] exposes seven representations: the embedding output at layer 0 and the outputs of its six transformer blocks. Mean pooling reduces each layer's token vectors to one vector per review. The checkpoints remain frozen throughout.

Both models receive the same balanced sample of 2,000 IMDB reviews in the same order[4]. This correspondence is required because CKA compares the representation of each review with that same review in the other model. The probe accuracies below use this 2,000-review sample as well, so they are slightly lower than the full-dataset results in the previous post.

Where the representations change

The diagonal comparison matches each pretrained layer with the same layer in the fine-tuned checkpoint.

CKA is 0.999 at the embedding layer and remains above 0.77 through layer 4. It then falls to 0.267 at layer 5 and 0.158 at layer 6. The lower-layer representations therefore remain relatively similar after fine-tuning, while the final two blocks diverge sharply.

The full matrix checks whether a fine-tuned layer resembles a pretrained layer at a different depth.

The upper-left block is bright, with the closest matches on or near the diagonal. The last two fine-tuned columns are dark across every pretrained layer. Under linear CKA, the fine-tuned top representations have no close match elsewhere in the pretrained model; their difference is not explained by a representation moving to another depth.

CKA and probe accuracy

The linear probes use the same sample, allowing representation similarity to be compared with sentiment decodability.

Layer	CKA \(s_\ell\)	Pretrained acc	Fine-tuned acc	Δ acc
0	0.999	0.773	0.787	+0.013
1	0.965	0.765	0.775	+0.010
2	0.901	0.762	0.795	+0.033
3	0.835	0.768	0.812	+0.043
4	0.779	0.780	0.860	+0.080
5	0.267	0.815	0.878	+0.063
6	0.158	0.798	0.887	+0.088

At layers 0 and 1, CKA stays near 1 and probe accuracy changes by only 1.3 and 1.0 percentage points. At layers 4 through 6, the representations are less similar and probe accuracy improves by 6.3 to 8.8 points. Across these seven layers, the largest representation differences coincide with the largest probe gains. This is an association, not evidence that one causes the other.

The pattern is consistent with earlier work on BERT-family models, which also locates the largest fine-tuning changes in upper layers[5][6].

Limitations

These results describe mean-pooled activations from one pair of checkpoints on 2,000 IMDB reviews. Different inputs, pooling methods, or kernels could change the values. Linear and RBF CKA often agree empirically[1], but CKA remains descriptive: a low score does not identify the mechanism that produced the difference or show that the model uses it. The probes add task information, but decodability is not a causal test either.