Framework optimizes language model training by learning which task combinations work best together

Distributing training effort across multiple tasks during supervised fine-tuning of large language models is a surprisingly difficult problem. Most practitioners rely on simple rules—sampling uniformly or proportional to dataset size—despite evidence that these strategies miss important interactions between tasks and squander computational resources on overlapping or redundant material.

The optimization challenge#

Researchers have long known that task relationships matter. Some tasks reinforce one another and boost overall model performance, while others are essentially redundant or compete for the same underlying skills. Yet existing approaches to task weighting lack a principled way to account for these dynamics.

A team led by Prateek Chanda proposes TaskPGM, a framework that treats task selection as a probabilistic graphical model problem. Tasks become nodes in a Markov random field, where unary potentials measure individual task utility and pairwise potentials capture how tasks relate to each other using behavioral divergences—metrics like Jensen-Shannon divergence and pointwise mutual information drawn from the predictive distributions of single-task fine-tuned models.

How it works#

The model balances two competing objectives: ensuring broad coverage of skill areas while minimizing redundancy. This framing allows the researchers to prove that the resulting objective function is weakly submodular, a property that guarantees approximation bounds for discrete variants of the problem. The framework was tested across multiple model families, including LLaMA-7B and Qwen2-7B, using the BIG-Bench Hard evaluation suite.

Results and interpretability#

Results show that TaskPGM-derived task mixtures outperform standard baseline strategies in fine-tuning performance. Beyond raw numbers, the framework produces interpretable structure—practitioners can visualize which tasks complement each other and which create redundancy, offering insight into model learning dynamics that black-box assignment methods cannot provide.

Note: This paper has been withdrawn by the author from arXiv.