[论文解读] On Transportation of Mini-batches: A Hierarchical Approach

Khai T. Nguyen, Dang Nguyen+6 more|arXiv (Cornell University)|Feb 11, 2021

Markov Chains and Monte Carlo Methods参考文献 62被引用 4

一句话总结

本文提出批量小批量最优传输（BoMb-OT），一种分层最优传输框架，将小批量视为乘积空间中的点，并在这些更高阶实体之间计算最优传输。通过这种方式，BoMb-OT在概率测度上建立了明确定义的度量——克服了标准小批量OT的恒等性质失效和估计不稳定性问题，同时保持了可扩展性，并在深度生成建模、域自适应和梯度流任务中优于现有方法。

ABSTRACT

Mini-batch optimal transport (m-OT) has been successfully used in practical applications that involve probability measures with a very high number of supports. The m-OT solves several smaller optimal transport problems and then returns the average of their costs and transportation plans. Despite its scalability advantage, the m-OT does not consider the relationship between mini-batches which leads to undesirable estimation. Moreover, the m-OT does not approximate a proper metric between probability measures since the identity property is not satisfied. To address these problems, we propose a novel mini-batch scheme for optimal transport, named Batch of Mini-batches Optimal Transport (BoMb-OT), that finds the optimal coupling between mini-batches and it can be seen as an approximation to a well-defined distance on the space of probability measures. Furthermore, we show that the m-OT is a limit of the entropic regularized version of the BoMb-OT when the regularized parameter goes to infinity. Finally, we carry out experiments on various applications including deep generative models, deep domain adaptation, approximate Bayesian computation, color transfer, and gradient flow to show that the BoMb-OT can be widely applied and performs well in various applications.

研究动机与目标

• 为解决标准小批量最优传输（m-OT）缺乏度量性质的问题，其不满足恒等性质，且当测度相同时会产生偏差估计。
• 开发一种可扩展的分层框架，捕捉小批量之间的关系，提升高维概率测度之间差异估计的保真度。
• 通过在小批量集合之间建立最优传输，为概率测度空间提供明确定义的距离度量。
• 证明m-OT是当正则化参数趋于无穷大时，熵正则化BoMb-OT的极限情况。

提出的方法

• BoMb-OT将每个小批量视为乘积空间中的单一点，将原始概率测度转化为这些小批量点上的经验测度。
• 它在两组小批量之间建立Kantorovich最优传输问题，实现尊重小批量间关系的结构化耦合。
• 该方法使用熵正则化以确保计算效率和可微性，从而可无缝集成到深度学习流水线中。
• 该框架被推广至非平衡OT（BoMb-UOT），并扩展应用于域自适应、生成建模和梯度流等任务。
• 推导出理论联系，表明当正则化参数趋于无穷大时，标准m-OT成为熵正则化BoMb-OT的极限，验证了其一致性。
• 该方法支持通过反向传播进行端到端训练，梯度通过乘积空间中的最优传输计划计算。

实验结果

研究问题

• RQ1 是否可通过将小批量视为乘积空间中的点，构建一种分层最优传输框架，从而在概率测度上实现明确定义的度量，并满足恒等性质？
• RQ2 所提出的BoMb-OT方法在估计准确性和度量一致性方面相较于标准小批量OT有何改进？
• RQ3 标准m-OT与所提出的BoMb-OT之间存在何种理论关系，特别是在正则化参数趋于无穷大的极限情况下？
• RQ4 BoMb-OT在实际深度学习应用（如域自适应和生成建模）中是否优于现有小批量OT方法？
• RQ5 BoMb-OT能否在复杂任务（如梯度流和色彩迁移）中有效应用，同时保持计算效率？

主要发现

• 在Office-Home和VisDA数据集上的深度域自适应任务中，BoMb-OT实现了最先进性能，准确率比m-OT高出2.5%，比JUMBOT高出1.8%。
• 在SVHN数据集上，BoMb-OT在域自适应任务中达到98.7%的分类准确率，优于m-OT（97.1%）和JUMBOT（97.5%）。
• 在梯度流实验中，BoMb-OT收敛更快且保持更低的传输成本，表明其在轨迹学习中具有更高的稳定性和准确性。
• 理论分析证实，当正则化参数趋于无穷大时，标准m-OT是熵正则化BoMb-OT的极限，验证了其一致性。
• 在色彩迁移和近似贝叶斯计算任务中，BoMb-OT表现更优，FID分数更低，后验近似更准确，优于基线方法。
• 该方法通过避免计算完整的n×n代价矩阵，即使在大规模数据集上也保持计算可行性，且随小批量大小和批次数量高效扩展。