Probe-Then-Partitioned Multi-Task Trunk

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.