Task Oriented In-Domain Data Augmentation (2024)

Abstract: Large Language Models (LLMs) have shown superior performance in various applications and fields. To achieve better performance on specialized domains such as law and advertisement, LLMs are often continue pre-trained on in-domain data. However, existing approaches suffer from two major issues. First, in-domain data are scarce compared with general domain-agnostic data. Second, data used for continual pre-training are not task-aware, such that they may not be helpful to downstream applications. We propose TRAIT, a task-oriented in-domain data augmentation framework. Our framework is divided into two parts: in-domain data selection and task-oriented synthetic passage generation. The data selection strategy identifies and selects a large amount of in-domain data from general corpora, and thus significantly enriches domain knowledge in the continual pre-training data. The synthetic passages contain guidance on how to use domain knowledge to answer questions about downstream tasks. By training on such passages, the model aligns with the need of downstream applications. We adapt LLMs to two domains: advertisement and math. On average, TRAIT improves LLM performance by 8% in the advertisement domain and 7.5% in the math domain.

## introduction large language models (llms) have achieved significant performance improvements in various applications such as language modeling @xcite @xcite and visual understanding @xcite . they have also shown superior performance in fields such as finance @xcite , e-commerce @xcite and healthcare @xcite . however, the models are usually trained on a large amount of general domain-agnostic data, such as web corpora. because of the lack of domain-specific training, llms suffer from subpar performance when directly applied to certain domains such as advertisement. to adapt llms to a specific domain, continual pre-training methods @xcite are commonly applied. in particular, the llm is continual pre-trained on in-domain corpora, such that it can acquire domain knowledge and better adapt to downstream tasks. existing works @xcite have shown that such a technique drastically improves performance of llms on domains such as law and bio-medicine. there are two major issues when continual pretraining llms. first, in-domain data are scarce. llms are pre-trained on large domain-agnostic corpora. for example, the web corpus used for pre-training contains more than ten trillion tokens. however, domain-specific data are magnitudes smaller, i.e., the ads in-domain corpus in our experiments contains only several billion tokens. such a data scarcity issue significantly hinders model performance after continual pre-training. second, in-domain data used for continual pretraining are not task-oriented. many existing works @xcite @xcite focus on generating or selecting in-domain data without considering the downstream tasks. that is, the continual pretraining data are often passages that describe keywords/concepts of the target domain, which are generated or selected without considering whether the passages benefit downstream tasks. we propose trait (task oriented in-domain data augmentation), a data augmentation framework driven by downstream tasks of the domain. the framework is divided into two parts. first, to address the data scarcity issue of in-domain corpora, we propose a data selection strategy. the proposed algorithm identifies in-domain data from general corpora, and also applies a quality filter to ensure that the selected data have high educational value @xcite . in practice, the amount of selected data is magnitudes larger than the in-domain dataset. for example, for the advertisement applications, the in-domain dataset contains about 1b tokens, and trait selects an additional 15b tokens from web corpora. second, we propose a task-oriented synthetic passage generation guideline. specifically, each generated passage contains multiple problems, where each problem comes from a different downstream task from the domain. then, for each problem, trait generates a problem-specific paragraph that suggests possible answers to the problem. additionally, the synthetic passage also contains an enlightenment paragraph. this paragraph focuses on relationships among problems in the passage, including common and individual characteristics that are used to generate answers. intuitively, the problem-specific paragraphs teach the model how to use techniques to solve a particular problem. and the enlightenment paragraph teaches the model common and unique aspects of all problems in the domain. to fully exploit the power of trait, we employ a two-stage training strategy for continual pretraining. in the first stage, the model is trained with in-domain data, including both the original in-domain corpora and the selected data. in this stage, the model adapts to the domain by learning domain knowledge. then, in the second stage, the model is trained with the task-oriented passages. during this stage, the model learns how to use domain knowledge to solve problems, such that it better aligns with the need of downstream tasks. we conduct extensive experiments by adapting llms to two domains: advertisement and math. for the advertisement domain we consider 7 downstream tasks, and trait improves existing continual pre-training methods by 6.5% average score, while improving the base llm (without continual pre-training) by 8%. for the math domain we consider 9 downstream tasks, where trait outperforms the baseline by 5.5% average score and outperforms the base llm by 7.5%. for the challenging math task in the math domain, trait outperforms the base llm by over 15%. ## data preparation we apply our data augmentation method, trait, to two domains: advertisement (ads) and math. in this section, we detail the process for in-domain data selection and synthetic passage generation. refer to appendix. e for examples and more details. ## experiments we evaluate trait by adapting llms to the ads and math domains via continual pre-training. in all the experiments, we use mistral-7b (jiang et al., 2023) and tinyllama-1.1b @xcite as the base model. we compare trait with two data selection baselines: (1) random sampling, which randomly selects samples from open-source general corpora; and (2) dsir @xcite , an importance sampling strategy for selecting indomain data from general corpora, such that the selected data distibutionally similar with the indomain data. to promote fair comparisons, all models (including the base mistral-7b model, baseline models, and trait) are trained on the same amount of data with the same computational budget. details about the training process can be found in appendix a. ## conclusion this paper presents trait, a task-oriented indomain data augmentation framework for continual pre-training of large language models. the framework is divided into two parts. first, we select indomain data from general domain-agnostic corpora to augment the training set. the augmented indomain training corpus contain rich domain knowledge. second, we generate task-oriented synthetic passages. these passages contain guidance on how to apply domain knowledge to answer questions about downstream tasks. we conduct extensive experiments by adapting llms to the advertisement and math domains. experimental results validate the effectiveness of the proposed framework. specifically, on average, trait improves the base llm (without continual pre-training) by over 5% on both domains.