---
title: 'TaskPGM: Task Probabilistic Graphical Model'
url: https://www.emergentmind.com/topics/taskpgm
type: topic
---

# TaskPGM: Task Probabilistic Graphical Model

TaskPGM, short for Task Probabilistic Graphical Model, is a framework for optimizing the task mixture used in large language model 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 [2507.12612].

## 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 [2507.12612].

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 [2507.12612].

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 [2507.12612].

## 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 [2507.12612].

The MRF assigns unary and pairwise potentials as
$$
\Psi_i = \beta\, s_i,\qquad \Psi_{ij} = \lambda\, L_{ij},
$$
where $\beta>0$ and $\lambda>0$ are scalar weights and $L$ is the graph Laplacian built from $\mathbf{S}$. Collecting these into a unary vector $\bm{\Psi}_{\text{un}}$ and pairwise matrix $\bm{\Psi}_{\text{pair}}$, the energy is
$$
\mathbb{E}(\mathbf{p}) =
-\bm{\Psi}_{\text{un}}^\top \mathbf{p}
+ \frac{1}{2}\mathbf{p}^\top \bm{\Psi}_{\text{pair}} \mathbf{p}.
$$
The optimization problem is
$$
\mathbf{p}^\ast = \arg\min_{\mathbf{p}\in\Delta_n}\ \mathbb{E}(\mathbf{p}).
$$
The negative linear term favors tasks with large unary mass, while the quadratic Laplacian term penalizes concentration on neighboring, hence similar, tasks [2507.12612].

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$ pulls the solution toward representative hubs, while small $\beta/\lambda$ promotes a more spread-out allocation [2507.12612].

## 3. Behavioral similarity and task affinities

TaskPGM defines task relationships through model behavior rather than through metadata or text embeddings. For each task $T_i$, a single-task finetuned model is obtained as
$$
\theta^\ast(T_i) := \theta_0 + \boldsymbol{\tau}(T_i),
$$
where $\theta_0$ is the base LLM and $\boldsymbol{\tau}(T_i)$ is the task vector. Let $P_{\theta^\ast(T_i)}(y\mid x)$ denote the predictive distribution [2507.12612].

One similarity construction is a PMI-style score between tasks $T_i$ and $T_j$:
$$
S_{ij} =
\frac{1}{2}
\Bigg[
\frac{1}{n}\sum_{k=1}^{n}
\log\frac{
P_{\theta^\ast(T_i)}(y^{T_j}_{k}\mid x^{T_j}_{k})
}{
P_{\theta^\ast(T_j)}(y^{T_j}_{k}\mid x^{T_j}_{k})
}
+
\frac{1}{m}\sum_{r=1}^{m}
\log\frac{
P_{\theta^\ast(T_j)}(y^{T_i}_{r}\mid x^{T_i}_{r})
}{
P_{\theta^\ast(T_i)}(y^{T_i}_{r}\mid x^{T_i}_{r})
}
\Bigg].
$$
The stated intuition is that if the model trained on $T_i$ assigns higher probability to correct labels on $T_j$’s data than the model trained on $T_j$ itself, then the similarity increases; the same logic is applied symmetrically in the reverse direction [2507.12612].

A second construction is based on Jensen–Shannon Divergence. For a sample $(x^{T_j}_k,y^{T_j}_k)\in\mathcal{D}_{T_j}$, define
$$
P_k = P_{\theta^\ast(T_i)}(\cdot\mid x^{T_j}_k),\quad
Q_k = P_{\theta^\ast(T_j)}(\cdot\mid x^{T_j}_k),\quad
M_k = \frac{1}{2}(P_k+Q_k),
$$
with per-sample divergence
$$
\mathrm{JSD}_{k}^{(j\leftarrow i)}
=
\frac{1}{2}KL(P_k\parallel M_k)
+
\frac{1}{2}KL(Q_k\parallel M_k).
$$
The paper then symmetrizes across both directions:
$$
S_{\mathrm{JSD}}(T_i;T_j)=
\frac{1}{2}
\left[
\frac{1}{n}\sum_{k=1}^n \mathrm{JSD}_{k}^{(j\leftarrow i)}
+
\frac{1}{m}\sum_{r=1}^m \mathrm{JSD}_{r}^{(i\leftarrow j)}
\right].
$$
JSD is described as symmetric, bounded in $[0,\log 2]$, and well-behaved even when supports differ. The framework may then transform JSD into a similarity before constructing the Laplacian [2507.12612].

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 [2507.12612].

## 4. Optimization, convexity, and theoretical guarantees

The objective
$$
\min_{\mathbf{p}\in\Delta_n}
\left(
-\bm{\Psi}_{\text{un}}^\top \mathbf{p}
+\frac{1}{2}\mathbf{p}^\top \bm{\Psi}_{\text{pair}}\mathbf{p}
\right)
$$
is a convex quadratic over the simplex when $\bm{\Psi}_{\text{pair}}$ is positive semidefinite. The paper states that the similarity-based Laplacian is PSD, but also introduces spectral shifting for more general pairwise constructions:
$$
\bm{\Psi}_{\text{psd}}
:=
\bm{\Psi}_{\text{pair}}
+
|\Lambda_{\min}(\bm{\Psi}_{\text{pair}})|\,I.
$$
This guarantees convexity, although the paper notes that a large shift can bias solutions toward more uniform mixtures [2507.12612].

Using the Lagrangian
$$
\mathcal{L}(\mathbf{p},\nu,\bm{\mu})=
-\bm{\Psi}_{\text{un}}^\top\mathbf{p}
+\frac{1}{2}\mathbf{p}^\top\bm{\Psi}_{\text{pair}}\mathbf{p}
+\nu(\mathbf{p}^\top\mathbf{1}_n-1)
-\bm{\mu}^\top\mathbf{p},
$$
the Karush–Kuhn–Tucker conditions yield, in the interior case $p_i>0$, the stationarity equation
$$
\bm{\Psi}_{\text{pair}}\mathbf{p}
=
\bm{\Psi}_{\text{un}}-\nu\mathbf{1}_n.
$$
If $\bm{\Psi}_{\text{pair}}$ is invertible, the closed-form interior solution is
$$
\mathbf{p}^\ast
=
\bm{\Psi}_{\text{pair}}^{-1}\bm{\Psi}_{\text{un}}
-
\frac{
\mathbf{1}_n^\top \bm{\Psi}_{\text{pair}}^{-1}\bm{\Psi}_{\text{un}} - 1
}{
\mathbf{1}_n^\top \bm{\Psi}_{\text{pair}}^{-1}\mathbf{1}_n
}
\,
\bm{\Psi}_{\text{pair}}^{-1}\mathbf{1}_n.
$$
When some coordinates become negative, the paper states that one may restrict to the support where $p_i>0$ and re-solve, or project onto the simplex with standard projection algorithms [2507.12612].

TaskPGM also introduces a budgeted or top-$k$ variant. For a subset $A\subseteq V$ with support restricted to $\bar A$, define the utility
$$
f(\bar A)
:=
\max_{\mathbf{p}\in\Delta_n^{\mathbb{R}^+};\ \operatorname{supp}(\mathbf{p})\subseteq \bar A}
\overline{\mathbb{E}(\mathbf{p})},
$$
where $\overline{\mathbb{E}}=-\mathbb{E}$. Over the family
$$
\mathcal{I}=\{S\subseteq V\mid |S|\le k\},
$$
the discrete problem is $\max_{A\in\mathcal I} f(\bar A)$. The paper proves that this set function is monotonic and weakly submodular, with submodularity ratio
$$
\gamma = \frac{c_\Omega}{C_\Omega} > 0,
$$
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 [2507.12612].

## 5. Training pipeline and implementation

The implementation described in the paper follows a six-stage pipeline. First, for each candidate task $T_i$, a single-task model $\theta^\ast(T_i)$ 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 $O(n^2)$ pair evaluations. Fourth, the graph Laplacian and MRF potentials are constructed, with optional spectral correction. Fifth, the optimized task proportions $\mathbf{p}^\ast$ are obtained from the quadratic program. Sixth, an instance budget $B$ is allocated by sampling
$$
\mathbf{k}=(k_1,\dots,k_n)\sim \mathrm{Multinomial}(B,\mathbf{p}^\ast),
$$
after which $k_i$ instances are drawn uniformly from task $i$’s dataset and the base LLM is finetuned on the resulting mixture [2507.12612].

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 $2\times 10^{-5}$, AdamW, weight decay 0.01, and gradient checkpointing. The evaluated base models are Llama-2-7B and Mistral-7B-v0.3 in bf16 [2507.12612].

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 $n$ single-task finetunes, with $n\approx 316$ 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 [2507.12612].

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 [2507.12612].

## 6. Empirical behavior, interpretability, and limitations

Empirically, TaskPGM is compared against three baselines: Random, which samples $B$ 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 [2507.12612].

The reported results are heterogeneous but informative. For Llama-2-7B with batch size 8 and $B=25$K, 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 $B=50$K, 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 $B=25$K, 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 $B=100$K 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 [2507.12612].

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 [2507.12612].

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 $(\beta,\lambda)$ 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 $\beta=307,\lambda=60$ is reported as robust in ablations, while the main experiments fix $\beta=20,\lambda=10$ [2507.12612].

TaskPGM is also presented as an interpretability tool. The optimized probabilities $\mathbf{p}^\ast$ 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,
$$
\mathrm{TV}(\mathbf{p}_k,\mathbf{p}_{k+1}^{(k)})
=
\frac{1}{2}\sum_{i=1}^k
\left|(\mathbf{p}_k)_i-(\mathbf{p}_{k+1})_i\right|,
$$
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 [2507.12612].

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 $O(n^2)$ 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{p}^\ast$ during training, quality-aware mixture optimization, and task-targeted mixtures for benchmarks such as math, code, or safety [2507.12612].

Source: https://www.emergentmind.com/topics/taskpgm