Fine-tuned network relies on generic representation to solve unseen cognitive task

TL;DR

This work fine-tuned GPT-2 on a context-dependent decision-making task, novel to the model but adapted from neuroscience literature, and compared its performance and internal mechanisms to a version of GPT-2 trained from scratch on the same task.

Abstract

Fine-tuning pretrained language models has shown promising results on a wide range of tasks, but when encountering a novel task, do they rely more on generic pretrained representation, or develop brand new task-specific solutions? Here, we fine-tuned GPT-2 on a context-dependent decision-making task, novel to the model but adapted from neuroscience literature. We compared its performance and internal mechanisms to a version of GPT-2 trained from scratch on the same task. Our results show that fine-tuned models depend heavily on pretrained representations, particularly in later layers, while models trained from scratch develop different, more task-specific mechanisms. These findings highlight the advantages and limitations of pretraining for task generalization and underscore the need for further investigation into the mechanisms underpinning task-specific fine-tuning in LLMs.

Figures (11)

• Figure 1: Task schema. Each trial of the context-dependent decision-making (CDDM) task is converted to text description and prompted to GPT-2. The model is trained to respond based on the context cue and sensory evidence, similar to the animal behavioral task.
• Figure 2: Last layer hidden states from different trials at different tokens after dimensionality reduction by UMAP, colored by token position (A) or different values of each behavioral variable (shown on top of each panel) (B)
• Figure 3: Accuracy of decoding each task-relevant variable from the hidden states at different tokens, for one example unit. Grey dashed lines indicate label-shuffled decoding accuracy as a baseline. Solid lines and shaded area indicate mean and standard deviation of decoding accuracy across 5 cross-validation folds. More example units shown in Figure \ref{['fig:supp fig 2 mixed selectivity']}
• Figure 4: Evaluation accuracy, measured in the percentage of correct repsonses, after zero-ablating each attention head for a pretrained network fine-tuned on the task (left) and a network trained on the task from scratch (right).
• Figure 5: Average attention weights ($\text{softmax}\left(\frac{QK^T}{\sqrt{d}}\right)$) across prompts from a fine-tuned network for 4 example attention heads that caused significant performance drop when zero-ablated. y-axis is source token, x-axis is destination token. Darker color indicates larger value.