TLDR: This research introduces a two-stage, cost-efficient pipeline for In-Context Learning (ICL) that reduces reliance on large language models (LLMs) for data labeling. The first stage, GraphSim, uses graph-based structural information to select better cross-task examples for initial LLM-based pseudo-labeling of a small dataset. The second stage, GLIP, then employs graph-based label propagation to efficiently label the remaining target data without further LLM queries. This combined approach achieves strong ICL performance on novel tasks while significantly lowering labeling costs, outperforming traditional similarity-based methods and LLM-only pseudo-labeling.
Large Language Models, or LLMs, have shown remarkable abilities to learn new tasks without needing extensive retraining. This capability, known as In-Context Learning (ICL), allows LLMs to perform tasks by simply looking at a few examples of input-output pairs, called demonstrations, during their inference process.
However, a significant hurdle for ICL is the need for high-quality, labeled examples specific to the task at hand. For new or less common tasks, gathering these examples can be very expensive and time-consuming. Current methods often rely on LLMs to generate “pseudo-labels” for data, but LLMs can be unreliable on unfamiliar tasks. Alternatively, human annotation is accurate but costly in terms of time and labor.
To tackle this challenge, researchers have explored using examples from well-established, high-resource tasks (known as source tasks) to help LLMs learn novel, low-resource tasks (target tasks). This approach is called cross-task ICL. While promising, a key limitation is that data distributions and label categories often differ significantly between source and target tasks, making it difficult to select truly relevant examples based purely on text similarity.
A Novel Two-Stage Pipeline for Efficient In-Context Learning
A new research paper introduces a cost-efficient, two-stage pipeline designed to reduce the reliance on LLMs for data labeling, making ICL more practical for novel tasks. This innovative framework, detailed in the paper “From Cross-Task Examples to In-Task Prompts: A Graph-Based Pseudo-Labeling Framework for In-context Learning”, combines the flexibility of using examples from different tasks with the scalability of an LLM-free label propagation method.
The pipeline consists of two main components: GraphSim for cross-task example selection and GLIP for graph-based label information propagation.
GraphSim: Intelligent Cross-Task Example Selection
The first stage, GraphSim, focuses on selecting the most relevant examples from a source task to help an LLM pseudo-label a small initial set of target task instances. Unlike traditional methods that rely solely on text embedding similarity, GraphSim recognizes that structural patterns within data can generalize across different domains, even when the text content itself varies.
GraphSim builds task-specific graphs for both source and target datasets. In these graphs, each data sample is a node, and connections (edges) are formed between similar samples. The similarity is initially based on cosine similarity of text embeddings. To enrich these representations, GraphSim uses two types of aggregation:
- Adjacency-based aggregation: This captures information from multi-hop neighbors in the graph, allowing each node to gather broader structural context.
- GNN-based aggregation: A collection of randomly initialized Graph Neural Networks (GNNs) are used to capture diverse structural features without requiring extensive training. The idea is that nodes with similar structural roles will have similar aggregated representations.
By combining original text embeddings with these structure-aware embeddings, GraphSim computes a more robust cross-task similarity score. This allows for the selection of source examples that are more transferable and lead to higher-quality initial pseudo-labels for the target task.
GLIP: Scaling Labeling with Graph-Based Propagation
After GraphSim helps an LLM pseudo-label a small subset of the target data, the second stage, GLIP (Graph-based Label Information Propagation), takes over. The goal here is to efficiently spread these initial labels to the remaining unlabeled target examples without needing more expensive LLM queries.
GLIP constructs a specialized graph for the target task, where each query-choice pair in a multiple-choice question becomes a node. This design ensures that the labeling process is invariant to the order of choices. Two types of edges are defined:
- Similarity-based positive edges: These connect nodes from different questions that are semantically related, helping to propagate label information across the dataset.
- Mutual exclusion negative edges: These connect nodes within the same question, enforcing the constraint that only one choice can be correct. This prevents multiple options for a single question from being incorrectly labeled as correct.
A Graph Neural Network is then trained on this graph using a combination of standard cross-entropy loss for the initially labeled nodes and a mutual exclusion constraint loss for the negative edges. Once trained, this GNN can predict labels for all the previously unlabeled target examples, creating a fully pseudo-labeled dataset ready for in-task ICL.
Empirical Validation and Efficiency
The researchers conducted extensive experiments across five different tasks and various LLMs, including LLaMA2-7B, LLaMA2-13B, GPT-3.5, LLaMA3-8B, and GPT-4o. The results consistently showed that GraphSim outperformed baselines that relied solely on embedding similarity for cross-task example selection. Furthermore, GLIP demonstrated superior performance compared to LLM-based pseudo-labeling for the remaining unlabeled data, often approaching the performance of an “Oracle” model that had access to perfectly labeled data.
A key finding was the pipeline’s efficiency. While LLM-based labeling costs increase linearly with the amount of data, the graph-based propagation in GLIP significantly reduces the need for expensive LLM calls. This makes the entire framework much more practical and scalable for real-world applications, especially in resource-constrained environments.
The study also included ablation studies, confirming the importance of both adjacency-based and GNN-based aggregations in GraphSim. It also explored the impact of the number of in-context examples, showing that while cross-task ICL can degrade with too many examples due to distribution shifts, GraphSim maintains greater robustness. The robustness of the GLIP component to its hyperparameter settings was also demonstrated.
Also Read:
- Decoding How Pre-Training and Context Shape In-Context Learning
- AI Agents Get Smarter: A Graph-Based Approach to Understanding Complex Tools
Looking Ahead
This research offers a significant step forward in making in-context learning more accessible and efficient for novel tasks. By intelligently leveraging structural information in data and combining it with scalable graph-based label propagation, the proposed pipeline provides a powerful tool for generating high-quality demonstrations without heavy reliance on costly LLM inference. Future work could explore optimal source task selection, applicability to other task formats beyond multiple-choice QA, and mixed-task in-context learning scenarios.
