Technical Summary: Knowledge Transfer Scaling Laws for 3D Medical Imaging

Problem Statement

The training of unified 3D medical foundation models requires mixing heterogeneous imaging domains (e.g., CT, MRI, PET) with distinct acquisition physics, resolutions, and clinical signals. While architectures and self-supervised objectives have been extensively studied, the optimal data mixture design remains an open question. Current practices rely on heuristic strategies, such as uniform sampling or data-proportional sampling, which treat all imaging domains as interchangeable. However, these approaches fail to account for two critical empirical observations:

1. Heterogeneous Scaling: Different medical domains exhibit variable rates of improvement (scaling) as training data increases.
2. Asymmetric Transfer: Knowledge transfer between domains is directional; training on one domain may significantly benefit another, while the reverse transfer may be negligible.

The core challenge is to determine the optimal allocation of a fixed pretraining token budget across these heterogeneous domains to minimize reconstruction loss and maximize downstream task performance, without resorting to exhaustive and computationally prohibitive mixture ablation studies.

Methodology

The authors formulate 3D medical data mixture design as a transfer-aware scaling-law optimization problem. The method relies on a controlled proxy objective—Masked Autoencoding (MAE) reconstruction loss—which is applied consistently across heterogeneous 3D domains.

1. Transfer-Aware Scaling Law

The authors derive a closed-form surrogate model to predict the MAE validation loss ($L_i$) for a target domain $i$ under a specific data mixture $h$ and total budget $T$. The model combines two components:

• Domain-Specific Self-Scaling: Each domain $i$ follows a power law when trained in isolation:
  $$L_i(T) = E_i + \frac{C_i}{T^{\beta_i}}$$
  Where $E_i$ is the asymptotic loss floor, $C_i$ controls reducible loss, and $\beta_i$ is the scaling exponent. Crucially, $\beta_i$ varies significantly across domains (observed range: 0.21 to 0.76), indicating that some domains benefit from additional data much more rapidly than others.

• Directed Cross-Domain Transfer: The authors introduce an asymmetric transfer matrix $\tau \in [0, 1]^{K \times K}$, where $\tau_{ij}$ quantifies how effectively training on source domain $j$ contributes to the loss reduction of target domain $i$.

  • $\tau_{ii} = 1$ by definition.
  • For $i \neq j$, $\tau_{ij}$ discounts the source budget. $T$ samples from domain $j$ are equivalent to $\tau_{ij}T$ effective samples for domain $i$.

The effective budget ($T_{eff, i}$) reaching target domain $i$ under mixture $h$ is the sum of direct and transferred contributions:
$$T_{eff, i}(h, T) = \sum_{j=1}^{K} \tau_{ij} h_j T$$

The final predictive surrogate for the loss on domain $i$ is:
$$\hat{L}_i(h, T) = E_i + \frac{C_i}{(T_{eff, i}(h, T))^{\beta_i}}$$

2. Estimation and Optimization

The framework operates in two stages to avoid exhaustive search:

1. Parameter Fitting: Using small-scale "proxy runs" (pure-domain training and directed source-to-target training), the authors fit the parameters $\{E_i, C_i, \beta_i\}$ and the transfer coefficients $\{\tau_{ij}\}$. This requires $K^2$ configurations (36 runs for 6 domains) at small budgets (10K–100K volumes).
2. Mixture Optimization: Once parameters are estimated, the optimal mixture weights $h^*$ are found by minimizing the weighted sum of predicted losses across all target domains, subject to budget constraints ($\sum h_j = 1, h_j \geq \epsilon$). This is solved via gradient-based constrained optimization (SLSQP).

Key Contributions

1. Formulation of Data Allocation: The paper formulates cross-domain data allocation in 3D medical imaging as a scaling-law optimization problem, utilizing MAE reconstruction loss as a controlled proxy across six heterogeneous domains (Abdominal CT/MRI, Brain T1/T2 MRI, Head-and-Neck CT/PET).
2. Transfer-Aware Scaling Law: A novel law that jointly models heterogeneous scaling exponents ($\beta_i$) and an asymmetric transfer matrix ($\tau_{ij}$). This yields a closed-form surrogate for compute-dependent mixture optimization derived from small-scale proxy runs.
3. Hub-and-Island Structure Discovery: The learned allocations reveal an interpretable structure where:
   • Hubs: Domains with broad outgoing transfer (e.g., ABD-CT) act as hubs; allocating budget here benefits multiple other domains simultaneously.
   • Islands: Domains with weak incoming transfer (e.g., HN-PET) act as islands; they require direct investment because knowledge from hubs cannot effectively transfer to them.
4. Empirical Validation: The derived "Transfer-Aware" allocation outperforms standard heuristics (Uniform, Data-Proportional, and Transfer-Naïve) and extrapolates accurately to unseen budgets.

Results

Scaling Law Performance

• Loss Reduction: The Transfer-Aware allocation achieves the lowest MAE loss across all tested budgets (10K to 200K). At a budget of 50K, it outperforms Data-Proportional sampling by 58% in mean MAE loss.
• Extrapolation: The model fitted on budgets up to 100K extrapolates to a held-out 200K budget with high accuracy (predicted loss: 0.0227 vs. observed: 0.0224; correlation $r=0.989$).
• Efficiency: The method requires only 36 proxy runs for estimation, compared to the thousands of runs required for an exhaustive grid search over the mixture simplex.

Downstream Clinical Tasks

The pretrained encoders were fine-tuned on supervised tasks including organ/lesion segmentation and disease classification.

• Segmentation: Transfer-Aware pretraining yielded the best Dice scores across all six segmentation benchmarks (AMOS, BraTS, HECKTOR, BHSD), outperforming Data-Proportional baselines by 0.6 to 4.7 points.
• Classification: The method achieved the highest accuracy and AUC on brain MRI (OASIS-1), brain PET (SynBrain-PET), and abdominal CT (OrganMNIST3D) classification tasks.
• Correlation: The ordering of performance in downstream tasks mirrored the ordering of MAE proxy losses, validating the proxy objective.

Structural Insights

• Asymmetry: The transfer matrix is highly asymmetric. For instance, ABD-CT transfers strongly to ABD-MRI ($\tau=0.56$) but weakly to HN-PET ($\tau=0.04$).
• Allocation Shift: The optimal strategy concentrates budget on the "Hub" (ABD-CT, ~24%) and the "Island" (HN-PET, ~56%), whereas Data-Proportional sampling allocates based on dataset size, and Uniform sampling allocates equally.
• Amplification: Domains with low direct allocation (the "floor" domains) achieved 3.7x to 6.1x effective budget amplification through cross-domain transfer, yet still outperformed baselines.

Significance and Claims

The paper claims that data allocation is a key design axis for scalable 3D medical foundation models. By moving from heuristic mixture choices to an analytical, transfer-aware optimization framework, practitioners can:

1. Reduce Compute Costs: Eliminate the need for exhaustive mixture ablation studies by using small-scale proxy runs to predict optimal large-scale mixtures.
2. Improve Representations: Generate stronger pretrained representations that transfer more effectively to diverse clinical tasks (segmentation and classification) across modalities.
3. Understand Domain Relationships: Reveal the intrinsic "hub-and-island" topology of medical imaging domains, providing actionable guidance on which domains serve as universal sources of knowledge and which require dedicated data collection.

The authors maintain modesty regarding limitations, noting that the linear composition of transfer contributions may not capture synergistic interactions at larger scales, and that the transfer matrix is currently treated as static rather than budget-dependent. However, the results suggest that first-order transfer dominates in current 3D medical settings, making the proposed framework a practical and effective tool for optimizing heterogeneous medical data mixtures.