优化切片方向让高维分布对比更高效，减少计算量。
Constrained Sliced Wasserstein Embedding
- 约束切片方向使一维传输计划逼近原空间最优解。
- 通过梯度训练切片参数，在图像/点云/蛋白序列上提升效果。
- 适合需要压缩高维嵌入且保持排列不变性的场景。
Sliced Wasserstein (SW) 距离通过将高维概率分布投影到多个一维分布来高效比较。但识别有效切片方向一直困难，常需大量切片以达良好性能，增加计算开销。本文提出一种约束学习方法，优化 SW 的切片方向：约束一维传输计划近似原空间最优计划，确保切片方向有意义。利用这些传输计划的连续松弛，实现基于梯度的原始-对偶优化，联合训练切片参数与其他模型参数。实验表明，该方法能将高维嵌入池化为固定长度的排列不变表示。在图像、点云和蛋白序列上的基础模型上，验证了其学习更有效切片方向的能力。代码见 https://github.com/Stranja572/constrainedswe。
原文摘要 · Abstract (English)
Sliced Wasserstein (SW) distances offer an efficient method for comparing high-dimensional probability measures by projecting them onto multiple 1-dimensional probability distributions. However, identifying informative slicing directions has proven challenging, often necessitating a large number of slices to achieve desirable performance and thereby increasing computational complexity. We introduce a constrained learning approach to optimize the slicing directions for SW distances. Specifically, we constrain the 1D transport plans to approximate the optimal plan in the original space, ensuring meaningful slicing directions. By leveraging continuous relaxations of these transport plans, we enable a gradient-based primal-dual approach to train the slicer parameters, alongside the remaining model parameters. We demonstrate how this constrained slicing approach can be applied to pool high-dimensional embeddings into fixed-length permutation-invariant representations. Numerical results on foundation models trained on images, point clouds, and protein sequences showcase the efficacy of the proposed constrained learning approach in learning more informative slicing directions. Our implementation code can be found at https://github.com/Stranja572/constrainedswe.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。