cs.LGJul 16, 2025

Learning Task Mixtures from Task Affinities: A Probabilistic Graphical Model for Supervised Fine-Tuning

Authors: Prateek Chanda, Saral Sureka, Parth Pratim Chatterjee, Krishnateja Killamsetty, Nikhil Shivakumar Nayak, Ganesh Ramakrishnan

Organizations: IIT Bombay · IBM Research · Red Hat AI Innovation · MIT-IBM Watson AI Lab

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

Abstract

Supervised fine-tuning performance for large language models depends strongly on how training budget is distributed across a heterogeneous set of tasks. In practice, mixtures are often fixed using simple heuristics (e.g., uniform or size-proportional sampling) that ignore task interactions, which can hurt transfer and waste budget on redundant sources. We introduce TaskPGM, a framework for learning continuous task mixtures via an energy-based model over tasks. Tasks form the nodes of a Markov random field: unary potentials capture per-task utility, and pairwise potentials encode inter-task relationships using behavioral divergences computed from predictive distributions of single-task fine-tuned models (e.g., Jensen--Shannon divergence and pointwise mutual information). Optimizing this objective yields mixtures that balance coverage against redundancy. We show that the resulting set function is weakly submodular under budget constraints, enabling approximation guarantees for discrete selection variants. Across multiple model families (LLaMA-7B, Qwen2-7B) and evaluation suites (BIG-Bench Hard), TaskPGM improves over standard mixing strategies and provides interpretable structure over task interactions.

Figures & tables

Figure 1 : Illustrative SST-2 instance highlighting compositional negation.
Figure 2 : MRF over tasks with similarities SijS_{ij} (left) and the learned mixture probabilities pi∗p_{i}^{\ast} from Eq. ( 2 ) (right; shown the corr. probabilities).
Figure 3 : Accuracy surfaces over β/λ\beta/\lambda ratios during greedy mixture construction (LLama-7B) For each downstream benchmark, the surface plots evaluation accuracy as a function of the greedy step kk (mixture size) and the weight ratio β/λ\beta/\lambda , with β\beta and λ\lambda varied on a log scale, illustrating how performance curvature varies across tasks.
SizeMethodLeaderboard
  BBHGPQAIFEvalMathMMLU-ProMUSR
25KRandom0.4909±0.00620.4909_{\pm 0.0062}0.3188±0.13500.3188_{\pm 0.1350}0.3285±0.00810.3285_{\pm 0.0081}0.1881±0.01000.1881_{\pm 0.0100}0.4157±0.00450.4157_{\pm 0.0045}0.4603±0.01800.4603_{\pm 0.0180}
25KUniform0.5013±0.00620.5013_{\pm 0.0062}0.3314±0.01360.3314_{\pm 0.0136}0.2926±0.00740.2926_{\pm 0.0074}0.2085±0.01040.2085_{\pm 0.0104}0.4161±0.00450.4161_{\pm 0.0045}0.4683±0.01800.4683_{\pm 0.0180}
25KEPM0.4970±0.00620.4970_{\pm 0.0062}0.3314±0.01360.3314_{\pm 0.0136}0.3094±0.00770.3094_{\pm 0.0077}0.2068±0.01050.2068_{\pm 0.0105}0.4146±0.00450.4146_{\pm 0.0045}0.4537±0.01800.4537_{\pm 0.0180}
25KLESS0.5173±0.00620.5173_{\pm 0.0062}0.3020±0.00940.3020_{\pm 0.0094}0.3297±0.00800.3297_{\pm 0.0080}0.2002±0.01010.2002_{\pm 0.0101}0.4230±0.00450.4230_{\pm 0.0045}0.3598±0.01700.3598_{\pm 0.0170}
25KOurs (PMI)0.5202 ±0.00620.3341 ±0.01360.3909±0.00850.3909_{\pm 0.0085}0.2096±0.01010.2096_{\pm 0.0101}0.4251 ±0.00450.4701 ±0.0180
Table 1 : Qwen2-7B: Instruct-tuning performance on Leaderboard (25K and 50K samples).
SizeMethodLeaderboard
  BBHGPQAIFEvalMathMMLU-ProMUSR
25KRandom0.3482 ±0.00590.2626 ±0.01280.3465 ±0.00000.0098 ±0.00270.1877 ±0.00360.3677 ±0.0172
25KUniform0.3501 ±0.00590.2701 ±0.01290.3501 ±0.00000.0151 ±0.00340.1768 ±0.00350.4027 ±0.0175
25KEPM0.3593 ±0.00590.2601 ±0.01270.3405 ±0.00000.0151 ±0.00330.1836 ±0.00350.4286 ±0.0177
25KLESS0.4059±0.00550.4059_{\pm 0.0055}0.2412±0.00880.2412_{\pm 0.0088}0.3417±0.00000.3417_{\pm 0.0000}0.0001±0.00040.0001_{\pm 0.0004}0.1816±0.00620.1816_{\pm 0.0062}0.4157±0.01750.4157_{\pm 0.0175}
25KOurs (PMI)0.4095 ±0.00540.2718 ±0.01290.3561 ±0.00000.0159 ±0.00340.1924 ±0.00360.4298 ±0.0177
Table 2 : Llama2-7B: Instruct-tuning performance on Leaderboard (25K and 50K samples).
Appendix figures & tables5 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 4 : Comparison of task similarity metrics using cosine similarity (top), and PMI and JSD-based heatmaps (bottom). Cosine scores are generally low and fail to distinguish task structure. PMI highlights asymmetric task alignment. JSD offers symmetric, bounded divergence and reveals clearer task groupings across models.
Figure 5 : Task similarity heatmaps for Qwen models computed using Jensen–Shannon divergence (left) and pointwise mutual information (right).
Figure 6 : Eigenvalue spectra of similarity matrices derived from (a) Jensen–Shannon divergence (JSD) and (b) pointwise mutual information (PMI) . The PMI-based matrix exhibits a steeper spectral decay , indicating a lower effective rank and thus a more compact embedding of similarity relationships.
DatasetLeaderboard
Size / MethodBBHGPQAIFEvalMathMMLU-ProMUSR
25K
Random0.3482 ±0.00590.2626 ±0.01280.3465 ±N/A0.0098 ±0.00270.1877 ±0.00360.3677 ±0.0172
Uniform0.3501 ±0.00590.2701 ±0.01290.3501 ±N/A0.0151 ±0.00340.1768 ±0.00350.4027 ±0.0175
EPM0.3593 ±0.00590.2601 ±0.01270.3405 ±N/A0.0151 ±0.00330.1836 ±0.00350.4286 ±0.0177
LESS0.4059±0.00550.4059_{\pm 0.0055}0.2412±0.00880.2412_{\pm 0.0088}0.3417±N/A0.3417_{\pm N/A}0.0001±0.00040.0001_{\pm 0.0004}0.1816±0.00620.1816_{\pm 0.0062}0.4157±0.01750.4157_{\pm 0.0175}
Appendix
Table 3 : Llama-2-7b: Instruction-tuning performance on Leaderboard subsets with β=20,λ=10\beta=20,\lambda=10 using batch size 8.
Dataset Size(Method)Leaderboard
 BBHGPQAIFEvalMathMMLU-ProMUSR
25K (PMI)
β\beta =14954 ; λ\lambda =2630.3637 ±0.00590.2685 ±0.01280.3405 ±N/A0.0159 ±0.00340.1869 ±0.00360.3849 ±0.0173
β\beta =5273 ; λ\lambda =1950.3536 ±0.00590.2718 ±0.01290.3609 ±N/A0.0166 ±0.00350.1823 ±0.00350.4021 ±0.0174
β\beta =2535 ; λ\lambda =1960.3659 ±0.00590.2701 ±0.01290.3357 ±N/A0.0128 ±0.00310.1890 ±0.00360.3929 ±0.0173
β\beta =307 ; λ\lambda =600.3605 ±0.00590.2735 ±0.01290.3681 ±N/A0.0159 ±0.00340.1881 ±0.00360.4074 ±0.0174
Appendix
Table 4 : Llama-2-7b: Instruction-tuning performance on Leaderboard subsets with varying β\beta and λ\lambda using batch size 8.

Explore similar work

CardsList

  1. Curvature-Guided Mixing for MLLM Adaptation

    Jun 23, 2026Jinglong Yang, Jiaxuan He, Wenjian Huang +2Large Language Model AdaptationCurvature-Aware Spectral Framework

    Research Institute of Trustworthy Autonomous Systems and Department of Computer Science and Engineering, Southern University of Science and Technology · Department of Computer Science, City University of Hong Kong

  2. A helps B while B hurts A: directed transfer in instruction-tuning mixture

    Sep 30, 2026Nima H. Siboni, Vahid RostamiInstruction TuningMixtures

    Juna.ai, Kastanienallee 32, 10435 Berlin, Germany · Computational Systems Neuroscience, Institute of Zoology, University of Cologne, Germany

  3. PPL-Factory: Task-Aware and Budget-Aware Data Selection from Language Modeling to Reasoning

    Jul 20, 2026Hang Zhang, Warren J. GrossLarge Language Model Fine-TuningLanguage Modeling

    Department of Electrical and Computer Engineering, McGill University Montreal, QC, Canada