GNN avi: Navigating the Information Flow in Large Language Models by Graph Neural Network (2024)

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.

## introduction large language models (llms) show remarkable in-context-learning (icl) capabilities by learning from prompts with demonstrations @xcite @xcite @xcite , with the exponential growth in model sizes. however, fine-tuning llms still remains essential for further enhancing their adaptability @xcite . prompt-based fine-tuning @xcite figure 1 : visualization of full parameter fine-tuning (fpft) and gnnavi from the perspective of information flow (top words to bottom words). without gn-navi, 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. recent advances in understanding the icl mechanism offer a new avenue for peft of llms. icl's success in leveraging few-shot demonstrations and prompts @xcite has motivated the adoption of prompt-based fine-tuning for moderately sized language models in a few-shot learning manner @xcite . recognizing the specific features of finetuning llms within the framework of icl, we propose gnnavi, a novel peft method designed expressly for prompt-based learning. our method draws inspiration from recent insights into the underlying process of icl from an information flow perspective, particularly the role of label words in the prompt @xcite . label words act as anchors with two functions: aggregating information from context words and directing this information to the last token for accurate predictions. gnnavi incorporates this understanding through the integration of a graph neural network (gnn) layer @xcite into llms, optimizing the promptbased fine-tuning process by navigating the information flow within prompts, as visualized in figure 1 . following the paths of information flow, we insert a gnn layer into the deep layers foot_1 of the llm. we treat the input text as a graph, where each token serves as a node, and connect these nodes according to the paths of information flow. the gnn layer aims to guide the information flow by aggregating information from neighbouring nodes. as a peft method, gnnavi adopts a lightweight fine-tuning strategy, updating only the parameters of the gnn layer. experimenting with few-shot training examples on gpt2-xl @xcite and llama2 @xcite , gnnavi achieves remarkable results with just 0.2% of the trainable parameters of the full model, consistently outperforming fpft and other peft methods across various classification tasks. additionally, we analyze the attention interaction between tokens and find that gnnavi demonstrates a more stable and clear information aggregation process compared to fpft. in summary, our contributions are: i) we pro-pose a novel peft method, gnnavi, inspired by the information flow perspective of llms. gn-navi effectively navigates the information aggregation process in llms. ii) we apply gnnavi to text classification tasks with few-shot training examples, outperforming baselines while updating only 0.2% to 0.5% of parameters. iii) our work sheds light on the application of gnns in nlp and provides novel insights for future research. to the best of our knowledge, we are the first to utilize gnns to enhance the performance of llms from the information flow perspective. prompt-based learning gpt-3 @xcite has sparked interest in prompt-based learning methods, and particularly in the icl paradigm. this surge in attention has fostered a multifaceted exploration into the factors influencing icl performance, including input perturbation @xcite , selection of demonstration @xcite , and calibration techniques @xcite . concurrently, there has been a deep dive into understanding the underlying mechanism of icl, employing diverse theoretical frameworks such as gradient descent @xcite , bayesian inference @xcite and information flow @xcite . following the route of icl, prompt-based fine-tuning has emerged as an effective strategy in scenarios with limited data @xcite . we leverage insights from these investigations into the icl mechanism and propose a tailored peft method for llms. ## results we report the results with 5 and 200 training examples in ## ablation study in §6.1 of this section, we delve into the influence of the position where the gnn layer is inserted in the llm. in §6.2, we investigate the effects of removing one of the information flow paths on performance. all of these studies are conducted using gnnavi-sage with 5 training samples per class under the experimental settings outlined in §4.2. the position where the gnn layer is inserted significantly impacts the model's performance. figure 5 illustrates the performance of gnnavi when the gnn layer is inserted at different locations in gpt2-xl. with the exception of emoc, all tasks exhibit lower performance when the gnn layer is added in the first 10 layers of gpt2-xl. performance improves as the gnn is added in deeper layers, reaching peak accuracy around the 44th layer. subsequently, accuracy declines until the last layer. this trend may stem from the gradual initiation of the information flow process in the early layers of llm, where the gnn's influence is limited due to insufficient token interaction. conversely, in the final layers, the information flow process is nearly complete, rendering it too late for the gnn to guide effectively. despite variations in performance changes across tasks, the average performance suggests that the optimal placement for the gnn layer is between the 38th and 42nd layers for gpt2-xl. ## further discussion: information flow while the attention mechanism in llm offers an information flow perspective for interpreting the model's working mechanism @xcite , it treats the input text as a fully connected graph. in contrast, gnnavi explicitly connects the context tokens to the label tokens for information aggregation and the label tokens to the final token for information distribution. thereby, the correct information flow is hardwired into the gnn. there is no need to learn it by adjusting the attention weights. to further investigate the differences in information flow between fpft and gnnavi, we utilize the saliency technique @xcite for interpretation. following the approach of @xcite , we compute the saliency score for each element of the attention matrix using a taylor expansion @xcite : ## conclusion in this work, we propose a novel peft method, gnnavi, leveraging gnn to navigate information flow within llms. specifically tailored for prompt-based fine-tuning, gnnavi significantly reduces the number of trainable parameters by simply adding a gnn layer into llms to guide the information flow within the prompt. gnnavi outperforms fpft and other peft methods across various classification tasks, even with few training examples. our work offers insights into handling llms from a graph perspective and presents a novel application of gnns in nlp. future work could explore different token connectivities for gnns or utilize gnns to control the information flow in llms.