Group Distributionally Robust Optimization with Flexible Sample Queries¶
会议: NeurIPS 2026
arXiv: 2505.15212
领域: 优化算法 / Optimization
关键词: 分布鲁棒优化 (GDRO), 有限建议预测 (PLA), 在线凸优化, 灵活样本查询, 随时可停算法 (Anytime)
一句话总结¶
将群体分布鲁棒优化建模为在线凸优化与非遗忘有限建议预测的双人博弈,提出了支持动态任意样本查询的统一 FTRL 算法,首创非遗忘 PLA 的高概率后悔界,并将 GDRO 优化误差提升至消除对数时间因子的最优阶。
研究背景与动机¶
传统的机器学习模型依赖于经验风险最小化(ERM),通常假设训练与测试样本服从同分布。然而在真实场景中,环境演化、协变量偏移或亚群组偏差经常导致测试分布发生漂移,使模型性能剧烈下滑。群体分布鲁棒优化(Group Distributionally Robust Optimization, GDRO)通过在有限个候选数据分布构成的预设不确定性集上最小化最坏分布风险,成为了应对分布偏移和提升亚群公平性的重要理论与算法范式。
然而,现有 GDRO 算法在采样机制上面临着严重的刚性约束:它们在每一轮迭代中只能处理固定数量的样本,要么是仅查询 1 个样本(如随机镜像下降或非遗忘在线博弈框架),要么是必须同时查询全部 \(m\) 个分布的样本。在诸如跨机构医疗诊断或不可信网络带宽下的不可知联邦学习(Agnostic Federated Learning)等实际场景中,受限于通信带宽或隐私准入,每一轮可获取的群组样本数通常是动态波动的(例如某轮只能访问 \(r_t\) 个群组,且 \(1 \le r_t \le m\))。现存方法若硬套 \(m\) 个样本方案则资源受限无法执行;若强行退化为单样本串行处理,则不仅丧失了并行采样加速的机会,还会引入更慢的收敛速度与额外的计算延迟。
这种采样灵活性与理论分析可行性之间的矛盾,根源于非遗忘环境中随时间变化的样本数对在线对抗博弈参与者带来的步长设计与后悔界分析难度。本文从博弈论和在线学习的视角切入,将具有灵活样本查询的 GDRO 重新表述为一个双人零和博弈问题。核心 idea:将 GDRO 建模为非遗忘在线凸优化(模型侧)与非遗忘有限建议预测(分布侧)的双人博弈,基于 FTRL 设计兼容单样本与多样本损失估计器的统一策略,实现首个支持任意时变样本查询量的高概率界 GDRO 算法。
方法详解¶
整体框架¶
本文将群体分布鲁棒优化问题转化为一个二人博弈框架:其中 \(\mathbf{w}\)-player 负责模型参数更新,求解一个带有随机梯度的非遗忘在线凸优化(Online Convex Optimization, OCO)问题;而 \(\mathbf{q}\)-player 负责维护各个分布的重要性权重,求解非遗忘有限建议预测(Prediction with Limited Advice, PLA)问题。在每一轮 \(t\) 中,环境提供当前轮允许查询的样本配额 \(r_t\)。两个玩家分别提交决策 \(\mathbf{w}_t\) 和 \(\mathbf{q}_t\),随后通过依赖性舍入(DepRound)等采样机制从选定的分布集合中采集样本,计算损失并互相反馈梯度与损失估计值。
为了克服传统随机镜像下降(SMD)在动态时变采样序列下加权输出难以分析且会引入额外 \(O(\log t)\) 累积因子的缺陷,本文创新性地为两个玩家均采用基于正则化跟随领导者(Follow-the-Regularized-Leader, FTRL)的随时可停(anytime)架构,并以均匀简单时间平均 \((\bar{\mathbf{w}}_t, \bar{\mathbf{q}}_t)\) 作为输出。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["环境输入:时变采样预算 r_t"] --> B["分布策略更新:基于负熵正则的 FTRL 计算 q_t"]
B --> C["分层采样机制:核心分布 c_t 与无放回补全集合 I_t"]
C --> D["样本查询与反馈:采集 C_t 中的样本并计算随机损失"]
D --> E["双路损失估计与累积:区分单样本 IX 与多样本无偏估计"]
D --> F["模型策略更新:复用 c_t 样本梯度的 FTRL 投影更新 w_t"]
E --> G["双人博弈闭环:输出无加权时间平均解 w_bar_t 与 q_bar_t"]
F --> G
关键设计¶
1. 统一非遗忘 PLA 决策更新:基于负熵 FTRL 的闭式权重分布生成
在每一轮迭代中,分布侧玩家需要根据历史估计的累积损失向量 \(\mathbf{L}_{t-1}\) 输出单纯形 \(\Delta_m\) 上的概率分布 \(\mathbf{q}_t\)。传统方法在应对单样本与多样本混合的非遗忘环境时难以收敛,本文使用负熵正则化项 \(\nu_q(\mathbf{q}) = \sum_{i=1}^m q_i \ln q_i\),将 FTRL 问题转化为闭式解: $\(q_{t,i} = \frac{\exp(-\eta_{q,t} L_{t-1,i})}{\sum_{k=1}^m \exp(-\eta_{q,t} L_{t-1,k})}\)$ 为了自适应于历史中非均匀且波动的查询配额序列 \(\{r_j\}_{j=1}^t\),算法配置了动态递减步长 \(\eta_{q,t} = \sqrt{\frac{\ln m}{m \sum_{j=1}^t \frac{1}{r_j}}}\)。该更新方式消除了对总迭代次数 \(T\) 的预先依赖,使算法具备随时可停特性,并且有效控制了正则化带来的稳定性折中。
2. 混合驱动的分层采样:结合主分布抽取与依赖性舍入的无放回集合构建
当环境给定查询配额 \(r_t \in [m]\) 后,如何从权重分布 \(\mathbf{q}_t\) 扩展到 \(r_t\) 个候选分布至关重要。本文设计了两阶段采样机制:首先严格按照概率向量 \(\mathbf{q}_t\) 从 \([m]\) 中抽取一个首要分布索引 \(c_t\);当 \(r_t \ge 2\) 时,再从剩余的 \(m-1\) 个分布集合 \([m] \setminus \{c_t\}\) 中,利用依赖性舍入算法(Dependent Rounding, DepRound)无放回且均匀随机地抽取额外的 \(r_t - 1\) 个分布索引,记为集合 \(\mathcal{I}_t\)。最终查询的分布全集为 \(C_t = \{c_t\} \cup \mathcal{I}_t\)。该机制保证了主要探索方向由当前对抗权重主导,同时其余配额在时间与空间上均等分摊,为后续构造有界的损失估计提供了精准的边际概率保证。
3. 区分采样规模的统一自适应损失估计器:融合隐式探索与无偏缩放
在非遗忘对抗环境下,若仅查询 1 个样本,\(q_{t,i}\) 过小会导致标准重要性加权估计器的方差爆炸;而当 \(r_t \ge 2\) 时,由于各臂有保底的被抽样概率,方差可以得到天然压制。为此,本文设计了自适应统一损失估计器 \(\tilde{\mathbf{s}}_t\): $\(\tilde{s}_{t,i} = \begin{cases} \displaystyle \frac{\hat{s}_{t,i}}{q_{t,i} + \gamma_t} \mathbb{I}[i \in C_t], & r_t = 1 \\[10pt] \displaystyle \frac{\hat{s}_{t,i}}{q_{t,i} + (1 - q_{t,i})\frac{r_t - 1}{m - 1}} \mathbb{I}[i \in C_t], & r_t \ge 2 \end{cases}\)$ 其中 \(\hat{s}_{t,i} = 1 - \ell(\mathbf{w}_t; \mathbf{z}_t^{(i)}) \in [0, 1]\),隐式探索(IX)参数设为 \(\gamma_t = \frac{1}{2}\eta_{q,t}\)。当 \(r_t=1\) 时引入受控的有偏隐式探索抑制极端方差;当 \(r_t \ge 2\) 时直接利用真实边际抽样概率构造完全无偏估计。最后将估计向量增量累加至统一的累积损失 \(\mathbf{L}_t = \mathbf{L}_{t-1} + \tilde{\mathbf{s}}_t\) 中,成功在单个进程内统一了 Exp3-IX 与 PrLiA 两种异构算法。
4. 样本复用的 OCO 求解器与均匀平均输出:彻底消除时间对数因子
对于模型参数玩家 \(\mathbf{w}_t\),本文并不产生额外的环境查询,而是直接复用 \(\mathbf{q}\)-player 抽取的首要分布样本 \(\mathbf{z}_t^{(c_t)}\) 构造无偏随机次梯度 \(\mathbf{g}_w(\mathbf{w}_t, \mathbf{q}_t) = \nabla \ell(\mathbf{w}_t; \mathbf{z}_t^{(c_t)})\)。采用强凸正则化项 \(\nu_w(\mathbf{w}) = \frac{1}{2}\|\mathbf{w}\|_2^2\),参数通过欧氏投影保持在假设空间 \(\mathcal{W}\) 内:\(\mathbf{w}_t = \Pi_{\mathcal{W}} [-\eta_{w,t} \mathbf{F}_{t-1}]\),其中梯度累加量为 \(\mathbf{F}_{t-1}\),步长设为 \(\eta_{w,t} = \frac{\sqrt{2}D}{\sqrt{5}G\sqrt{t}}\)。最终算法输出采用未加权的简单算术平均 \((\bar{\mathbf{w}}_t, \bar{\mathbf{q}}_t) = (\frac{1}{t}\sum_{j=1}^t \mathbf{w}_j, \frac{1}{t}\sum_{j=1}^t \mathbf{q}_j)\)。这一策略避免了 SMD 历史加权累积产生的 \(\sum_{j=1}^t \frac{1}{j}\) 发散项,从而将高概率优化误差界的分子从 \(O(\log t)\) 直接精炼消除。
损失函数 / 训练策略¶
算法以零和博弈的目标函数 \(\phi(\mathbf{w}, \mathbf{q}) = \sum_{i=1}^m q_i R_i(\mathbf{w})\) 为核心,通过每一轮交替接收反馈以在线方式协同更新。完整的训练循环如下: 1. 环境揭示当轮采样预算 \(r_t\); 2. \(\mathbf{q}\)-player 计算权重 \(\mathbf{q}_t\),抽取 \(c_t\) 并通过 DepRound 补齐 \(\mathcal{I}_t\),将查询集 \(C_t\) 发送给环境; 3. \(\mathbf{w}\)-player 计算参数 \(\mathbf{w}_t\); 4. 环境反馈 \(C_t\) 中各个分布的样本损失 \(\hat{s}_{t,i}\) 以及针对 \(c_t\) 的随机梯度 \(\nabla \ell(\mathbf{w}_t; \mathbf{z}_t^{(c_t)})\); 5. 分别更新损失估计累积量 \(\mathbf{L}_t\) 与梯度累积量 \(\mathbf{F}_t\),并即时维护输出均值 \(\bar{\mathbf{w}}_t, \bar{\mathbf{q}}_t\)。
实验关键数据¶
主实验¶
实验在合成二分类数据集(\(m=20\) 个分布,维度 500)与真实世界糖尿病多分类数据集(Diabetes,包含 99,493 个样本、114 个特征,基于种族、年龄与性别划分为 \(m=12\) 个群组)上展开,重复 5 次求平均。评估指标为在所有分布上的最坏情况风险(Max Risk)。
在动态资源场景(每轮随机采样 \(r_t \in [1, m-1]\))下,基线方法 SMD(\(m\)) 无法运行,各算法在两数据集上的收敛指标与性能对比见下表:
| 数据集 | 算法配置 | 每轮更新样本数 | 迭代收敛速度 (Iterations to Converge) | 最终稳态最坏风险 (Max Risk) | 相同运行时间下的收敛效率 |
|---|---|---|---|---|---|
| Synthetic (\(m=20\)) | Online(1) | 1 (固定) | 较慢 (\(\mathcal{O}(\sqrt{m \log m / t})\)) | 0.385 | 低(串行更新吞吐不足) |
| Synthetic (\(m=20\)) | Online(1)' | \(r_t\) (串行多次) | 较快(启发式拆分) | 0.382 | 较慢(缺乏步长理论自适应) |
| Synthetic (\(m=20\)) | HYB (Ours) | \(r_t\) (灵活并行) | 极快 (\(\mathcal{O}(\frac{1}{t}\sqrt{\sum \frac{m}{r_j}\log m})\)) | 0.372 | 最高(计算与采样高度并行) |
| Synthetic (\(m=20\)) | UNI (Ours) | \(r_t\) (灵活并行) | 极快 (\(\mathcal{O}(\frac{1}{t}\sqrt{\sum \frac{m}{r_j}\log m})\)) | 0.372 | 最高(曲线与 HYB 完全吻合) |
| Diabetes (\(m=12\)) | Online(1) | 1 (固定) | 较慢 | 0.812 | 低 |
| Diabetes (\(m=12\)) | Online(1)' | \(r_t\) (串行多次) | 中等 | 0.806 | 中等(无步长保证) |
| Diabetes (\(m=12\)) | HYB (Ours) | \(r_t\) (灵活并行) | 最快 | 0.795 | 最高 |
| Diabetes (\(m=12\)) | UNI (Ours) | \(r_t\) (灵活并行) | 最快 | 0.795 | 最高 |
消融实验¶
固定资源配额消融实验中,验证不同固定采样量 \(r \in [m]\) 对收敛速度与样本复杂度的影响。以提出的统一算法 UNI 为核心,考察在 Synthetic 数据集(\(r \in \{1, 5, 10, 15, 20\}\))和 Diabetes 数据集(\(r \in \{1, 4, 7, 10, 12\}\))上的收敛行为:
| 实验配置 (UNI Algorithm) | 采样规模 \(r\) | 理论误差上界阶数 | 达到收敛所需迭代数 (Relative Iterations) | 最终样本复杂度 (Sample Complexity) | 收敛特性说明 |
|---|---|---|---|---|---|
| Synthetic (\(m=20\)) | \(r=1\) | \(\mathcal{O}(\sqrt{m\log m / t})\) | \(1.00\times\) (基准) | \(\mathcal{O}(m\log m / \epsilon^2)\) | 递减步长初期降损快于固定步长 Online(1) |
| Synthetic (\(m=20\)) | \(r=5\) | \(\mathcal{O}(\sqrt{m\log m / (5t)})\) | \(\sim 0.22\times\) | \(\mathcal{O}(m\log m / \epsilon^2)\) | 迭代数显著减少,总样本消耗持平 |
| Synthetic (\(m=20\)) | \(r=10\) | \(\mathcal{O}(\sqrt{m\log m / (10t)})\) | \(\sim 0.11\times\) | \(\mathcal{O}(m\log m / \epsilon^2)\) | 与全量更新相比兼具轻量与快速收敛 |
| Synthetic (\(m=20\)) | \(r=20\) (\(m\)) | \(\mathcal{O}(\sqrt{\log m / t})\) | \(\sim 0.05\times\) | \(\mathcal{O}(m\log m / \epsilon^2)\) | 最终精度与经典 SMD(\(m\)) 完全一致 |
| Diabetes (\(m=12\)) | \(r=1\) | \(\mathcal{O}(\sqrt{m\log m / t})\) | \(1.00\times\) (基准) | \(\mathcal{O}(m\log m / \epsilon^2)\) | 单样本稳健基线 |
| Diabetes (\(m=12\)) | \(r=4\) | \(\mathcal{O}(\sqrt{m\log m / (4t)})\) | \(\sim 0.27\times\) | \(\mathcal{O}(m\log m / \epsilon^2)\) | 资源敏感环境下的最优权衡点 |
| Diabetes (\(m=12\)) | \(r=7\) | \(\mathcal{O}(\sqrt{m\log m / (7t)})\) | \(\sim 0.15\times\) | \(\mathcal{O}(m\log m / \epsilon^2)\) | 显著优于启发式串行重采样 |
| Diabetes (\(m=12\)) | \(r=12\) (\(m\)) | \(\mathcal{O}(\sqrt{\log m / t})\) | \(\sim 0.08\times\) | \(\mathcal{O}(m\log m / \epsilon^2)\) | 达成最快迭代收敛率 |
关键发现¶
- 并行采样显著压缩训练物理耗时:在相同运行时间内,UNI 和 HYB 均能比 Online(1) 与串行启发式 Online(1)' 达到更低的最大风险。这证明了能够同时处理多个动态样本的算法在计算和 IO 层面具有压倒性吞吐优势。
- 样本复杂度高度一致:无论是混合算法 HYB 还是统一算法 UNI,当固定每轮样本量为 \(r\) 时,虽然单轮耗费样本量不同,但最终实现误差 \(\epsilon\) 所需的全局样本总数稳定保持在 \(\mathcal{O}(m \log(m) / \epsilon^2)\),完美衔接了以往 \(r=1\) 与 \(r=m\) 的极端边界情况。
- UNI 与 HYB 表现高度等价但实现更轻简:在所有实验轨迹中,维护单一累积损失向量的统一算法(UNI)与分别维护单/多样本状态的混合算法(HYB)在收敛曲线与数值上高度契合,从而验证了统一表述在不牺牲任何收敛速率的前提下显著降低了维护复杂度。
亮点与洞察¶
- 博弈视角的桥梁作用:将采样受限的 GDRO 精确解耦为 OCO 与有限建议预测(PLA)的二人零和博弈,通过集中不等式将随机性从对抗对抗性中剥离,使得组合与在线学习的前沿工具能够模块化无缝接入。
- 巧妙的统一损失估计机制:利用 IX 修正项在 \(r_t=1\) 时压制分母为 0 的方差爆炸,而在 \(r_t \ge 2\) 时依靠无放回采样的天然下界构造无偏估计,用一条平滑过渡的更新式统一了两种看似截然不同的 bandit 与 expert 算法。
- 时序对数因子的彻底消除:敏锐地捕捉到了 SMD 加权平均输出导致的历史累积对数惩罚,果断替换为 FTRL 结合简单均匀时间平均输出,一举将随时可停优化误差从 \(\mathcal{O}(\log(t)\sqrt{m\log m / (rt)})\) 改进为紧致的 \(\mathcal{O}(\sqrt{m\log m / (rt)})\)。
局限与展望¶
- 理论对序列遗忘性的假设约束:本文针对全部时间点同时成立的时间均匀高概率界(Theorem 5)在证明时需要假设样本容量序列 \(\{r_t\}\) 是不依赖于内部随机性的无意识(oblivious)序列,未来需要探索如何放宽这一假设以完全适应自适应攻击者。
- 模型参数更新的样本利用率:目前 \(\mathbf{w}\)-player 仅复用了由 \(\mathbf{q}_t\) 抽取的主样本 \(c_t\) 的梯度,而没有将额外采样的 \(r_t - 1\) 个样本的梯度充分加权融合,未来可通过方差缩减(SVRG/SARAH)等技术进一步加速模型侧收敛。
- 拓展至大语言模型对齐与 MRO:该动态灵活查询机制天然契合具有异构数据源混合的大语言模型预训练、强化学习多偏好对齐(RLHF)以及超额风险极小化(MRO)框架,具备广阔的拓展潜力。
相关工作与启发¶
- vs Nemirovski et al. (2009): 经典方法使用随机镜像下降(SMD)求解 GDRO,必须在每轮全量查询所有 \(m\) 个分布;本文将其彻底扩展为任意且动态波动的样本量 \(r_t \in [m]\),并在 \(r=m\) 时仍保持同阶最优样本复杂度。
- vs Zhang et al. (2023): 该工作引入非遗忘在线博弈求解 GDRO,但严格限制每轮仅能查询 1 个样本;本文不仅兼容了 \(r=1\) 的情景,还支持多样本并行查询加速,并将 anytime 算法的收敛速率提升了 \(\mathcal{O}(\log t)\) 因子。
- vs Seldin et al. (2014): Seldin 等人提出的 PrLiA 算法仅在遗忘对手与期望意义下给出了后悔界;本文首次建立起了面向非遗忘对手(Non-oblivious Adversary)的 PLA 高概率后悔分析,填补了该领域的理论空白。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ 成功打通了时变样本查询机制与 GDRO 理论保证,并给出了首个非遗忘 PLA 高概率界。
- 实验充分度: ⭐⭐⭐⭐☆ 在合成数据集与真实多属性公平性医疗数据集上均进行了深入的时变与固定资源对比验证。
- 写作质量: ⭐⭐⭐⭐⭐ 理论推导清晰严谨,从博弈建模、单算法扩展到统一算法的演进逻辑层层递进。
- 价值: ⭐⭐⭐⭐⭐ 为动态异构资源下的稳健机器学习、联邦学习及多源数据混合优化提供了关键的基础理论支撑。