STEPH: Sparse Task Vector Mixup with Hypernetworks for Efficient Knowledge Transfer in WSI Prognosis

Conference: CVPR 2026 arXiv: 2603.10526 Code: GitHub Area: Medical Imaging / Computational Pathology Keywords: Whole-Slide Image, Survival Analysis, Cross-Cancer Knowledge Transfer, Task Vector, Hypernetwork, Model Merging

TL;DR

STEPH proposes a model merging framework based on Task Vector Mixup (TVM) and hypernetwork-driven sparse aggregation, which efficiently transfers knowledge from multiple cancer-type-specific prognosis models into a target cancer model. It achieves a mean C-Index of 0.6949 across 13 TCGA datasets (+5.14% vs. cancer-type-specific learning, +2.01% vs. ROUPKT), while requiring only a single-model forward pass at inference—far more efficient than multi-model representation transfer approaches.

Background & Motivation

Background: Pathology whole-slide images (WSIs) are gigapixel-scale and serve as a primary data source for cancer prognosis (survival analysis). Multiple instance learning (MIL)-based cancer-type-specific models are the dominant paradigm; however, each cancer type typically has only ~1,000 training samples, and high tumor heterogeneity limits generalization.

Limitations of Prior Work: (1) Cancer-type-specific learning—limited data and high heterogeneity lead to poor generalization of single-cancer models; (2) Multi-cancer joint training—the sheer volume of WSIs incurs prohibitive computational costs and introduces privacy risks; (3) Representation transfer (ROUPKT)—routes and aggregates WSI representations from multiple source models, but inference requires running all source models, with cost scaling linearly with the number of sources.

Key Challenge: How can a single model efficiently absorb cross-cancer knowledge without resorting to joint training (high computation) or multi-model inference (high inference cost)?

Goal: Transfer prognosis knowledge from multiple cancer types into a target model via model merging, enabling lightweight and efficient cross-cancer knowledge transfer.

Key Insight: The task vector \(\tau_t = \mathcal{M}_t - \mathcal{M}_0\) encodes cancer-type-specific prognosis knowledge. Unlike model merging in MTL—which aims to preserve multi-task capability (requiring resolution of task interference)—the goal in WSI prognosis is to enhance generalization on the target task, achieved by interpolating source and target task vectors via mixup to obtain better optimization directions.

Core Idea: Apply mixup interpolation to each source–target task vector pair to absorb beneficial knowledge, then use a hypernetwork to learn input-adaptive sparse aggregation weights, and finally merge into a single enhanced model.

Method

Overall Architecture

Pre-trained initialization \(\mathcal{M}_0\) → independent fine-tuning per cancer type to obtain \(\mathcal{M}_t, \{\mathcal{M}_{s_i}\}\) → compute task vectors \(\tau_t, \{\tau_{s_i}\}\) → TVM applies mixup interpolation to each \((\tau_t, \tau_{s_i})\) pair → hypernetwork-driven sparse aggregation selects top-\(K\) mixed vectors and computes a weighted sum to yield \(\tau_t^*\) → final model \(\mathcal{M}_t^* = \mathcal{M}_0 + \tau_t^*\).

Key Designs

1. Task Vector Mixup (TVM)

   • Function: Interpolates source and target task vectors via mixup to obtain optimization directions that incorporate knowledge from both.
   • Mechanism: For each source–target pair \((\tau_t, \tau_{s_i})\), the interpolation is computed as \(\tau_{\text{mix}} = \lambda_i \tau_t + (1-\lambda_i) \tau_{s_i}\). The mixing coefficient \(\lambda_i\) is adaptively predicted by a hypernetwork \(\mathcal{H}_{\text{mix}}\) conditioned on input WSI features (rather than being fixed), constrained to \([0,1]\) via sigmoid. The hypernetwork employs a mean-MIL encoder to process bag-of-patches features.
   • Design Motivation: Grounded in Vicinal Risk Minimization (VRM)—task vectors are cumulative gradients, and their mixup approximates the gradient of training on virtual mixed data, yielding better-generalizing models. Loss landscape visualization and SAR analysis confirm that \(\lambda \in [0.7, 0.8]\) leads to lower training and test losses.

2. Sparse Task Vector Aggregation

   • Function: Selects the most beneficial top-\(K\) mixed task vectors from \(m\) candidates for weighted aggregation.
   • Mechanism: A second hypernetwork \(\mathcal{H}_{\text{agg}}\) (sharing the MIL encoder with \(\mathcal{H}_{\text{mix}}\) but with an independent output head) produces aggregation weights \(w = \{w_i \geq 0\}\); the top-\(K\) weights are selected and the aggregated task vector is computed as \(\tau_t^* = \sum_j w_j \tau_{\text{mix},j}\). An auxiliary loss \(\mathcal{L}_{\text{agg}} = (\log \sum_i e^{w_i})^2\) suppresses excessively large weights.
   • Design Motivation: Not all source cancer knowledge is beneficial to the target—some models may be poorly trained or inherently conflicting. Sparse selection (inspired by MoE) filters out redundant and harmful knowledge. Input-adaptive weights \(w\) are more flexible than globally fixed weights, as different WSI samples may benefit from different source cancer types.

3. Hypernetwork-Driven Dynamic Weights

   • Function: Learns input-conditioned \(\lambda\) and \(w\) via hypernetworks, replacing globally fixed parameters.
   • Mechanism: \(\mathcal{H}_{\text{mix}}\) and \(\mathcal{H}_{\text{agg}}\) share a mean-MIL encoder (to reduce parameters) with separate fully connected output heads. The training objective combines NLL survival loss with auxiliary regularization terms: \(\mathcal{L}_{\text{mix}} = \sum_j \lambda_j^2/K\) encourages absorption of source knowledge, and \(\mathcal{L}_{\text{agg}} = (\log \sum_i e^{w_i})^2\) prevents weight explosion.
   • Design Motivation: WSI prognosis datasets are small (~1,000 samples), making grid search over fixed parameters prone to overfitting small validation sets. Hypernetwork-predicted dynamic parameters are more robust than fixed \(\lambda/w\); applying the hypernetwork scheme to existing model merging methods yields an average improvement of 14.5%.

Loss & Training

NLL survival loss + auxiliary losses (\(\beta=0.05\), \(\gamma\) via cross-validation); \(K=5\); \(m=12\) (12 source cancer types); 5-fold CV; UNI for patch feature extraction.

Key Experimental Results

Main Results — Mean C-Index across 13 TCGA Datasets

Method	Category	Mean C-Index
Vanilla (cancer-type-specific)	Cancer-type-specific	0.6609
Fine-tuned (cancer-type-specific)	Cancer-type-specific	0.6611
ROUPKT	Representation Transfer	0.6812
Model Avg.	Model Merging	0.5804
AdaMerging	Model Merging	0.5689
TIES AM	Model Merging	0.6396
Surgery AM	Model Merging	0.5943
Iso-C AM	Model Merging	0.5699
STEPH	Model Merging	0.6949

Ablation Study

Configuration	Mean C-Index
w/o mixup, fix \(\lambda=0\) (source only)	0.6860
w/o mixup, fix \(\lambda=1\) (target only)	0.6851
w/ mixup, trainable \(\lambda\)	0.6921
w/ mixup, hypernetwork \(\lambda\)	0.6949
w/o sparsity	0.6912
w/ sparsity, trainable \(w\)	0.6490
w/ sparsity, hypernetwork \(w\)	0.6949

Hypernetwork Scheme Applied to Existing Methods

Method	Original	+Hypernetwork Aggregation	Gain
AdaMerging	0.5689	0.6877	+20.9%
TIES	0.6396	0.6802	+6.3%
Surgery	0.5943	0.6668	+12.2%
Iso-C	0.5699	0.6761	+18.6%

Key Findings

• STEPH outperforms cancer-type-specific learning on 12 of 13 datasets, with a mean improvement of 5.14% and a maximum single-dataset gain of 11.4% (BRCA).
• Existing general-purpose model merging methods (AdaMerging, TIES, etc.) perform poorly on WSI prognosis (0.57–0.64), as they are designed for multi-task retention rather than single-task enhancement.
• Input-adaptive weights via hypernetworks are the key contribution—applying them to any existing method yields an average improvement of 14.5%.
• SAR analysis reveals that TVM improvements are primarily attributable to attention layers rather than embedding layers, indicating that MIL attention aggregation benefits more from cross-cancer knowledge than instance-level encoding.
• Visualization of \(\lambda\) training dynamics shows that KIPAN, COADREAD, and BLCA exhibit \(\lambda_i < 0.3\) with relatively large \(w_i\), confirming that BRCA genuinely absorbs beneficial knowledge from these specific cancer types.

Highlights & Insights

1. Model merging for single-task enhancement rather than MTL: Unlike mainstream model merging research (aimed at acquiring multi-task capability), STEPH targets enhanced generalization on a single task—a shift in objective that introduces fundamentally different methodological requirements (from resolving task interference to mining beneficial knowledge).
2. VRM theoretical framework supports TVM: Task vector mixup is not simple parameter averaging; it approximates training on mixed virtual data, providing a principled theoretical foundation.
3. Hypernetworks as a general-purpose enhancement: Applying hypernetwork-driven aggregation to four existing methods yields an average improvement of 14.5%, demonstrating the strong generality of the input-adaptive mechanism itself.

Limitations & Future Work

1. Reliance on TCGA datasets, where some cancer types have very few samples (<400), may introduce evaluation bias.
2. Experiments are based on standard attention-based MIL architectures; more advanced MIL methods (e.g., graph-based) have not been validated.
3. STEPH still requires training data to learn merging weights; training-free model merging is a promising future direction.
4. \(K=5\) (top-5 mixed vectors) is globally fixed; adaptive \(K\) selection remains unexplored.

Related Work & Insights

• vs. ROUPKT: ROUPKT runs all source models at inference to obtain representations for routing aggregation, with cost scaling linearly in the number of sources. STEPH merges all knowledge into a single model during training and requires only one forward pass at inference, representing a qualitative leap in efficiency.
• vs. AdaMerging/TIES: General model merging methods target multi-task retention and focus on resolving interference; STEPH targets single-task enhancement and focuses on mining beneficial knowledge—a difference in objective that leads to substantially divergent methodologies.
• vs. data mixup: Classical mixup interpolates in input/feature space; STEPH performs mixup in parameter space (task vectors), representing an interesting extension of the mixup principle.

Rating

⭐⭐⭐⭐

• Novelty ⭐⭐⭐⭐: The perspective of using model merging for single-task enhancement is novel; TVM is grounded in VRM theory.
• Experimental Thoroughness ⭐⭐⭐⭐⭐: Comprehensive coverage across 13 datasets, multiple baselines, ablations, visualizations, and hyperparameter analyses.
• Writing Quality ⭐⭐⭐⭐: Problem formulation is clear; theoretical analysis and visualizations provide sufficient supporting evidence.
• Value ⭐⭐⭐⭐: Offers an efficient solution for cross-cancer knowledge transfer in computational pathology; the hypernetwork aggregation scheme demonstrates broad generality.

Related Papers

• [CVPR 2026] Sparse Task Vector Mixup with Hypernetworks for Efficient Knowledge Transfer in Whole-Slide Image Prognosis
• [CVPR 2026] Cross-Slice Knowledge Transfer via Masked Multi-Modal Heterogeneous Graph Contrastive Learning for Spatial Gene Expression Inference
• [CVPR 2026] Association of Radiologic PPFE Change with Mortality in Lung Cancer Screening Cohorts
• [CVPR 2026] OmniFM: Toward Modality-Robust and Task-Agnostic Federated Learning for Heterogeneous Medical Imaging
• [CVPR 2026] MedGRPO: Multi-Task Reinforcement Learning for Heterogeneous Medical Video Understanding