AvAtar: Learning to Align via Active Optimal Transport
Abstract
Alignment plays a fundamental role in many machine learning problems, such as multi-network analysis, multimodal learning, and point cloud registration. Recent works increasingly leverage optimal transport (OT) for distributional alignment, whose effectiveness largely depends on sparse supervision that is hard or costly to obtain in practice. Existing works, however, largely overlook how to actively acquire high-quality supervision to improve their alignment performance under OT frameworks. In this paper, we propose a principled active alignment framework for optimal transport alignment called AvAtar. We quantify the informativeness of a candidate by measuring its gradient-based impact on the global alignment result, computed as the gradient propagation from the global alignment result to all possible supervisions of the candidate through the entropy-regularized OT formulation. While differentiating through OT is challenging given its constrained nature, we leverage the adjoint-state method to reformulate the computation to a linear system solvable by the conjugate gradient method with linear complexity and guaranteed convergence. By encoding the global alignment result via effective utility functions, AvAtar is applicable to general alignment problems under the OT framework. Extensive experiments on three representative alignment tasks demonstrate the effectiveness, scalability, and generalizability of the proposed AvAtar.
Lay Summary
Many AI systems need to match related information across different sources, such as linking users across social networks or connecting images with their text descriptions. Recent alignment methods often rely on a mathematical framework called optimal transport. However, these methods are highly sensitive to the small amount of human supervision they receive, and collecting such labels can be expensive. We introduce AvAtar, a method that intelligently decides which examples are most valuable to ask humans to label in optimal-transport-based alignment systems. Instead of selecting examples randomly, Instead of choosing examples based only on uncertainty or randomness, AVATAR estimates how much a new label would improve the overall alignment structure learned by the model. This allows the method to focus on labels that have the greatest global impact on the matching results. We design the method to remain efficient even for very large datasets, making it practical for real-world applications. Experiments on multiple tasks, including network matching and image-text alignment, show that AvAtar consistently improves performance while requiring less human supervision.