Technical Summary: AVATAR – Learning to Align via Active Optimal Transport

1. Problem Definition

Alignment is a fundamental task in machine learning, encompassing multi-network analysis, multimodal learning, and point cloud registration. The goal is to identify meaningful correspondences between two sets of data points. Recently, Optimal Transport (OT) has emerged as a robust framework for solving these problems by inferring object-level alignment from a transport plan derived from a cost function and marginal constraints.

However, OT-based alignment methods are highly sensitive to the quantity and quality of supervision. In practice, obtaining high-quality supervision (e.g., ground-truth correspondences) is often costly or infeasible, leading to a reliance on sparse supervision. While active learning aims to select the most informative samples to query, existing active alignment methods suffer from three key limitations when applied to OT:

1. Lack of OT-specific tailoring: Existing methods do not utilize key OT components (e.g., cost functions, marginal constraints) effectively.
2. Absence of principled impact quantification: There is no rigorous method to quantify how a new supervision signal propagates through the constrained OT formulation to affect the global alignment result.
3. Poor generalizability: Prior active alignment strategies are often task-specific (e.g., designed only for network consistency) and do not generalize well to other OT-based tasks or domains.

The paper addresses the problem of Active OT-based Alignment: selecting a fixed budget of source objects to query from an oracle (e.g., a human annotator) to maximize the improvement in alignment performance for the remaining unlabeled objects.

2. Methodology: AVATAR

The authors propose AVATAR (Active Optimal Transport via Alignment via Active Optimal Transport), a principled framework designed intrinsically for OT-based alignment. The core innovation lies in quantifying the informativeness of a candidate query by measuring its posterior impact on the global alignment result via gradient propagation.

2.1. Gradient-Based Query Impact

AVATAR evaluates a candidate source object $x_i$ by estimating how querying its true alignment would change a global utility function $f(T)$, where $T$ is the transport plan.

• Pairwise Impact: The impact of querying a specific pair $(x_i, y_j)$ is defined as the gradient of the utility function with respect to the supervision signal $H_{i,j}$ (i.e., $\nabla_{H_{i,j}} f$).
• Chain Rule Decomposition: This gradient is decomposed into the impact of the supervision on the cost matrix ($\nabla_{H_{i,j}} \tilde{C}$) and the impact of the cost matrix on the utility function ($\nabla_{\tilde{C}} f$).
• The Challenge: Differentiating through the OT formulation is difficult because the transport plan $T$ is an implicit function of the cost matrix $\tilde{C}$, defined by a large-scale constrained optimization problem. Explicitly computing the Jacobian $\frac{dT}{d\tilde{C}}$ is computationally intractable ($O((nm)^2)$).

2.2. Adjoint-State Method and Linear Systems

To overcome the differentiation challenge, AVATAR leverages the adjoint-state method to reformulate the gradient computation.

• Linear System Reformulation: Instead of inverting a massive Jacobian, the method derives a linear system of size $(n+m)$ involving adjoint vectors $y_\alpha$ and $y_\beta$. The gradient $\nabla_{\tilde{C}} f$ is computed as:
  $$\nabla_{\tilde{C}} f = \frac{1}{\epsilon} T \odot (y_\alpha \mathbf{1}_m^\top + \mathbf{1}_n y_\beta^\top - \nabla_T f)$$
  where $\epsilon$ is the entropic regularization weight.
• Solving the System: The coefficient matrix of this linear system is singular but positive-semidefinite. The authors prove that the Conjugate Gradient (CG) method can solve this system with guaranteed convergence to the global optimum.
• Complexity: By exploiting the sparsity of the transport plan $T$ (which is typically sparse in OT applications), the CG method achieves linear time complexity $O(K(n+m))$ per query batch, where $K$ is the number of CG iterations.

2.3. Posterior Object Query Impact

For a source object $x_i$, the total impact is the aggregation of pairwise impacts weighted by the current transport plan probabilities $T_{i,j}$ (acting as a posterior distribution over potential targets):
$$I(p_i) = \sum_{j=1}^m T_{i,j} I(p_{ij})$$
This allows AVATAR to select the source object with the highest expected impact on the global alignment quality.

2.4. Utility Functions

AVATAR is task-agnostic, utilizing effective utility functions $f(T)$ to encode alignment quality:

• $f_{L2}$: Encourages deterministic alignment (approximating a permutation matrix).
• $f_{entropy}$: Encourages alignment consistent with the entropy-regularized OT objective.
• $f_{consist}$: Specifically designed for Network Alignment, enforcing consistency principles where neighbors of aligned nodes should share similar alignment results.

3. Key Contributions

1. Problem Formalization: The paper is the first to formalize the problem of active learning specifically for OT-based alignment.
2. Novel Method (AVATAR): Proposes a method to quantify candidate informativeness via gradient propagation through the entropy-regularized OT formulation, utilizing the adjoint-state method to handle the implicit constraints.
3. Theoretical Analysis: Provides proofs establishing the correctness of the gradient computation, the linear time complexity of the solution, and the linear convergence rate of the Conjugate Gradient method used.
4. Empirical Validation: Demonstrates effectiveness across three diverse alignment tasks: Network Alignment (NA), Image-Text Retrieval, and Image-Text Grounding.

4. Experimental Results

Experiments were conducted on 8 datasets using 4 OT-based alignment algorithms and 9 baseline active learning methods.

• Effectiveness: AVATAR consistently outperformed existing active learning approaches (including Random, Entropy, Margin, and task-specific methods like GibbsMatchings) under the same query budget.
  • In Network Alignment, AVATAR improved Mean Reciprocal Rank (MRR) by up to 2.8% over baselines.
  • In Image-Text tasks, AVATAR achieved state-of-the-art Recall@1, often outperforming baselines by up to 2.6%.
• Efficiency and Scalability: AVATAR demonstrated a strong balance between performance and efficiency. It achieved up to a 25x speedup in query time compared to some baselines (e.g., TopMatchings) while improving MRR by up to 8.1%. Scalability tests on synthetic graphs confirmed linear scaling with respect to the number of nodes.
• Convergence: Empirical analysis showed that the Conjugate Gradient method used to solve the adjoint system converges rapidly and is less sensitive to the entropic regularization weight $\epsilon$ than the Sinkhorn algorithm used for the OT solver itself.
• Ablation Studies:
  • Sparse vs. Dense: Using sparse matrix operations significantly improved efficiency (up to 8.6x speedup) without compromising accuracy.
  • Posterior Aggregation: Using the transport plan $T$ as a posterior weight for aggregation proved superior to uniform aggregation.

5. Significance and Claims

The paper claims that AVATAR provides a principled, task-agnostic framework for active learning in the context of Optimal Transport. Its significance lies in:

• Bridging the Gap: It addresses the critical limitation of existing active alignment methods by explicitly modeling the propagation of supervision through the OT formulation.
• Scalability: By reformulating the gradient computation into a solvable linear system with linear complexity, it makes active learning feasible for large-scale alignment problems where previous methods were computationally prohibitive.
• Generalizability: The framework is applicable to diverse alignment tasks (networks, multimodal data) by simply changing the utility function, rather than requiring task-specific query heuristics.

The authors conclude that AVATAR effectively improves alignment performance with minimal supervision, offering a robust solution for scenarios where high-quality labels are scarce or expensive to obtain.