Prior-Informed Generalizable Low-Rank Representation Learning via Regularized Adversarial Reweighting
Abstract
Modern learning problems increasingly leverage data from heterogeneous sources to improve robustness and generalizability under distribution shift, often with high-dimensional parameters admitting low-rank structure. Existing group distributionally robust optimization (group-DRO) approaches optimize against adverse source mixtures, but extending this principle to low-rank matrix estimation with a broad class of losses leads to a challenging nonconvex-concave problem for which methods and theory remain limited. We propose Prior-Informed generalizable Low-rank representation Learning via Adversarial Reweighting (PILLAR), which smooths the group-DRO objective by regularizing adversarial source weights toward a user-specified reference distribution while learning a shared low-rank representation. By tempering extreme adversarial reweighting, PILLAR improves estimation stability and reduces the potential conservativeness of group-DRO. Using Burer-Monteiro factorization, we optimize the nonconvex-concave objective by bifactored gradient descent with mirror ascent. We establish linear convergence to a statistical neighborhood of the population saddle point and derive an error bound that separates low-rank estimation from adversarial reweighting. We further show that PILLAR improves on vanilla group DRO when the latter's excess conservativeness dominates source heterogeneity and prior misspecification. Simulations and an electronic health record study of suicide-attempt risk demonstrate improved worst-source and transported predictive accuracy.
Related
- Joint Contrastive-Adversarial Representation Learning for Unsupervised Domain Adaptation
- A Composite Representation Theorem for Solution Mappings in Regularized Composite Optimization with Applications to Hausdorff Lipschitz Stability
- PaECTER: Patent-level representation learning using citation-informed transformers
- Fast Calibration of Generalized Posteriors via Dirichlet Reweighting
- Q-MMR: Off-Policy Evaluation via Recursive Reweighting and Moment Matching
BibTeX
@misc{zhan2026priorinformed,
title = {{Prior-Informed Generalizable Low-Rank Representation Learning via Regularized Adversarial Reweighting}},
author = {Keyao Zhan and Junwei Lu and Tianxi Cai},
howpublished = {arXiv},
year = {2026},
month = oct,
doi = {10.48550/arxiv.2610.14414},
eprint = {2610.14414},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2610.14414},
note = {CC BY-NC-ND 4.0},
}
RIS
TY - UNPB
TI - Prior-Informed Generalizable Low-Rank Representation Learning via Regularized Adversarial Reweighting
AU - Keyao Zhan
AU - Junwei Lu
AU - Tianxi Cai
PY - 2026
DA - 2026/10/10
PB - arXiv
DO - 10.48550/arxiv.2610.14414
UR - https://arxiv.org/abs/2610.14414
AB - Modern learning problems increasingly leverage data from heterogeneous sources to improve robustness and generalizability under distribution shift, often with high-dimensional parameters admitting low-rank structure. Existing group distributionally robust optimization (group-DRO) approaches optimize against adverse source mixtures, but extending this principle to low-rank matrix estimation with a broad class of losses leads to a challenging nonconvex-concave problem for which methods and theory remain limited. We propose Prior-Informed generalizable Low-rank representation Learning via Adversarial Reweighting (PILLAR), which smooths the group-DRO objective by regularizing adversarial source weights toward a user-specified reference distribution while learning a shared low-rank representation. By tempering extreme adversarial reweighting, PILLAR improves estimation stability and reduces the potential conservativeness of group-DRO. Using Burer-Monteiro factorization, we optimize the nonconvex-concave objective by bifactored gradient descent with mirror ascent. We establish linear convergence to a statistical neighborhood of the population saddle point and derive an error bound that separates low-rank estimation from adversarial reweighting. We further show that PILLAR improves on vanilla group DRO when the latter's excess conservativeness dominates source heterogeneity and prior misspecification. Simulations and an electronic health record study of suicide-attempt risk demonstrate improved worst-source and transported predictive accuracy.
ER -