# Probe-Then-Partitioned Multi-Task Trunk

- ID: 1237
- Canonical URL: https://synthcore.org/idea/1237/probe-then-partitioned-multi-task-trunk
- API JSON: https://synthcore.org/api/idea/1237.json
- API Markdown: https://synthcore.org/api/idea/1237.md
- Verification status: failed_benchmark
- Source: [arXiv:2607.21426](https://arxiv.org/abs/2607.21426)
- Category: architecture
- Solves: accuracy, generalization, scalability
- ML areas: transformer, moe-routing, mlp, training-dynamics
- Math tags: information-theory, geometry, statistics, optimization
- Ratings: usefulness 7/10; difficulty 5/10; novelty 5/10

## Idea description

Train a cheap shared multi-task probe briefly, extract one semantic embedding per task, and use density-based clustering to determine which tasks should share a neural trunk. After clustering, replace the globally shared trunk by one trunk per discovered cluster, with task heads remaining separate; this preserves cooperation among related tasks while isolating destructive task interactions.

## Mathematical statement

The paper formulates joint clustering and representation learning as maximizing the sum of task-wise mutual informations, $$\max_{p^{\mathrm{CU}}_1,\ldots,p^{\mathrm{CU}}_K,p^{\mathrm{SU}}_1,\ldots,p^{\mathrm{SU}}_N,\mathcal{T}_1,\ldots,\mathcal{T}_K}\sum_{k=1}^{K}\sum_{i\in\mathcal{T}_k} I(\bm{Y}_i;Z_i).$$ Here, $N$ is the number of tasks, $K$ is the number of clusters or cooperative units, $\mathcal{T}_k\subseteq\{1,\ldots,N\}$ is the set of tasks assigned to cluster $k$, $Z_i$ is the semantic target for task $i$, and $\bm{Y}_i$ is its decoded or task-specific representation. The implementation replaces unavailable mutual information with a supervised proxy: for task loss $\ell_i$, use $\widehat I_i=-\mathbb{E}[\ell_i]$ up to a task-dependent constant, so maximizing the information objective corresponds to minimizing task losses while restricting parameter sharing through the partition. Each task receives an embedding $e_i\in\mathbb{R}^d$, formed from its probe-stage hidden representation and normalized gradient signature. The task distance is cosine distance $d_{ij}=1-(e_i^    op e_j)/(\|e_i\|_2\|e_j\|_2)$. For density clustering, the $m$-core distance is $\mathrm{core}_m(i)=d_{i,(m)}$, the distance from task $i$ to its $m$th nearest task, and the mutual-reachability distance is $d^{\mathrm{mr}}_{ij}=\max\{\mathrm{core}_m(i),\mathrm{core}_m(j),d_{ij}\}$. A minimum-spanning-tree hierarchy over these distances produces dense task groups; ambiguous noise tasks are initially assigned singleton trunks.

## Key formulas

- $$\max_{p^{\mathrm{CU}}_1,\ldots,p^{\mathrm{CU}}_K,p^{\mathrm{SU}}_1,\ldots,p^{\mathrm{SU}}_N,\mathcal{T}_1,\ldots,\mathcal{T}_K}\quad \sum_{k=1}^{K}\sum_{i\in\mathcal{T}_k} I(\bm{Y}_i;Z_i)$$
- $$d_{ij}=1-\frac{e_i^\top e_j}{\|e_i\|_2\|e_j\|_2},\qquad \mathrm{core}_m(i)=d_{i,(m)},\qquad d^{\mathrm{mr}}_{ij}=\max\{\mathrm{core}_m(i),\mathrm{core}_m(j),d_{ij}\}$$
- $$\mathcal{L}_{\mathrm{clustered}}(\theta,\{\phi_k\},\{\psi_i\})=\sum_{k=1}^{K}\sum_{i\in\mathcal{T}_k}\lambda_i\,\ell_i\!\left(h_i\!\left(g_k(x;\theta_k);\psi_i\right),y_i\right)$$
- $$e_i=\operatorname{norm}\!\left(\frac{1}{B}\sum_{b=1}^{B}\operatorname{Pool}\big(h_i(x_b)\big)\;\Big\|\;\frac{1}{B}\sum_{b=1}^{B}\frac{\nabla_{\theta}\ell_i(x_b)}{\|\nabla_{\theta}\ell_i(x_b)\|_2}\right)$$

## Implementation notes

1. Integration point: use a standard multi-task Transformer or MLP with a shared probe trunk $g(x;\theta)$ and separate task heads $h_i(\cdot;\psi_i)$. Train this model for $E_{\mathrm{init}}$ epochs, where $E_{\mathrm{init}}$ is 5--10% of the planned training budget. During this phase, collect for every task $i$ a task embedding $e_i$: concatenate the mean-pooled hidden activation at the trunk output with a normalized average gradient of the shared trunk, then L2-normalize the concatenation. 2. Pseudocode: initialize shared trunk $\theta$ and heads $\psi_i$; for each probe epoch, sample a batch for every task, compute $L=\sum_i\lambda_i\ell_i(h_i(g(x_i;\theta)),y_i)$, and update $(\theta,\psi_i)$; after probing, compute all $e_i$; calculate pairwise cosine distances $d_{ij}$; calculate $m$-nearest-neighbor core distances and $d^{\mathrm{mr}}_{ij}$; run HDBSCAN with minimum cluster size 2 or 3; produce groups $\mathcal{T}_k$. Replace $g$ by independent cluster trunks $g_k(x;\theta_k)$, copy the probe weights into every $\theta_k$, attach task head $i$ to the trunk associated with its cluster, and continue end-to-end optimization using $\mathcal{L}_{\mathrm{clustered}}$. 3. The paper-derived decisions are the staged partition and the information-maximization target. Compute embeddings, distances, and density clusters from observed activations and gradients; do not backpropagate through the discrete clustering. Treat HDBSCAN noise points as singleton groups initially, then optionally merge each singleton into the cluster producing the lowest held-out loss after one validation epoch. To control parameter growth, use low-rank adapters in each cluster trunk or cluster only the upper half of Transformer blocks. 4. First experiment: use NYUv2 or Multi-MNIST with 4--8 tasks, comparing one globally shared trunk, independent task models, random task partitions, and probe-then-partition at equal total parameter count and FLOPs. Measure final mean accuracy, worst-task accuracy, pairwise gradient-conflict rate, validation loss, and training cost. Success is lower worst-task loss and reduced negative-gradient cosine without more than 1.5x inference cost; test robustness by halving the probe budget and permuting task labels.

## Verification

- Status: failed_benchmark
- Mechanism evidence: yes
- Mechanism confirmed: no
- Practical verdict: inconclusive
- Verdict: Built a NumPy mechanism verification and a PyTorch probe-then-partitioned multi-task trunk MVP. The core predictions manifested: cosine-distance identities were exact, the small-noise embedding law E[d]≈7σ² matched within 2.3% for σ≤0.06, and the analytic partition gain matched the angle sweep to numerical precision. In the learned toy regression, probe partitioning reduced mean MSE from 0.340 to 0.032 and worst-task MSE from 0.992 to 0.078, but the discovered partition used singleton trunks and was not parameter/FLOP matched.

### Mechanism check

- Verdict: Built a NumPy mechanism verification and a PyTorch probe-then-partitioned multi-task trunk MVP. The core predictions manifested: cosine-distance identities were exact, the small-noise embedding law E[d]≈7σ² matched within 2.3% for σ≤0.06, and the analytic partition gain matched the angle sweep to numerical precision. In the learned toy regression, probe partitioning reduced mean MSE from 0.340 to 0.032 and worst-task MSE from 0.992 to 0.078, but the discovered partition used singleton trunks and was not parameter/FLOP matched.
- Confidence: 8/10
- Limitations: This is a tiny synthetic regression experiment, not NYUv2 or Multi-MNIST. HDBSCAN itself was not used; a lightweight MST-cut approximation was implemented. The learned comparison did not equalize parameter count, inference FLOPs, or probe-plus-training cost, and gradient-conflict rate, validation generalization, label-permutation robustness, and half-budget probing were not tested.

### Practical benchmark

- Paired seeds: 8
- Baseline mean: 0
- Idea mean: 0
- p-value: 0.54355
- Paired wins: 3/8
- Benchmark verdict: no significant win

## Artifacts

- [bench_report.json](https://synthcore.org/code/497/bench_report.json)
- [bench_run.py](https://synthcore.org/code/497/bench_run.py)
- [experiment.py](https://synthcore.org/code/497/experiment.py)
- [report.md](https://synthcore.org/code/497/report.md)
- [report_bench_2026-08-31T160400.md](https://synthcore.org/code/497/report_bench_2026-08-31T160400.md)
- [results.json](https://synthcore.org/code/497/results.json)
- [torch_mini.py](https://synthcore.org/code/497/torch_mini.py)
- [torch_results.json](https://synthcore.org/code/497/torch_results.json)
- [Download all files as ZIP](https://synthcore.org/download/497)

## Disclaimer

AI-generated research hypothesis, automatically tested. Not peer-reviewed.