Probing Layers › Task Localization

Task localization is a more targeted application of layer probing. Rather than asking "what is encoded at each layer?", it asks "at which specific layer(s) is the information most concentrated?" This matters for transfer learning, pruning, and deciding which intermediate representations to use for downstream tasks.

The simplest localization metric is the argmax of the probe accuracy curve: the layer that achieves the highest accuracy for a given task. But peak accuracy alone can be misleading. A more informative measure is the area under the probe curve (AUPC), which considers accuracy across all layers and rewards models where information is accessible from many layers (robust encoding) versus only one (fragile encoding).

Another useful measure is the layer of first saturation: the earliest layer at which the probe accuracy reaches, say, 95% of its eventual maximum. This identifies the minimum depth needed for a task, which has implications for efficient inference. If POS tagging saturates at layer 3 of a 12-layer model, you can potentially cut computation short for tasks that only require POS information.

The saturation analysis has motivated a line of research on early exit inference, where the model produces a prediction as soon as its representations at an intermediate layer are confident enough, without running all remaining layers. For tasks that saturate early, early exit can reduce computation by 50% or more with minimal accuracy loss. Layer probing thus has direct engineering applications beyond pure interpretability.

You can also compute a center of mass for the probe curve: a weighted average of layer indices, with probe accuracy as weights. This gives a single number summarizing how deep the information is located, on average. A center of mass near 3 means the property is mostly in early layers; a center near 10 means it requires deep processing. Comparing centers of mass across tasks reveals the relative depth of different types of linguistic knowledge.

Layer localization has direct consequences for how you should do transfer learning. The common practice of fine-tuning the entire model or taking the final layer's [CLS] token works well when your downstream task aligns with the information at the top layers. But for tasks that are primarily syntactic (like constituency parsing), you might get better performance, or faster convergence, by starting from a middle-layer representation.

This insight underlies layer selection in parameter-efficient fine-tuning. Adapter layers and LoRA modules are often inserted only at certain transformer layers. Understanding which layers encode task-relevant information helps you target the insertions where they matter most, rather than distributing them uniformly.

In practice, the question of which layer to use for a downstream task is often answered by validation performance rather than probing results. But probing provides a principled starting point: begin your search at the layer where probe accuracy peaks for the most closely related linguistic property. For a text classification task that relies primarily on sentiment, start from the upper layers. For a morphological analysis task, consider the lower layers. Probing reduces the search space from all layers to a small set of plausible candidates.

Another practical consideration is the interaction between layer selection and model fine-tuning. When you fine-tune a model on a downstream task, the weights of all layers shift. A probing result from the pre-trained model may not predict which layers will be most useful after fine-tuning, especially if the fine-tuning task requires substantially different representations from the pre-training objective. In these cases, probing the fine-tuned model (rather than only the pre-trained one) gives a more accurate picture of where the task-relevant information resides.

A powerful visualization is a heatmap of probe accuracy, where rows are tasks and columns are layers. This gives you an overview of the model's internal organization and lets you identify which layers are "specialists" (high accuracy on a narrow set of tasks) versus "generalists" (moderate accuracy on many tasks).