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

arXiv: https://arxiv.org/abs/2507.12612 · HTML: https://mlsanity.com/paper/2507.12612/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 28 of its 40 references that are in the map.
- First posted on arXiv 16 Jul 2025.

## Builds on (most cited first, 28 of 28)

- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [LLaMA: Open and Efficient Foundation Language Models](https://mlsanity.com/paper/2302.13971/) (2023)
- [Chain-of-Thought Prompting Elicits Reasoning in Large Language Models](https://mlsanity.com/paper/2201.11903/) (2022)
- [GLUE: A Multi-Task Benchmark and Analysis Platform for Natural Language Understanding](https://mlsanity.com/paper/1804.07461/) (2018)
- [Scaling Instruction-Finetuned Language Models](https://mlsanity.com/paper/2210.11416/) (2022)
- [Convex optimization](https://mlsanity.com/paper/2106.01946/) (2021)
- [Understanding Black-box Predictions via Influence Functions](https://mlsanity.com/paper/1703.04730/) (2017)
- [SuperGLUE: A Stickier Benchmark for General-Purpose Language Understanding Systems](https://mlsanity.com/paper/1905.00537/) (2019)
- [Multitask Prompted Training Enables Zero-Shot Task Generalization](https://mlsanity.com/paper/2110.08207/) (2021)
- [Tulu 3: Pushing Frontiers in Open Language Model Post-Training](https://mlsanity.com/paper/2411.15124/) (2024)
- [Super-NaturalInstructions: Generalization via Declarative Instructions on 1600+ NLP Tasks](https://mlsanity.com/paper/2204.07705/) (2022)
- [The Flan Collection: Designing Data and Methods for Effective Instruction Tuning](https://mlsanity.com/paper/2301.13688/) (2023)
- [An Empirical Study of Example Forgetting during Deep Neural Network Learning](https://mlsanity.com/paper/1812.05159/) (2018)
- [Deep Learning on a Data Diet: Finding Important Examples Early in Training](https://mlsanity.com/paper/2107.07075/) (2021)
- [LESS: Selecting Influential Data for Targeted Instruction Tuning](https://mlsanity.com/paper/2402.04333/) (2024)
- [Task Arithmetic in the Tangent Space: Improved Editing of Pre-Trained Models](https://mlsanity.com/paper/2305.12827/) (2023)
- [Task2Vec: Task Embedding for Meta-Learning](https://mlsanity.com/paper/1902.03545/) (2019)
- [RegMix: Data Mixture as Regression for Language Model Pre-training](https://mlsanity.com/paper/2407.01492/) (2024)
- [Geometric Dataset Distances via Optimal Transport](https://mlsanity.com/paper/2002.02923/) (2020)
- [Coverage-centric Coreset Selection for High Pruning Rates](https://mlsanity.com/paper/2210.15809/) (2022)
- [RETRIEVE: Coreset Selection for Efficient and Robust Semi-Supervised Learning](https://mlsanity.com/paper/2106.07760/) (2021)
- [D2 Pruning: Message Passing for Balancing Diversity and Difficulty in Data Pruning](https://mlsanity.com/paper/2310.07931/) (2023)
- [TSDS: Data Selection for Task-Specific Model Finetuning](https://mlsanity.com/paper/2410.11303/) (2024)
- [Improving Influence-based Instruction Tuning Data Selection for Balanced Learning of Diverse Capabilities](https://mlsanity.com/paper/2501.12147/) (2025)
- [SMART: Submodular Data Mixture Strategy for Instruction Tuning](https://mlsanity.com/paper/2403.08370/) (2024)
- [DELIFT: Data Efficient Language model Instruction Fine Tuning](https://mlsanity.com/paper/2411.04425/) (2024)
- [Speculative Coreset Selection for Task-Specific Fine-tuning](https://mlsanity.com/paper/2410.01296/) (2024)
- [A Single Linear Layer Yields Task-Adapted Low-Rank Matrices](https://mlsanity.com/paper/2403.14946/) (2024)

## Built on by (most cited first, 1 of 1)

- [A helps B while B hurts A: directed transfer in instruction-tuning mixture](https://mlsanity.com/paper/2609.39702/) (2026)

## Abstract (arXiv)

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.