TEGEE: Task dEfinition Guided Expert Ensembling for Generalizable and Few-shot Learning

TL;DR

TEGEE tackles the core bottleneck of in-context learning by explicitly extracting task definitions and guiding expert ensembling. It deploys a dual 3B-model design—one for task-definition extraction and one for learning from demonstrations—augmented with a LoRA-based dynamic expert pool and a retrieval-enforced weighting scheme. Empirical results on SuperNI show TEGEE achieving performance on par with LLaMA2-13B and outperforming 7B baselines, while enabling continual few-shot learning through continual pool augmentation. The work highlights the primacy of task-definition extraction in ICL and demonstrates how modular, task-aware ensembling can extend few-shot learning to many-shot regimes with practical continual learning benefits.

Abstract

Large Language Models (LLMs) exhibit the ability to perform in-context learning (ICL), where they acquire new tasks directly from examples provided in demonstrations. This process is thought to operate through an implicit task selection mechanism that involves extracting and processing task definitions from these demonstrations. However, critical questions remain: Which is more essential -- task extraction or definition? And how can these capabilities be further improved? To address these questions, we propose \textbf{TEGEE} (Task Definition Guided Expert Ensembling), a method that explicitly extracts task definitions and generates responses based on specific tasks. Our framework employs a dual 3B model approach, with each model assigned a distinct role: one focuses on task definition extraction, while the other handles learning from demonstrations. This modular approach supports the hypothesis that extracting task definitions is more vital than processing the task itself. Empirical evaluations show that TEGEE performs comparably to the larger LLaMA2-13B model. By leveraging a modular design, our approach extends traditional ICL from few-shot to many-shot learning, supporting an unlimited number of demonstrations and enhancing continual learning capabilities.

Figures (14)

• Figure 1: Pipeline of TEGEE. The preparation stage focuses on generating a pool of task-based experts, each trained using task-specific data. ICL firstly generates task definitions from given demonstrations using a task definition generator, which can be a custom generator, GPT, or human. Secondly, we utilize the task definitions to retrieve the task-based expert pool and match the top three experts that are most similar. Next, the expert ensemble step initializes the fusion expert by performing a weighted average of the matched experts' weights. Finally, the data format is reconstructed for training the fusion expert and performing inference on the last input field.
• Figure 2: Performance comparison of TEGEE across four distinct expert base model sizes and five task definition extractors with 5 demonstration inputs. The expert base models include LLaMA2-7B and LLaMA2-13B, and Falcon7B. Baseline ICL results are denoted by dashed lines. Notably, our paradigm employing a 6B parameters demonstrates superior performance to the LLaMA2-7B ICL and exhibits comparable results to the LLaMA2-13B ICL.
• Figure 3: The Performances trend alongside increases in the number of demonstrations.
• Figure 4: GPT3 & GPT-3.5
• Figure 5: GPT3 & GPT4