Technical Summary: TaskPGM – Learning Task Mixtures from Task Affinities

Problem Statement

Supervised Fine-Tuning (SFT) for Large Language Models (LLMs) relies heavily on the distribution of the training budget across a heterogeneous set of instruction tasks. Current practices often rely on simple heuristics, such as uniform sampling or dataset-size proportional weighting. These methods fail to account for task interactions, leading to two primary inefficiencies:

1. Redundancy: Over-weighting tasks that induce nearly identical behavioral updates, wasting computational budget.
2. Missed Transfer: Under-utilizing tasks that offer complementary information or cross-domain transfer benefits.

Existing automated approaches (e.g., influence-based retrieval like LESS, or submodular selection like SMART) often optimize for discrete subsets rather than continuous mixture proportions, rely on semantic similarity that ignores downstream model behavior, or require expensive iterative training. The core challenge is to automatically determine a continuous task mixture that is simultaneously representative (high utility) and diverse (low redundancy) to maximize downstream performance under a fixed budget.

Methodology: TaskPGM

The authors propose TaskPGM (Task Mixtures via Probabilistic Graphical Modelling), a framework that casts SFT data mixing as an energy minimization problem over a dense Markov Random Field (MRF).

1. Probabilistic Graphical Model Formulation

Tasks are modeled as nodes in an MRF. The goal is to find an optimal sampling distribution $p^* \in \Delta_n$ (where $\Delta_n$ is the probability simplex) that maximizes an energy function $E(p)$. The energy function consists of two components:

• Unary Potentials ($\Psi_{un}$): Encode the utility or representativeness of individual tasks. This is derived from the total similarity mass of a task relative to the entire collection.
• Pairwise Potentials ($\Psi_{pair}$): Encode redundancy and diversity. These terms penalize the co-occurrence of highly similar tasks in the mixture.

The optimization objective is formulated as:
$$\max_{p \in \Delta_n} E(p) = \Psi_{un}^\top p - \frac{1}{2} p^\top \Psi_{pair} p$$
This is equivalent to minimizing a convex quadratic function (under specific conditions) subject to simplex constraints.

2. Behavioral Task Affinity Metrics

A key innovation is the method for computing pairwise task affinities ($S_{ij}$). Instead of relying on semantic similarity (embeddings) or metadata, TaskPGM computes affinities based on behavioral divergences derived from the predictive distributions of single-task fine-tuned models.

• Pointwise Mutual Information (PMI): Measures the mutual information between the predictive distributions of two tasks when evaluated on each other's datasets.
• Jensen–Shannon Divergence (JSD): A symmetric measure of the discrepancy between the predictive distributions of two models on cross-task examples.

These metrics capture how tasks shape model behavior on out-of-task examples, directly addressing cross-task transfer and redundancy.

3. Optimization and Solvers

• Continuous Mixture: The authors derive a closed-form optimum for the continuous mixture weights using Karush-Kuhn-Tucker (KKT) conditions.
  • If the pairwise similarity matrix is not Positive Semi-Definite (PSD), they apply spectrum shifting (adding a constant to the diagonal) to ensure convexity, which introduces a regularization term biasing the solution toward uniformity if the non-PSD structure is severe.
  • When the inverse is ill-conditioned or non-negativity constraints are violated, they employ Projected Gradient Descent (PGD) on the probability simplex.
• Discrete/Budgeted Variant: For scenarios requiring the selection of a subset of tasks (top-$k$), the authors lift the continuous optimization to a discrete set function. They prove this function is monotonic and weakly submodular under standard regularity conditions (Restricted Strong Convexity and Smoothness), enabling efficient greedy construction with approximation guarantees.

Key Contributions

1. Probabilistic Framework for Mixture Design: A novel formulation of SFT data mixing as energy minimization on an MRF, explicitly trading off representativeness (unary terms) and diversity (pairwise terms).
2. Behavior-Based Affinity Metrics: The proposal to compute task relationships using predictive distribution divergences (JSD/PMI) from single-task fine-tuned models, moving beyond static semantic similarity to capture functional task dependencies.
3. Theoretically Backed Optimization:
   • Derivation of a closed-form solution for continuous mixtures via KKT conditions.
   • Establishment of weak submodularity for the discrete budgeted variant, supporting greedy selection with theoretical guarantees.
4. Empirical Validation: Demonstration of consistent performance gains across multiple model families (Qwen2-7B, Llama-2-7B) and benchmarks (BIG-Bench Hard, MMLU, etc.) compared to heuristics and state-of-the-art selection baselines.

Experimental Results

The framework was evaluated on 309 instruction tasks using Qwen2-7B and Llama-2-7B models with budgets of 25K and 50K instances.

• Performance Gains: TaskPGM (using both PMI and JSD variants) consistently outperformed uniform sampling, size-proportional (EPM), and random baselines.
  • On Qwen2-7B (25K budget), the PMI variant improved performance by nearly 3% over heuristic baselines on complex reasoning benchmarks.
  • On Llama-2-7B, the PMI-derived mixture outperformed Uniform and EPM by 2% to over 4% on BBH.
• Comparison with SOTA: TaskPGM remained competitive with or superior to LESS (a strong influence-function based baseline). Notably, TaskPGM achieves this with a single general-purpose checkpoint, whereas LESS constructs distinct, target-aware models for each benchmark.
• Efficiency: TaskPGM is computationally more efficient than LESS. The similarity computation is a one-time amortized cost (~16–34 GPU-hours), after which mixtures for any budget can be derived in seconds via the closed-form solution. In contrast, LESS requires ~320 GPU-hours to build gradient datastores for each target.
• Scalability: Increasing the budget from 25K to 50K yielded consistent gains for TaskPGM, whereas heuristic methods (Uniform, EPM) exhibited fragmented and non-robust behavior as budgets grew.

Significance and Claims

The paper claims that TaskPGM provides a principled, interpretable, and efficient approach to SFT data mixing.

• Interpretability: The learned weights and similarity heatmaps reveal task clusters, redundancy, and "attractor" tasks, offering insight into why certain sources help.
• Generalization: Unlike methods that optimize for specific targets, TaskPGM produces a single mixture that improves performance across a broad spectrum of downstream tasks (reasoning, language understanding, knowledge).
• Robustness: The gains are driven by functional task selection based on model behavior rather than favorable sampling variance, as evidenced by consistent improvements across different base models and similarity metrics.

The authors conclude that principled, similarity-driven mixing is essential for scalable instruction tuning, often matching or exceeding heuristic mixtures with substantially less data. They acknowledge limitations regarding the fixed pool of tasks and the computational cost of affinity estimation, suggesting scalable approximations as a direction for future work.