Task Knowledge Injection via Interpolations and Reinstatement for Large Language Model Generalization (2025)

Abstract: Large language models have shown tremendous potential across various NLP tasks, and instruction tuning has been widely adopted to elicit their superior performance. However, instruction tuning may overly tailor the models to task-specific formats, potentially compromising their generalization on unseen tasks. We attribute the issue to the spurious correlations learned between inputs and targets. We propose explicit task knowledge injection to mitigate these shortcuts with latent task adaptation and knowledge reinstatement. Latent tasks serve as interpolations between new tasks and facilitate knowledge sharing with joint adaptation enabling the model to build task knowledge more smoothly. Knowledge reinstatement helps optimize building new knowledge with prior knowledge. Specifically, we retrieve input-relevant latent tasks and jointly learn the task and the relevant latent tasks. Moreover, we prompt the model to recall the forms of inputs corresponding to the target and build the task knowledge through the reinstatement of prior knowledge while learning the new task.We conduct extensive experiments on state-of-the-art large language models including Llama3.1-8B and Vicuna-13B across 1000+ instruction-following tasks to demonstrate the effectiveness of our method. The results demonstrate our method improves generalization on both in-domain and out-of-domain unseen tasks.

## introduction pre-trained large language models @xcite @xcite @xcite have become a cornerstone in many nlp tasks due to their impressive generalization capabilities @xcite @xcite . these models can be prompted with arbitrary demonstrations to accomplish various tasks. to further en-* co-corresponding authors. current techniques that help improve instruction tuning generalization can be categorized into four groups: (1) data-based methods @xcite @xcite @xcite @xcite incorporate broad-coverage tasks, contextual demonstrations, chain-of-thought demonstrations, similar task augmentation, more complex and high-quality data into the training data to improve generalization, which requires carefully curated examples. unlike previous studies, we attribute the issue of generalization on unseen tasks to learned spurious correlations between inputs and targets @xcite @xcite . specifically, a task can have multiple verbalizations, and if a model is trained with only a few instructions, the sequence-to-sequence training paradigm may cause the model to learn shortcuts like certain words or sentence structures. we address this by explicitly injecting relevant knowledge related to current tasks to help the model build knowledge more effectively about the tasks rather than just superficial patterns. to this end, we propose task knowledge injection with latent task adaptation and knowledge reinstatement. our method first retrieves relevant latent tasks by comparing the task input with latent task keys. latent tasks are interpolations between the tasks represented as tuples of keys (vectors) and corresponding knowledge (compiled as weight increments), which are initialized randomly. then we make the model jointly learn the current task and relevant latent tasks by re-parameterizing the model parameters and weight increments. this process facilitates knowledge sharing and enables the model to build task knowledge more smoothly. moreover, our method prompts the model to recall which forms of inputs in the existing model memory might generate the target. we make the model learn through the reinstatement of the prior knowledge during the current task learning, which helps build more comprehensive knowledge about the task. we conduct extensive experiments on publicly available large language models, vicuna-13b and llama3.1-8b across 1000+ instruction-following tasks. the experimental results demonstrate that our approach improves generalization on both indomain and out-of-domain unseen tasks. the contributions of our paper are summarized as follows: fine-tuning methods for improved generalization. we compile the methods related to instruction-tuning, robustness, generalization, and continual learning into (1) data-based, (2) parameter-efficient fine-tuning (peft), (3) adversarial training, (4) regularization-based methods. ## problem formulation and motivation in this paper, we aim to fine-tune existing models to continually improve their performance on downstream tasks. given a set of tasks @xmath0, ..., t k } along with their training examples instruction tuning may cause the model to learn spurious correlations between inputs and targets under the sequence-to-sequence training paradigm @xcite , resulting in limited performance. inspired by @xcite , which suggests that congruency in reinstatement and memory augmentation facilitates the assimilation and construction of new knowledge, we propose task knowledge injection with latent task adaptation and knowledge reinstatement. latent tasks serve as interpolations between new tasks, allowing the model to learn new knowledge more smoothly. joint adaptation with relevant latent tasks that share common knowledge provides useful inductive biases, helping the model build a more comprehensive understanding of the task. fine-tuning often leads to knowledge drift, where previously learned information is overwritten. by reinstating prior knowledge while learning new tasks, the model preserves critical learned representations and strengthens its understanding of the new task in a way that aligns with its existing knowledge base. ## task knowledge injection in this section, we introduce our task knowledge injection method. we introduce an overview first in subsection 4.1. we introduce latent task adaptation in subsection 4.2 and then introduce knowledge reinstatement in subsection 4.3. ## conclusion in this paper, we propose a novel instruction-tuning method called task knowledge injection. it is a data-efficient knowledge injection approach that enhances task generalization without relying on additional models or data. we conduct extensive experiments on recent large language models, including llama3.1-8b and vicuna-13b, across 1000+ tasks from super natural instructions and flan datasets. the experimental results demonstrate the superior performance of our method and each component on both in-domain and out-of-domain datasets. our method enables the model to establish correlations between new tasks and between new tasks and prior knowledge, facilitating building task knowledge more effectively, and thereby improving generalization performance. we conduct detailed analyses revealing that joint task adaptation plays a more critical role in the knowledge injection process. we also show that incorporating more relevant latent tasks leads to slightly improved performance while weakening knowledge sharing results in degraded performance.