GNNavi: Navigating the Information Flow in Large Language Models by Graph Neural Network

TL;DR

GNNavi tackles the efficiency bottleneck of fine-tuning large language models by integrating a Graph Neural Network (GNN) layer into decoder-only LLMs to steer the information flow within prompts. By treating the input sequence as a graph and hard-wiring the flow through label-word anchors, GNNavi updates only the GNN parameters (about $0.2\%$ to $0.5\%$ of the full model) yet outperforms full parameter fine-tuning and other PEFT methods in few-shot text classification tasks. The method leverages two GNN variants, GCN and GraphSAGE, inserted after a mid-to-deep layer ($l$-th decoder layer) to aggregate context and distribute information to the final token. This approach yields strong few-shot performance on GPT2-XL and Llama2 across SST-2, EmoC, TREC, Amazon, and AGNews, while significantly reducing computational demands and revealing clearer information-flow dynamics than FPFT.

Abstract

Large Language Models (LLMs) exhibit strong In-Context Learning (ICL) capabilities when prompts with demonstrations are used. However, fine-tuning still remains crucial to further enhance their adaptability. Prompt-based fine-tuning proves to be an effective fine-tuning method in low-data scenarios, but high demands on computing resources limit its practicality. We address this issue by introducing a prompt-based parameter-efficient fine-tuning (PEFT) approach. GNNavi leverages insights into ICL's information flow dynamics, which indicates that label words act in prompts as anchors for information propagation. GNNavi employs a Graph Neural Network (GNN) layer to precisely guide the aggregation and distribution of information flow during the processing of prompts by hardwiring the desired information flow into the GNN. Our experiments on text classification tasks with GPT-2 and Llama2 show GNNavi surpasses standard prompt-based fine-tuning methods in few-shot settings by updating just 0.2% to 0.5% of parameters. We compare GNNavi with prevalent PEFT approaches, such as prefix tuning, LoRA and Adapter in terms of performance and efficiency. Our analysis reveals that GNNavi enhances information flow and ensures a clear aggregation process.

Figures (7)

• Figure 1: Visualization of Full Parameter Fine-tuning (FPFT) and GNNavi from the perspective of information flow (top words to bottom words). Without GNNavi, tokens interact with every preceding word in FPFT, leading to confusion in information flow. Conversely, in GNNavi, label words aggregate information from preceding words (blue path), and the final token aggregates information from the label words (pink path), resulting in a clearer information aggregation process.
• Figure 2: Visualization of GNNavi with an example of sentiment analysis, where label words and the last token are highlighted in blue and pink, respectively. a) The GNN layer is integrated into a decoder-only LLM. The LLM processes a prompt containing demonstrations and generates the next token as the prediction. b) The input text is transformed into a graph, with tokens as nodes and information flow paths as edges. c) Visualizing the working mechanism of the GNN: Node representations are updated by aggregating information from neighboring nodes. To maintain simplicity, not all nodes are listed.
• Figure 3: Results of average accuracy with different number of training examples. The x-axis denotes the number of training examples per class.
• Figure 4: The improvement gained by adding training examples for GNNavi-SAGE, compared to using 5 training examples per class.
• Figure 5: Performance Comparison with GNN inserted at various positions in GPT2-XL.