TaskPGM: Task Probabilistic Graphical Model

• TaskPGM is a framework that optimizes LLM finetuning mixtures by modeling tasks as nodes in a Markov Random Field (MRF).
• It leverages unary potentials for task usefulness and pairwise potentials to penalize redundancy, solving a convex quadratic program with clear theoretical guarantees.
• Empirical results show that with larger budgets, TaskPGM can outperform uniform or dataset-size baselines on benchmarks like MMLU, demonstrating practical gains.

TaskPGM, short for Task Probabilistic Graphical Model, is a framework for optimizing the task mixture used in LLM finetuning by selecting continuous sampling proportions over tasks rather than relying on heuristic recipes such as uniform sampling or sampling proportional to dataset size. It models tasks as nodes in a dense Markov Random Field (MRF), defines unary potentials that reward representativeness and pairwise potentials that penalize redundancy, and derives the resulting mixture by minimizing an energy over the probability simplex. The method constructs task affinities from behavioral divergences between single-task finetuned models, notably Jensen–Shannon Divergence (JSD) and a PMI-style score computed from predictive distributions. In this formulation, mixture design becomes a quadratic program with linear constraints, with a closed-form interior solution and a budgeted variant that is monotone and weakly submodular (Chanda et al., 16 Jul 2025).

1. Definition and scope

TaskPGM addresses the problem of determining, for a collection of task datasets, what continuous task proportions should be used to form a finetuning mixture. The framework is explicitly task-level: it optimizes a probability vector over tasks, then samples instances within each task according to that optimized distribution. This distinguishes it from uniform allocation across tasks, size-based allocation proportional to dataset cardinality, and random sampling from the union of all instances (Chanda et al., 16 Jul 2025).

The central design choice is to treat each task as a node in a dense undirected MRF. The node-level, or unary, component rewards tasks that are broadly useful to other tasks, while the edge-level, or pairwise, component penalizes simultaneously assigning large weight to tasks that are behaviorally similar. This creates a formal trade-off between representativeness and diversity. The paper states that minimizing this energy over the simplex yields an optimal continuous mixture $\mathbf{p}^\ast$, and that a single ratio $\beta/\lambda$ controls the balance between attraction to high-utility tasks and repulsion among redundant ones (Chanda et al., 16 Jul 2025).

The framework is intended for LLM finetuning mixtures rather than for example-level data curation alone. Its empirical evaluation is conducted on Llama-2-7B and Mistral-7B-v0.3, using candidate tasks drawn from Flan 2021, T0, Chain-of-Thought data, Tulu v3, and GLUE / SuperGLUE, with 316 tasks selected as the pool for mixture optimization (Chanda et al., 16 Jul 2025).

2. Formal task-mixture model

Let the candidate task set be

$\mathcal{T} = [T_1, T_2, \dots, T_n],$

with associated datasets $\mathcal{D}_i$. A task mixture is a probability vector

$\mathbf{p} = (p_1,\dots,p_n)^\top \in \Delta_n,$

where

$$\Delta_n = \left\{ \mathbf{p} \in \mathbb{R}^n \mid \mathbf{p}^\top \mathbf{1}_n = 1,\ \mathbf{p}\ge 0 \right\}.$$

The corresponding assignment tuple is

$$\boldsymbol{\Pi}_n := \left\{ \langle T_i, p_{[i]} \rangle \right\}_{i=1}^n.$$

Task relations are encoded in a symmetric similarity matrix $\mathbf{S}\in\mathbb{R}^{n\times n}$, with total similarity mass

$s_i := \sum_{j=1}^n S_{ij}.$

Tasks with large $s_i$ are therefore globally similar or connected to many others (Chanda et al., 16 Jul 2025).

The MRF assigns unary and pairwise potentials as

$\beta/\lambda$0

where $\beta/\lambda$1 and $\beta/\lambda$2 are scalar weights and $\beta/\lambda$3 is the graph Laplacian built from $\beta/\lambda$4. Collecting these into a unary vector $\beta/\lambda$5 and pairwise matrix $\beta/\lambda$6, the energy is

$\beta/\lambda$7

The optimization problem is

$\beta/\lambda$8

The negative linear term favors tasks with large unary mass, while the quadratic Laplacian term penalizes concentration on neighboring, hence similar, tasks (Chanda et al., 16 Jul 2025).

This decomposition gives the framework its main interpretation. The unary part is a representativeness term: a task that is behaviorally central in the task graph is encouraged. The pairwise part is a diversity term: it discourages allocating high probability to multiple tasks that lie in the same similarity neighborhood. The paper explicitly notes that large $\beta/\lambda$9 pulls the solution toward representative hubs, while small $\mathcal{T} = [T_1, T_2, \dots, T_n],$0 promotes a more spread-out allocation (Chanda et al., 16 Jul 2025).

3. Behavioral similarity and task affinities

TaskPGM defines task relationships through model behavior rather than through metadata or text embeddings. For each task $\mathcal{T} = [T_1, T_2, \dots, T_n],$1, a single-task finetuned model is obtained as

$\mathcal{T} = [T_1, T_2, \dots, T_n],$2

where $\mathcal{T} = [T_1, T_2, \dots, T_n],$3 is the base LLM and $\mathcal{T} = [T_1, T_2, \dots, T_n],$4 is the task vector. Let $\mathcal{T} = [T_1, T_2, \dots, T_n],$5 denote the predictive distribution (Chanda et al., 16 Jul 2025).

One similarity construction is a PMI-style score between tasks $\mathcal{T} = [T_1, T_2, \dots, T_n],$6 and $\mathcal{T} = [T_1, T_2, \dots, T_n],$7:

$\mathcal{T} = [T_1, T_2, \dots, T_n],$8

The stated intuition is that if the model trained on $\mathcal{T} = [T_1, T_2, \dots, T_n],$9 assigns higher probability to correct labels on $\mathcal{D}_i$0’s data than the model trained on $\mathcal{D}_i$1 itself, then the similarity increases; the same logic is applied symmetrically in the reverse direction (Chanda et al., 16 Jul 2025).

A second construction is based on Jensen–Shannon Divergence. For a sample $\mathcal{D}_i$2, define

$\mathcal{D}_i$3

with per-sample divergence

$\mathcal{D}_i$4

The paper then symmetrizes across both directions:

$\mathcal{D}_i$5

JSD is described as symmetric, bounded in $\mathcal{D}_i$6, and well-behaved even when supports differ. The framework may then transform JSD into a similarity before constructing the Laplacian (Chanda et al., 16 Jul 2025).

This behavioral construction is a defining feature of TaskPGM. Similarity is not inferred from dataset names, domains, or latent task metadata, but from how separately finetuned models behave on one another’s data. A plausible implication is that the resulting graph is closer to transfer structure than to surface-form resemblance, which the paper connects to the visibility of task clusters in PMI and JSD heatmaps (Chanda et al., 16 Jul 2025).

4. Optimization, convexity, and theoretical guarantees

The objective

$\mathcal{D}_i$7

is a convex quadratic over the simplex when $\mathcal{D}_i$8 is positive semidefinite. The paper states that the similarity-based Laplacian is PSD, but also introduces spectral shifting for more general pairwise constructions:

$\mathcal{D}_i$9

This guarantees convexity, although the paper notes that a large shift can bias solutions toward more uniform mixtures (Chanda et al., 16 Jul 2025).

Using the Lagrangian

$\mathbf{p} = (p_1,\dots,p_n)^\top \in \Delta_n,$0

the Karush–Kuhn–Tucker conditions yield, in the interior case $\mathbf{p} = (p_1,\dots,p_n)^\top \in \Delta_n,$1, the stationarity equation

$\mathbf{p} = (p_1,\dots,p_n)^\top \in \Delta_n,$2

If $\mathbf{p} = (p_1,\dots,p_n)^\top \in \Delta_n,$3 is invertible, the closed-form interior solution is

$\mathbf{p} = (p_1,\dots,p_n)^\top \in \Delta_n,$4

When some coordinates become negative, the paper states that one may restrict to the support where $\mathbf{p} = (p_1,\dots,p_n)^\top \in \Delta_n,$5 and re-solve, or project onto the simplex with standard projection algorithms (Chanda et al., 16 Jul 2025).

TaskPGM also introduces a budgeted or top-$\mathbf{p} = (p_1,\dots,p_n)^\top \in \Delta_n,$6 variant. For a subset $\mathbf{p} = (p_1,\dots,p_n)^\top \in \Delta_n,$7 with support restricted to $\mathbf{p} = (p_1,\dots,p_n)^\top \in \Delta_n,$8, define the utility

$\mathbf{p} = (p_1,\dots,p_n)^\top \in \Delta_n,$9

where $$\Delta_n = \left\{ \mathbf{p} \in \mathbb{R}^n \mid \mathbf{p}^\top \mathbf{1}_n = 1,\ \mathbf{p}\ge 0 \right\}.$$0. Over the family

$$\Delta_n = \left\{ \mathbf{p} \in \mathbb{R}^n \mid \mathbf{p}^\top \mathbf{1}_n = 1,\ \mathbf{p}\ge 0 \right\}.$$1

the discrete problem is $$\Delta_n = \left\{ \mathbf{p} \in \mathbb{R}^n \mid \mathbf{p}^\top \mathbf{1}_n = 1,\ \mathbf{p}\ge 0 \right\}.$$2. The paper proves that this set function is monotonic and weakly submodular, with submodularity ratio

$$\Delta_n = \left\{ \mathbf{p} \in \mathbb{R}^n \mid \mathbf{p}^\top \mathbf{1}_n = 1,\ \mathbf{p}\ge 0 \right\}.$$3

using Restricted Strong Convexity and Restricted Smoothness of the quadratic form. The practical consequence is that greedy task discovery under a cardinality constraint admits approximation guarantees analogous to those used in submodular optimization, although exact submodularity is not claimed (Chanda et al., 16 Jul 2025).

5. Training pipeline and implementation

The implementation described in the paper follows a six-stage pipeline. First, for each candidate task $$\Delta_n = \left\{ \mathbf{p} \in \mathbb{R}^n \mid \mathbf{p}^\top \mathbf{1}_n = 1,\ \mathbf{p}\ge 0 \right\}.$$4, a single-task model $$\Delta_n = \left\{ \mathbf{p} \in \mathbb{R}^n \mid \mathbf{p}^\top \mathbf{1}_n = 1,\ \mathbf{p}\ge 0 \right\}.$$5 is trained. In the core experiments this uses full-parameter finetuning, with no adapters, for 3 epochs, effective batch size 64, and 8 H100 GPUs. Second, predictive distributions are computed on each task’s own dataset and on other tasks’ datasets. Third, all pairwise task similarities are computed using either the PMI-style construction or the JSD-based construction; the paper emphasizes caching self-distributions and cross-distributions and parallelizing the $$\Delta_n = \left\{ \mathbf{p} \in \mathbb{R}^n \mid \mathbf{p}^\top \mathbf{1}_n = 1,\ \mathbf{p}\ge 0 \right\}.$$6 pair evaluations. Fourth, the graph Laplacian and MRF potentials are constructed, with optional spectral correction. Fifth, the optimized task proportions $$\Delta_n = \left\{ \mathbf{p} \in \mathbb{R}^n \mid \mathbf{p}^\top \mathbf{1}_n = 1,\ \mathbf{p}\ge 0 \right\}.$$7 are obtained from the quadratic program. Sixth, an instance budget $$\Delta_n = \left\{ \mathbf{p} \in \mathbb{R}^n \mid \mathbf{p}^\top \mathbf{1}_n = 1,\ \mathbf{p}\ge 0 \right\}.$$8 is allocated by sampling

$$\Delta_n = \left\{ \mathbf{p} \in \mathbb{R}^n \mid \mathbf{p}^\top \mathbf{1}_n = 1,\ \mathbf{p}\ge 0 \right\}.$$9

after which $$\boldsymbol{\Pi}_n := \left\{ \langle T_i, p_{[i]} \rangle \right\}_{i=1}^n.$$0 instances are drawn uniformly from task $$\boldsymbol{\Pi}_n := \left\{ \langle T_i, p_{[i]} \rangle \right\}_{i=1}^n.$$1’s dataset and the base LLM is finetuned on the resulting mixture (Chanda et al., 16 Jul 2025).

For mixture finetuning, the reported setup uses 1 epoch over the sampled mixture, with effective batch sizes 8 or 64 depending on the experiment, learning rate $$\boldsymbol{\Pi}_n := \left\{ \langle T_i, p_{[i]} \rangle \right\}_{i=1}^n.$$2, AdamW, weight decay 0.01, and gradient checkpointing. The evaluated base models are Llama-2-7B and Mistral-7B-v0.3 in bf16 (Chanda et al., 16 Jul 2025).

The task pool is assembled from several instruction-tuning sources: Flan 2021, T0, Chain-of-Thought data, Tulu v3, and GLUE / SuperGLUE. From these, 316 tasks are selected for optimization. The paper notes that the dominant computational expense is the $$\boldsymbol{\Pi}_n := \left\{ \langle T_i, p_{[i]} \rangle \right\}_{i=1}^n.$$3 single-task finetunes, with $$\boldsymbol{\Pi}_n := \left\{ \langle T_i, p_{[i]} \rangle \right\}_{i=1}^n.$$4 in the experiments, whereas similarity-matrix construction and the quadratic solve are comparatively light, especially once single-task models and predictive distributions have been cached (Chanda et al., 16 Jul 2025).

This implementation makes the framework scalable at the level of a few hundred tasks. The paper also notes the limitation that pairwise similarity construction is quadratic in the number of tasks, which may become problematic when the task pool grows to the thousands (Chanda et al., 16 Jul 2025).

6. Empirical behavior, interpretability, and limitations

Empirically, TaskPGM is compared against three baselines: Random, which samples $$\boldsymbol{\Pi}_n := \left\{ \langle T_i, p_{[i]} \rangle \right\}_{i=1}^n.$$5 instances uniformly from the union of all instances; Uniform, which allocates budget equally across tasks; and EPM, or Examples Proportional Mixing, which allocates budget proportional to task dataset size. Evaluation is performed on MMLU and on Open LLM Leaderboard subsets including BBH, GPQA, IFEval, Math, MMLU-Pro, and MUSR (Chanda et al., 16 Jul 2025).

The reported results are heterogeneous but informative. For Llama-2-7B with batch size 8 and $$\boldsymbol{\Pi}_n := \left\{ \langle T_i, p_{[i]} \rangle \right\}_{i=1}^n.$$6K, MMLU is 0.3913 for Random, 0.3479 for Uniform, 0.3802 for EPM, 0.4242 for TaskPGM with PMI, and 0.3926 for TaskPGM with JSD. For Llama-2-7B with $$\boldsymbol{\Pi}_n := \left\{ \langle T_i, p_{[i]} \rangle \right\}_{i=1}^n.$$7K, the corresponding MMLU scores are 0.4108, 0.3725, 0.3801, 0.4056, and 0.4074, respectively. For Mistral-7B with batch size 8 and $$\boldsymbol{\Pi}_n := \left\{ \langle T_i, p_{[i]} \rangle \right\}_{i=1}^n.$$8K, the paper reports a counterexample to uniform dominance: Random achieves 0.4539 on MMLU, Uniform 0.4376, EPM 0.4364, TaskPGM with PMI 0.3903, and TaskPGM with JSD 0.3783. However, for Mistral-7B with $$\boldsymbol{\Pi}_n := \left\{ \langle T_i, p_{[i]} \rangle \right\}_{i=1}^n.$$9K and batch size 8, the ranking reverses sharply: Random is 0.4476, Uniform 0.4486, EPM 0.4505, TaskPGM with PMI 0.5476, and TaskPGM with JSD 0.5301 (Chanda et al., 16 Jul 2025).

These results support two technical points. First, TaskPGM is not uniformly superior in every small-budget regime, and the paper explicitly records cases where heuristic mixtures perform better. Second, with larger budgets and suitable training configurations, the optimized mixture can substantially outperform Uniform and EPM; the paper also states that gains of 4+ percentage points on MMLU and MUSR are common in its broader evaluations (Chanda et al., 16 Jul 2025).

The ablations distinguish PMI and JSD as two different views of task geometry. PMI-based similarity is reported to have sharper eigenvalue decay and lower effective rank, which the paper associates with strong performance on global benchmarks such as MMLU. JSD-based similarity is described as having a more gradual spectrum and can dominate on specialized tasks such as GPQA and IFEval. Hyperparameter sweeps over $\mathbf{S}\in\mathbb{R}^{n\times n}$0 further indicate that very high ratios can favor MMLU while hurting leaderboard tasks, whereas very low ratios approach near-uniform mixtures and lose structural benefit; a moderate ratio such as $\mathbf{S}\in\mathbb{R}^{n\times n}$1 is reported as robust in ablations, while the main experiments fix $\mathbf{S}\in\mathbb{R}^{n\times n}$2 (Chanda et al., 16 Jul 2025).

TaskPGM is also presented as an interpretability tool. The optimized probabilities $\mathbf{S}\in\mathbb{R}^{n\times n}$3 can be read directly as task importances. The paper compares cosine similarity of task vectors with PMI and JSD heatmaps, arguing that cosine similarity is noisy whereas PMI and JSD reveal clearer task blocks. For task discovery, it introduces total variation distance between successive mixtures,

$\mathbf{S}\in\mathbb{R}^{n\times n}$4

to quantify how the addition of a new task reshapes the mixture. Adding high-unary tasks induces large TV jumps, while adding low-unary tasks produces more gradual change (Chanda et al., 16 Jul 2025).

The stated limitations are consequential. The approach requires a single-task finetuned model for each candidate task, which is expensive even if embarrassingly parallel. The quality of the result depends on the similarity metric; PMI and JSD work well empirically, but other objectives may require other metrics. Spectral correction can introduce an unwanted uniform bias. The optimized mixture inherits biases in the candidate task pool rather than correcting them. Finally, the $\mathbf{S}\in\mathbb{R}^{n\times n}$5 cost of pairwise similarity construction is a scaling bottleneck for very large task inventories. The future directions suggested in the paper are more efficient task representations, dynamic updating of $\mathbf{S}\in\mathbb{R}^{n\times n}$6 during training, quality-aware mixture optimization, and task-targeted mixtures for benchmarks such as math, code, or safety (Chanda et al., 16 Jul 2025).