On the Power of Source Screening for Learning Shared Feature Extractors

Leo Muxing Wang; Connor Mclaughlin; Lili Su

cs.LGFeb 18, 2026

On the Power of Source Screening for Learning Shared Feature Extractors

Authors: Leo Muxing Wang, Connor Mclaughlin, Lili Su

Organizations: Northeastern University

Abstract

Learning with shared representation is widely recognized as an effective way to separate commonalities from heterogeneity across various heterogeneous sources. Most existing work includes all related data sources via simultaneously training a common feature extractor and source-specific heads. It is well understood that data sources with low relevance or poor quality may hinder representation learning. In this paper, we further dive into the question of which data sources should be learned jointly by focusing on the traditionally deemed ``good'' collection of sources, in which individual sources have similar relevance and qualities with respect to the true underlying common structure. Towards tractability, we focus on the linear setting where sources share a low-dimensional subspace. We find that source screening can play a central role in statistically optimal subspace estimation. We show that, for a broad class of problem instances, training on a carefully selected subset of sources suffices to achieve minimax optimality, even when a substantial portion of data is discarded. We formalize the notion of an informative subpopulation, develop algorithms and practical heuristics for identifying such subsets, and validate their effectiveness through both theoretical analysis and empirical evaluations on synthetic and real-world datasets.

Figures & tables

		SOTA ( Niu et al., 2024 )
General Cases	UB	$O\left(\sqrt{\frac{d\lambda_{1}}{N\lambda_{k}^{2}}}+\sqrt{\frac{Md}{N^{2}\lambda_{k}^{2}}}\right)$
	LB	$\Omega\left(\sqrt{\frac{d}{N\lambda_{k}}}+\sqrt{\frac{Md}{N^{2}\lambda_{k}^{2}}}\right)$
Well-represented Cases	UB	$\Theta\left(\sqrt{\frac{dk}{N}}+\sqrt{\frac{Mdk^{2}}{N^{2}}}\right)$
	LB

Table 1: Statistical error rate depends on the spectrum of

D

. Here, UB stands for Upper Bound, and LB stands for Lower Bound.

Figure 1: Subspace estimation error (

||\mathrm{sin}(B^{\star},B)||

) in a clustered setting. While pooling the full population maximizes sample size, uneven representation introduces bias. Conversely, a smaller balanced subset recovers the latent basis more effectively across all tested estimators. See Section 6 for setup and estimator details.

Figure 2: (Left) Performance on Clustered

\alpha_{i}

setting. (Right) Performance on Heterogeneous Gaussian

\alpha_{i}

setting.

Figure 3: (Left) Ablation over latent dimensionality $k$ . (Right) Ablation over full dimensionality.

Figure 4: (Left) Ablation over clustered sources assignment proportion. (Right) Ablation over the number of sources $M$ .

	ACSIncome (%)	CelebA (%)
Full Population	72.8	89.5
Random Selection	71.2	88.3
Power-of-Choice	73.0	89.8
Ours ( $k=3$ )	73.4	90.0
Ours ( $k=5$ )	73.9	90.3
Ours ( $k=7$ )	74.2	90.5

Table 2: Binary Classification Accuracy on Income and CelebA.

Appendix figures & tables3 assets

Supplementary material from the paper’s appendix.

Appendix

Reference	Method	General cases	Well-represented
Upper Bound:
( Du et al., 2021 )	ERM	-	$O(\sqrt{\frac{dk^{2}}{N}})$
( Tripuraneni et al., 2021 )	MoM	$O\Big(\sqrt{\frac{d}{N\lambda_{k}^{2}}}\Big)$	$O\Big(\sqrt{\frac{dk^{2}}{N}}\Big)$
( Duchi et al., 2022 )	Spectral	$O\Big(\sqrt{\frac{d}{N\lambda_{k}^{2}}}\Big)$	$O\Big(\sqrt{\frac{dk^{2}}{N}}\Big)$
( Collins et al., 2021 )	AltMin	-	$O\Big(\sqrt{\frac{dk^{2}}{N}}\Big)$
( Thekumparampil et al., 2021 )	AltMin	$O\Big(\sqrt{\frac{dk\lambda_{1}}{N\lambda_{k}^{2}}}\Big)$	$O\Big(\sqrt{\frac{dk^{2}}{N}}\Big)$

Appendix

Table 3: A summary of the high-probability statistical rate for estimating

B^{\star}

. Here “general cases” refer to cases with a general source diversity matrix

D

, while “well-represented” cases assume

\lambda_{1}=\Theta(\lambda_{k})=\Theta(1/k)

. The abbreviations for the methods are as follows: “ERM”: empirical risk minimization; “MoM”: method-of-moments estimator; “AltMin”: alternating minimization algorithm initialized from MoM; “AdaptRep”: adaptive representation learning; “ARMUL”: adaptive and robust multi-task learning.

	Mean Accuracy (%)
Selection Method	$d_{h}=64$	$d_{h}=256$
Full Population	79.91	79.49
Random Selection	78.08	77.87
Power-of-Choice	79.70	79.41
Ours	80.13	79.96

Appendix

Table 4: Classification accuracy on ACSIncome with a two-layer neural network feature extractor at hidden dimensions

d_{h}\in\{64,256\}

. Our source-screening procedure consistently outperforms all baselines, demonstrating that the benefits extend beyond the linear regime.

Selection Method	Mean Accuracy	Std. Accuracy
Subsampling	79.96	0.21
Reweight V1	79.38	0.03
Reweight V2	79.58	0.06
Full Population	79.49	0.06

Appendix

Table 5: Comparison of client selection strategies on the ACS Income Dataset (two-layer neural network).

Shared Representation Learning

See all 9 papers

Learning from Biased and Costly Data Sources: Minimax-optimal Data Collection under a Budget
Feb 19, 2026Michael O. Harding, Vikas Singh, Kirthevasan KandasamyMinimax Estimation Selection Bias
Department of Statistics, University of Wisconsin-Madison · Department of Biostatistics, University of Wisconsin-Madison · Department of Computer Sciences, University of Wisconsin-Madison
Sharp Regret Bounds and a Task-Covariance Correction for Spectral Representation Learning
Sep 14, 2026Dier Tang, Jing Yee Tan, Guangyue HanRepresentation Learning Self-Supervised Learning
Department of Mathematics, The University of Hong Kong
MIRA: Mid-training Rubric Anchoring for Source-Aware Data Selection
May 28, 2026Haowen Wang, Yaxin Du, Jian Yang +9Data Selection Efficient Language Model Training
1Beihang University · 2IQuest Research · 3Shanghai Jiao Tong University +2

Feb 19, 2026J/K move · Enter open · S save

Learning from Biased and Costly Data Sources: Minimax-optimal Data Collection under a Budget

Michael O. Harding, Vikas Singh, Kirthevasan Kandasamy

Department of Statistics, University of Wisconsin-Madison · Department of Biostatistics, University of Wisconsin-Madison · Department of Computer Sciences, University of Wisconsin-Madison

Minimax Estimation Selection Bias

Open paper AskPDF arXiv

Data collection is a critical component of modern statistical and machine learning pipelines, particularly when data must be gathered from multiple heterogeneous sources to study a target population of interest. In many use cases, such as medical studies or political polling, different sources incur different sampling costs. Observations often have associated group identities - for example, health markers, demographics, or political affiliations - and the relative composition of these groups may differ substantially, both among the source populations and between sources and target population. In this work, we study multi-source data collection under a fixed budget, focusing on the estimation of population means and group-conditional means. We show that naive data collection strategies (e.g. attempting to "match" the target distribution) or relying on standard estimators (e.g. sample mean) can be highly suboptimal. Instead, we develop a sampling plan which maximizes the effective sample size - the total sample size divided by

D_{χ^2}(q\mid\mid\overline{p}) + 1

, where

q

is the target distribution,

\overline{p}

is the aggregated source distribution, and

D_{χ^2}

is the

χ^2

-divergence. We pair this sampling plan with a classical post-stratification estimator and upper bound its risk. We provide matching lower bounds, establishing that our approach achieves the budgeted minimax optimal risk. Our techniques also extend to prediction problems when minimizing the excess risk, providing a principled approach to multi-source learning with costly and heterogeneous data sources.

Abstract

Shared Representation Learning

Explore similar work

Learning from Biased and Costly Data Sources: Minimax-optimal Data Collection under a Budget

Sharp Regret Bounds and a Task-Covariance Correction for Spectral Representation Learning

MIRA: Mid-training Rubric Anchoring for Source-Aware Data Selection