Technical Summary: Cohort-Based Active Modality Acquisition (CAMA)

Problem Definition

Real-world multimodal machine learning often operates under constraints where certain data modalities are missing, costly to acquire, or invasive to obtain. While existing paradigms like Active Learning (AL), Active Feature Acquisition (AFA), and Active Modality Acquisition (AMA) focus on optimizing acquisition for individual samples (often during training or at the test-time sample level), they frequently fail to address the specific challenge of cohort-level resource allocation.

In many practical scenarios, such as healthcare, budgets are resource-specific (e.g., a fixed number of MRI scans or a grant for proteomics testing) rather than flexible per-patient funds. The central problem addressed is: Given a cohort of $N$ samples where only a subset of modalities is initially available, and a fixed budget $\beta$ for acquiring one specific, costly missing modality, which subset of samples should be selected to maximize a global performance metric (e.g., AUROC, AUPRC) for the entire cohort?

This setting, termed Cohort-based Active Modality Acquisition (CAMA), differs from prior work by shifting the decision level from the individual sample to the cohort, aiming to optimize the global discriminative performance of the cohort rather than just individual prediction accuracy.

Methodology

1. Problem Formulation

The authors formalize CAMA as an optimization problem. Let $D = \{(x_i, y_i)\}_{i=1}^N$ be a dataset where each sample $i$ has $P$ modalities, but only a subset $P^{avail}_i$ is observed. The goal is to select a subset $S \subseteq \{1, \dots, N\}$ with $|S| = \beta$ to acquire the missing modality.

• Scores: The framework distinguishes between:
  • $s^{avail}_i$: Score using available modalities.
  • $s^{acquired}_i$: Score using available + acquired modalities (the "oracle" score).
  • $\{s^{imp}_{i,k}\}_{k=1}^K$: A set of $K$ imputed scores estimating $s^{acquired}_i$ using a generative model.
• Objective: Maximize a global metric $Metric(y, s(S))$, where $s(S)$ is the vector of scores updated for samples in $S$.

2. Acquisition Function (AF) Strategies

Since solving the optimization problem exactly is computationally intractable, the authors propose heuristic Acquisition Functions (AFs) to rank samples. These are categorized into:

• Oracle Strategies: Theoretical upper bounds assuming perfect knowledge of true labels and $s^{acquired}_i$. They greedily select samples yielding the largest marginal gain in the global metric.
• Upper-Bound Heuristics: Label-agnostic strategies that assume knowledge of $s^{acquired}_i$ but not $y_i$. They rely on metrics like true uncertainty reduction, rank change, or KL-Divergence between probability distributions.
• Imputation-Based Strategies (Proposed): These utilize a generative model $f_{imp}$ to estimate the counterfactual scores $\{s^{imp}_{i,k}\}$ without actually acquiring the data. Strategies include:
  • Expected KL-Divergence: Prioritizing samples where the imputed distribution diverges most from the baseline.
  • Expected Rank Change: Prioritizing samples expected to shift significantly in the cohort's ranking.
  • Expected Uncertainty Reduction: Prioritizing samples where uncertainty is expected to decrease most.
  • Expected Probability: Prioritizing samples with high average imputed probability.
• Baselines: Strategies using only pre-acquisition information (e.g., baseline uncertainty) or random selection.

3. Model Architecture

The proposed architecture combines a discriminative late-fusion model with a generative imputation module (Figure 2):

• Discriminative Component: A late-fusion model $f$ with modality-specific encoders (e.g., ViT for images, BERT for text) and a fusion classifier. It produces $s^{avail}_i$.
• Generative Component: A generative model $f_{imp}$ (implemented as Diffusion Transformers or Beta-Conditional VAEs) trained to approximate the conditional distribution $p(z^{(k)} | \{z^{(m)}\}_{m \in P^{avail}_i})$ for missing modalities.
• Training: The model is trained end-to-end with a combined loss of binary cross-entropy (discriminative) and a generative loss (variational bound or denoising objective). Crucially, during training, the generative outputs are decoupled from the classifier; the classifier is trained only on available modalities via attention masks to prevent data leakage, while the generative module learns to reconstruct latent embeddings.

Key Contributions

1. The CAMA Setting: The formalization of a novel test-time, cohort-level acquisition setting that addresses the prioritization of samples for a single, fixed modality under a global budget constraint.
2. Theoretical Framework for AFs: The derivation of acquisition functions based on standard evaluation metrics (AUROC, AUPRC), including the definition of upper-bound heuristics and imputation-based strategies.
3. Novel Architecture: A hybrid architecture integrating discriminative late fusion with generative imputation (specifically adapting IMDer [64] concepts) to enable counterfactual reasoning at test time.
4. Comprehensive Evaluation: Extensive empirical validation across diverse datasets (UK Biobank, MIMIC-Symile, MIMIC-HAIM, MOSEI) with up to 15 modalities and 100,000 samples.

Experimental Results

The authors evaluated their methods on four datasets, focusing on binary classification tasks.

• Performance: The Imputation-based KL-Divergence strategy consistently outperformed all non-oracle baselines (including uncertainty-based and random selection) across most datasets and metrics.
• Comparison to Heuristics: While the "True" (oracle knowledge) heuristics provided upper bounds, the imputation-based KL-Divergence strategy often approached or, in some specific cases (e.g., MOSEI), exceeded the performance of the corresponding label-agnostic upper-bound heuristics.
• Scalability: The method demonstrated scalability on the UK Biobank dataset (approx. 100k samples, 15 modalities), successfully guiding the acquisition of proteomics data for Systemic Lupus Erythematosus (SLE) prediction.
• Ablation Studies: Key architectural decisions, such as Layer Normalization, label smoothing, and decoupling the generative module from the classifier during training, were found to be essential for stability and performance.
• Efficiency: While DDPMs offered strong performance, Beta-Conditional VAEs provided a faster inference alternative with a trade-off in performance.

Significance and Claims

The paper claims that CAMA provides an effective approach for optimizing modality acquisition in resource-constrained settings, particularly where a single, expensive modality must be allocated across a large cohort.

• Real-World Relevance: The authors highlight the UK Biobank proteomics project as a prime example, where a fixed budget necessitated strategic selection of patients for expensive testing. Their method offers a data-driven framework for such decisions.
• Resource Efficiency: By prioritizing samples where the acquisition of a missing modality yields the highest global gain (rather than just local uncertainty reduction), the method maximizes the utility of limited resources.
• Limitations: The authors acknowledge that their current work focuses on a single specified modality and fixed costs, excluding the more complex problem of jointly selecting both the modality and the sample or handling heterogeneous costs. They also note that their evaluation uses synthetic missingness patterns (except for UKB) and does not yet include a dedicated fairness analysis.

In conclusion, the work establishes that imputation-based acquisition functions, specifically those leveraging expected KL-Divergence, can effectively guide cohort-level resource allocation, bridging the gap between theoretical acquisition strategies and practical, budget-constrained multimodal learning applications.