Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

KA2L: A Knowledge-Aware Active Learning Framework for LLMs

Haoxuan Yin, Bojian Liu, Chen Tang, Yangfan Wang, Lian Yan, Jingchi Jiang

Abstract

Fine-tuning large language models (LLMs) with high-quality knowledge has been shown to enhance their performance effectively. However, there is a paucity of research on the depth of domain-specific knowledge comprehension by LLMs and the application of targeted active learning to improve their expertise. To address this gap, we introduce the Knowledge-Aware Active Learning (KA2L) framework. This framework assesses LLMs' mastery of specific knowledge points to aid in constructing unanswerable or unknowable questions through latent space analysis. This active learning strategy enhances training efficiency by focusing on knowledge the model has yet to master, thereby minimizing redundancy in learning already acquired information. This study innovatively employs a knowledge distribution probing technique to examine the hidden states of specific Transformer layers and identify the distribution of known and unknown knowledge within the LLM. Additionally, a hidden-state decoding method is proposed to generate numerous unknown questions in natural language from the latent knowledge space. In our experiments, we selected nine open-source LLMs to validate the effectiveness of the proposed framework. Results indicate that KA2L not only significantly reduces 50% annotation and computation costs across two open-domain and one vertical-domain dataset but also achieves better performance, offering valuable insights into active learning strategies for LLMs. The code is available at https://anonymous.4open.science/r/KA2L-F15C.

KA2L: A Knowledge-Aware Active Learning Framework for LLMs

Abstract

Fine-tuning large language models (LLMs) with high-quality knowledge has been shown to enhance their performance effectively. However, there is a paucity of research on the depth of domain-specific knowledge comprehension by LLMs and the application of targeted active learning to improve their expertise. To address this gap, we introduce the Knowledge-Aware Active Learning (KA2L) framework. This framework assesses LLMs' mastery of specific knowledge points to aid in constructing unanswerable or unknowable questions through latent space analysis. This active learning strategy enhances training efficiency by focusing on knowledge the model has yet to master, thereby minimizing redundancy in learning already acquired information. This study innovatively employs a knowledge distribution probing technique to examine the hidden states of specific Transformer layers and identify the distribution of known and unknown knowledge within the LLM. Additionally, a hidden-state decoding method is proposed to generate numerous unknown questions in natural language from the latent knowledge space. In our experiments, we selected nine open-source LLMs to validate the effectiveness of the proposed framework. Results indicate that KA2L not only significantly reduces 50% annotation and computation costs across two open-domain and one vertical-domain dataset but also achieves better performance, offering valuable insights into active learning strategies for LLMs. The code is available at https://anonymous.4open.science/r/KA2L-F15C.
Paper Structure (39 sections, 5 equations, 3 figures, 11 tables, 1 algorithm)

This paper contains 39 sections, 5 equations, 3 figures, 11 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Work
  3. Active Learning for LLMs
  4. LLM Hallucination Detection
  5. LLM Hidden State Decoding
  6. Method
  7. Problem Formulation
  8. Knowledge Distribution Probing
  9. Hidden State and Model Output Sampling
  10. Semantic Entropy Calculation and Label Construction
  11. Classifier Training and Inference
  12. Unknown Question Augmentation via Hidden State Decoding
  13. Downstream Tasks
  14. Experimental Setup
  15. Models, Datasets, and Evaluation Metrics
  16. ...and 24 more sections

Figures (3)

  • Figure 1: KA2L Workflow: (a) Knowledge Distribution Probing: Training Phase: For each question in the sampled question set, sample its hidden state once and its textual outputs multiple times. Perform semantic clustering on the textual outputs to calculate Semantic Entropy (SE). The SE is then binarized using a dynamic threshold to serve as labels for the classifier. An MLP classifier is trained using these hidden states and the binarized SE (BiSE). Inference Phase: For a new set of questions, sample their $l$-th layer hidden states. These are then classified by the MLP classifier, representing the knowledge distribution as "Known" and "Unknown" knowledge. (b) Unknown Question Augmentation: Sample the hidden states ($h_{unk}$) of questions identified as "Unknown" from the knowledge distribution. These are then transformed and decoded into multiple similar questions. (c) Downstream Tasks: This knowledge distribution guides dataset construction and model fine-tuning. Many existing methods can be applied, such as LoRA hu2022lora and P-tuning liu-etal-2022-p.
  • Figure 2: Classifier AUROC score across different hidden layers (embedding layer excluded) on the Trivia_QA dataset.
  • Figure 3: PCA Analysis on Mistral-7B-Instruct-v0.3