Relational Knowledge Distillation Using Fine-tuned Function Vectors

TL;DR

This work tackles how relational knowledge can be distilled and manipulated inside LLMs by turning implicit task representations into explicit vectors. It introduces fine-tuned function vectors (FFV) to refine relation representations with minimal data, and a composite function vector (CFV) that linearly combines FFVs to transfer relational knowledge to novel tasks, all through activation-steering without updating the base model. FFVs yield substantial gains in zero-shot relation tasks and align more closely with human relational similarity judgments, while CFVs enable one-shot and cross-domain analogy solving, including challenging far-analogy and SAT-style problems, across multiple models. The results demonstrate that activation patching and linear vector representations can enhance interpretability and reasoning in LLMs, suggesting a scalable path toward explicit, human-aligned relational knowledge in AI systems.

Abstract

Representing relations between concepts is a core prerequisite for intelligent systems to make sense of the world. Recent work using causal mediation analysis has shown that a small set of attention heads encodes task representation in in-context learning, captured in a compact representation known as the function vector. We show that fine-tuning function vectors with only a small set of examples (about 20 word pairs) yields better performance on relation-based word-completion tasks than using the original vectors derived from causal mediation analysis. These improvements hold for both small and large language models. Moreover, the fine-tuned function vectors yield improved decoding performance for relation words and show stronger alignment with human similarity judgments of semantic relations. Next, we introduce the composite function vector - a weighted combination of fine-tuned function vectors - to extract relational knowledge and support analogical reasoning. At inference time, inserting this composite vector into LLM activations markedly enhances performance on challenging analogy problems drawn from cognitive science and SAT benchmarks. Our results highlight the potential of activation patching as a controllable mechanism for encoding and manipulating relational knowledge, advancing both the interpretability and reasoning capabilities of large language models.

Figures (11)

• Figure 1: Illustration of in-context learning and zero-shot task with function vector intervention in LLMs. The context includes a sequence of word pairs instantiating the antonym relation (e.g., new : old, short : long, etc.). Then, a query word ("rich") is provided and the model needs to predict the next word. If the model predicts "poor" with the highest probability, that will be counted as a correct response under the top-1 criterion of accuracy. The pre-trained LLM’s weights remain unchanged during in-context learning, yet some form of learning evidently takes place, as performance on the task is enhanced when demonstration examples are provided.
• Figure 2: Illustration of the training procedure for the fine-tuned function vector. For a relation $z$, the function vector $\mathbf{v_z}$ is updated using its word pairs (e.g., A : B). The final loss incorporates the cross-entropy from this word pair as well as a weighted L2 norm of the vector for regularization.
• Figure 3: Illustration of the composite function vector for a four-term analogy task (e.g., (nose : scent :: antenna : ?). The composite function vector, computed as the weighted sum of the fine-tuned function vectors $\mathbf{v}$ and the posterior distribution $w$ given the source pair (e.g., nose : scent), is injected into the LLM to influence its output predictions for the target. (e.g., (antenna : ?).
• Figure 4: Zero-shot evaluation results for GPT-J: Top-1 prediction accuracy for the 6 representative relations (top left) and problem-solving relations (top right) in the simple-task dataset, and Top-5 prediction accuracy for the SemEval dataset by relation type (bottom). Baseline performance (grey) reflects running GPT-J directly on the Zero-shot task. Compared to the initial FV (green), injecting the fine-tuned FV (orange) into GPT-J yielded the highest average accuracies. The error bars indicate the standard deviation across five runs. FV = function vector.
• Figure 5: Relational dissimilarity matrices of human judgments and model predictions. Each row and column represents a word pair, and each cell represents the pairwise dissimilarity between the row's word pair and the column's word pair. The warmer the cell's color, the higher its dissimilarity.