Source-linked AI summary
Unsupervised Alignment of Embeddings with Wasserstein Procrustes
Edouard Grave, Armand Joulin, Quentin Berthet
TL;DR
The paper addresses unsupervised alignment of high-dimensional point clouds, including word embeddings for bilingual lexicon induction. It jointly estimates an orthogonal transformation and permutation, initializes optimization with a graph-matching convex relaxation, and uses a stochastic algorithm. On unsupervised word translation, the method reaches state-of-the-art-level performance while using an efficient optimization procedure.
Problem
High-dimensional embedding alignment is needed for applications such as unsupervised translation and computer vision, while existing Wasserstein- and adversarial-based approaches can be difficult and unstable to optimize.
Method
The method jointly estimates an orthogonal matrix and permutation matrix, using a convex graph-matching relaxation for initialization and a stochastic algorithm for large-scale optimization.
Results
The method performs on par with state-of-the-art unsupervised approaches for word translation, with results comparable to ICP and significantly better than adversarial training.
Takeaways & Limitations
The approach provides a general, scalable way to align high-dimensional embeddings and achieves state-of-the-art-level unsupervised word-translation performance.
Takeaways & Limitations
The problem is non-convex, and the relationship between graph matching and point-cloud alignment, including where the approach succeeds or fails, remains incompletely understood.
Abstract
from arXiv · showhide
We consider the task of aligning two sets of points in high dimension, which has many applications in natural language processing and computer vision. As an example, it was recently shown that it is possible to infer a bilingual lexicon, without supervised data, by aligning word embeddings trained on monolingual data. These recent advances are based on adversarial training to learn the mapping between the two embeddings. In this paper, we propose to use an alternative formulation, based on the joint estimation of an orthogonal matrix and a permutation matrix. While this problem is not convex, we propose to initialize our optimization algorithm by using a convex relaxation, traditionally considered for the graph isomorphism problem. We propose a stochastic algorithm to minimize our cost function on large scale problems. Finally, we evaluate our method on the problem of unsupervised word translation, by aligning word embeddings trained on monolingual data. On this task, our method obtains state of the art results, while requiring less computational resources than competing approaches.
1 Introduction
Aligning high-dimensional embedding clouds matters across language and vision, but existing unsupervised methods can be difficult and unstable to optimize. The paper proposes jointly estimating a transformation and word correspondence, using convex initialization and scalable stochastic optimization.
- High-dimensional embedding alignment supports unsupervised translation, point-set registration, and structure-from-motion.
- Unsupervised alignment methods based on Wasserstein distance or adversarial training can require sophisticated, hard, and sometimes unstable optimization.
- The proposed approach jointly learns a transformation and correspondence between two point clouds.
- A convex relaxation initializes a stochastic algorithm that scales to large datasets and converges in a few minutes.
- The method is evaluated on toy experiments and bilingual lexicon induction, where it is competitive with state-of-the-art unsupervised approaches.
2 Approach
The approach aligns two point clouds by jointly estimating an orthogonal transformation and unknown one-to-one correspondences, combining Procrustes analysis with Wasserstein distance. It uses convex relaxation for initialization and stochastic mini-batch optimization to scale the nonconvex problem.
- Procrustes: Procrustes analysis estimates an orthogonal transformation between matched point sets, preserving distances and admitting a closed-form SVD solution.The solution is Q*=UV^T when X^T Y=USV^T.
- Wasserstein distance: Wasserstein distance formulates unknown correspondences as an optimal assignment problem over permutation matrices enforcing one-to-one mappings.The exact assignment can be solved by the Hungarian algorithm, while Sinkhorn provides an approximate Earth Mover Distance solver.
- Joint formulation: The proposed formulation jointly learns an orthogonal matrix Q and a permutation matrix P because neither the correspondence nor the transformation is known.The resulting objective is not jointly convex, although each variable has an exact solution when the other is fixed.
- Stochastic optimization: Stochastic optimization samples mini-batches, solves their optimal matching, computes a gradient for Q, and projects the update back onto the orthogonal matrices.This avoids full matching costs that scale as O(n^3) exactly or O(n^2) up to logarithmic terms approximately.
- Convex relaxation: The method relaxes permutation matrices to the Birkhoff polytope, a convex-hull relaxation associated with graph matching, to initialize optimization.The underlying graph-matching formulation is NP-hard in general, motivating the convex relaxation.
- Similarity measure: CSLS is used as a similarity measure between vectors from two sets, with nearest-neighbor averaging rather than a temperature parameter controlling its adjustment.Its free parameter, the number of nearest neighbors, is described as simpler to set than ISF's temperature.
3 Related work
Prior work addresses embedding alignment through supervised transformations, distribution matching, adversarial training, and graph-matching relaxations. This paper's formulation connects these lines by jointly aligning points and transforming embeddings.
- Unsupervised embedding alignment methods have minimized distribution distances using Wasserstein distance or adversarial training.These approaches avoid relying entirely on bilingual anchor points or noisy supervision.
- Supervised word-vector alignment commonly uses a linear transformation, with orthogonal constraints shown to suit normalized embeddings.Orthogonality preserves distances between points under the transformation.
- The paper formulates embedding alignment similarly to graph matching, whose permutation relaxation replaces permutation matrices with doubly-stochastic matrices.Such relaxations can equal exact isomorphism under certain graph conditions but generally lack convergence guarantees.
- The proposed formulation minimizes Wasserstein distance to estimate a transformation between two point distributions while jointly learning alignment.This connects explicit point matching with distribution-based alignment.
4 Experiments
The experiments evaluate the method on controlled embedding variations and unsupervised word translation. Convex-relaxation initialization improves stochastic optimization, while the method remains competitive with strong unsupervised baselines and trades batch size against speed and approximation quality.
- 4.1 Toy experiments: Toy datasets vary embedding training through random seeds, data splits, context-window sizes, and model algorithms.The embeddings are trained on English News Crawl data, enabling ground-truth matching between corresponding vectors.
- 4.1 Toy experiments: Convex-relaxation initialization consistently reaches good solutions, whereas random initialization succeeds less often as the toy problems become harder.Even a relaxed solution that is not itself good can provide a useful initialization for the stochastic gradient method.
- 4.2 Unsupervised word translation: The word-translation evaluation uses the same datasets, code, and word vectors as Conneau et al. (2017), comparing Procrustes, adversarial training, and ICP.The task infers a bilingual dictionary by aligning word vectors trained on two monolingual corpora.
- 4.2 Unsupervised word translation: After refinement, the method performs on par with ICP and significantly better than adversarial training.The authors attribute comparable ICP performance to the use of a similar loss function, while emphasizing the difference in initialization requirements.
- 4.2 Unsupervised word translation: Without refinement, the method outperforms the GAN-based approach on 6 out of 8 language pairs for both NN and CSLS, but remains worse than supervised alignment.The authors also report similar performance for ISF and CSLS with their method and Procrustes, while noting ISF's sensitivity to its temperature parameter.
- 4.2 Unsupervised word translation: Larger batches improve the approximation of squared Wasserstein distance and generally improve performance, while smaller batches converge faster through O(b2) Sinkhorn computations.Table 4 reports precision at 1 after 4,000 iterations using nearest-neighbor retrieval.
5 Conclusion
The paper presents an efficient stochastic approach for high-dimensional embedding alignment, initialized by a convex relaxation, and validates it on toy examples and unsupervised word translation.
- The approach combines a convex relaxation for initialization with an efficient stochastic algorithm for the non-convex, computationally expensive alignment problem.
- On unsupervised word translation, the method achieves performance on par with the state of the art.
- Toy examples and unsupervised word translation provide validation settings for the proposed method.
- The paper identifies open questions about when the approach is guaranteed to work or provably fail, and about improving the relaxation procedure.