Saved in:
| Main Authors: | , , , , , |
|---|---|
| Format: | Preprint |
| Published: |
2025
|
| Subjects: | |
| Online Access: | https://arxiv.org/abs/2507.12612 |
| Tags: |
Add Tag
No Tags, Be the first to tag this record!
|
| _version_ | 1866913978573127680 |
|---|---|
| author | Chanda, Prateek Sureka, Saral Chatterjee, Parth Pratim Killamsetty, Krishnateja Nayak, Nikhil Shivakumar Ramakrishnan, Ganesh |
| author_facet | Chanda, Prateek Sureka, Saral Chatterjee, Parth Pratim Killamsetty, Krishnateja Nayak, Nikhil Shivakumar Ramakrishnan, Ganesh |
| contents | The performance of finetuned large language models (LLMs) hinges critically on the composition of the training mixture. However, selecting an optimal blend of task datasets remains a largely manual, heuristic driven process, with practitioners often relying on uniform or size based sampling strategies. We introduce TASKPGM, a principled and scalable framework for mixture optimization that selects continuous task proportions by minimizing an energy function over a Markov Random Field (MRF). Task relationships are modeled using behavioral divergences such as Jensen Shannon Divergence and Pointwise Mutual Information computed from the predictive distributions of single task finetuned models. Our method yields a closed form solution under simplex constraints and provably balances representativeness and diversity among tasks. We provide theoretical guarantees, including weak submodularity for budgeted variants, and demonstrate consistent empirical improvements on Llama 2 and Mistral across evaluation suites such as MMLU and BIGBench. Beyond performance, TASKPGM offers interpretable insights into task influence and mixture composition, making it a powerful tool for efficient and robust LLM finetuning. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2507_12612 |
| institution | arXiv |
| publishDate | 2025 |
| record_format | arxiv |
| spellingShingle | Learning What Matters: Probabilistic Task Selection via Mutual Information for Model Finetuning Chanda, Prateek Sureka, Saral Chatterjee, Parth Pratim Killamsetty, Krishnateja Nayak, Nikhil Shivakumar Ramakrishnan, Ganesh Machine Learning Artificial Intelligence 68T50 I.2.7; I.2.6; I.2.4 The performance of finetuned large language models (LLMs) hinges critically on the composition of the training mixture. However, selecting an optimal blend of task datasets remains a largely manual, heuristic driven process, with practitioners often relying on uniform or size based sampling strategies. We introduce TASKPGM, a principled and scalable framework for mixture optimization that selects continuous task proportions by minimizing an energy function over a Markov Random Field (MRF). Task relationships are modeled using behavioral divergences such as Jensen Shannon Divergence and Pointwise Mutual Information computed from the predictive distributions of single task finetuned models. Our method yields a closed form solution under simplex constraints and provably balances representativeness and diversity among tasks. We provide theoretical guarantees, including weak submodularity for budgeted variants, and demonstrate consistent empirical improvements on Llama 2 and Mistral across evaluation suites such as MMLU and BIGBench. Beyond performance, TASKPGM offers interpretable insights into task influence and mixture composition, making it a powerful tool for efficient and robust LLM finetuning. |
| title | Learning What Matters: Probabilistic Task Selection via Mutual Information for Model Finetuning |
| topic | Machine Learning Artificial Intelligence 68T50 I.2.7; I.2.6; I.2.4 |
| url | https://arxiv.org/abs/2507.12612 |