Add unsupervised multilingual alignement

Summary: Add a script for unsupervised multilingual alignment.

Reviewed By: EdouardGrave

Differential Revision: D17180273

fbshipit-source-id: edbb139ff9474ef325a43bb16e9c0cf1a76e0900

alignment/README.md

@@ -16,9 +16,10 @@ The details of this approach can be found in [1].
 ### Unsupervised alignment
 
 The script `unsup_align.py` aligns word embeddings from two languages without requiring any supervision.
-The details of this approach can be found in [2].
+Additionally, the script `unsup_multialign.py` aligns multiple languages to a common space with no supervision.
+The details of these approaches can be found in [2] and [3] respectively.
 
-In addition to NumPy, the unsupervised method requires the [Python Optimal Transport](https://pot.readthedocs.io/en/stable/) toolbox.
+In addition to NumPy, the unsupervised methods require the [Python Optimal Transport](https://pot.readthedocs.io/en/stable/) toolbox.
 
 ### Download
 
@@ -39,7 +40,7 @@ If you use the supervised alignment method, please cite:
 }
 ```
 
-If you use the unsupervised alignment method, please cite:
+If you use the unsupervised bilingual alignment method, please cite:
 
 [2] E. Grave, A. Joulin, Q. Berthet, [*Unsupervised Alignment of Embeddings with Wasserstein Procrustes*](https://arxiv.org/abs/1805.11222)
 
@@ -51,3 +52,16 @@ If you use the unsupervised alignment method, please cite:
     year={2018}
 }
 ```
+
+If you use the unsupervised alignment script `unsup_multialign.py`, please cite:
+
+[3] J. Alaux, E. Grave, M. Cuturi, A. Joulin, [*Unsupervised Hyperalignment for Multilingual Word Embeddings*](https://arxiv.org/abs/1811.01124)
+
+```
+@article{alaux2018unsupervised,
+  title={Unsupervised hyperalignment for multilingual word embeddings},
+  author={Alaux, Jean and Grave, Edouard and Cuturi, Marco and Joulin, Armand},
+  journal={arXiv preprint arXiv:1811.01124},
+  year={2018}
+}
+```

alignment/unsup_multialign.py

@@ -0,0 +1,198 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+#
+# Copyright (c) 2019-present, Facebook, Inc.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import io, os, ot, argparse, random
+import numpy as np
+from utils import *
+
+parser = argparse.ArgumentParser(description=' ')
+
+parser.add_argument('--embdir', default='data/', type=str)
+parser.add_argument('--outdir', default='output/', type=str)
+parser.add_argument('--lglist', default='en-fr-es-it-pt-de-pl-ru-da-nl-cs',  type=str,
+        help='list of languages. The first element is the pivot. Example: en-fr-es to align English, French and Spanish with English as the pivot.')
+
+parser.add_argument('--maxload', default=20000, type=int, help='Max number of loaded vectors')
+parser.add_argument('--uniform', action='store_true', help='switch to uniform probability of picking language pairs')
+
+# optimization parameters for the square loss
+parser.add_argument('--epoch', default=2, type=int, help='nb of epochs for square loss')
+parser.add_argument('--niter', default=500, type=int, help='max number of iteration per epoch for square loss')
+parser.add_argument('--lr', default=0.1, type=float, help='learning rate for square loss')
+parser.add_argument('--bsz', default=500, type=int, help='batch size for square loss')
+
+# optimization parameters for the RCSLS loss
+parser.add_argument('--altepoch', default=100, type=int, help='nb of epochs for RCSLS loss')
+parser.add_argument('--altlr', default=25, type=float, help='learning rate for RCSLS loss')
+parser.add_argument("--altbsz", type=int, default=1000, help="batch size for RCSLS")
+
+args = parser.parse_args()
+
+###### SPECIFIC FUNCTIONS ######
+
+def getknn(sc, x, y, k=10):
+    sidx = np.argpartition(sc, -k, axis=1)[:, -k:]
+    ytopk = y[sidx.flatten(), :]
+    ytopk = ytopk.reshape(sidx.shape[0], sidx.shape[1], y.shape[1])
+    f = np.sum(sc[np.arange(sc.shape[0])[:, None], sidx])
+    df = np.dot(ytopk.sum(1).T, x)
+    return f / k, df / k
+
+
+def rcsls(Xi, Xj, Zi, Zj, R, knn=10):
+    X_trans = np.dot(Xi, R.T)
+    f = 2 * np.sum(X_trans * Xj)
+    df = 2 * np.dot(Xj.T, Xi)
+    fk0, dfk0 = getknn(np.dot(X_trans, Zj.T), Xi, Zj, knn)
+    fk1, dfk1 = getknn(np.dot(np.dot(Zi, R.T), Xj.T).T, Xj, Zi, knn)
+    f = f - fk0 -fk1
+    df = df - dfk0 - dfk1.T
+    return -f / Xi.shape[0], -df.T / Xi.shape[0]
+
+
+def GWmatrix(emb0):
+    N = np.shape(emb0)[0]
+    N2 = .5* np.linalg.norm(emb0, axis=1).reshape(1, N)
+    C2 = np.tile(N2.transpose(), (1, N)) + np.tile(N2, (N, 1))
+    C2 -= np.dot(emb0,emb0.T)
+    return C2
+
+
+def gromov_wasserstein(x_src, x_tgt, C2):
+    N = x_src.shape[0]
+    C1 = GWmatrix(x_src)
+    M = ot.gromov_wasserstein(C1,C2,np.ones(N),np.ones(N),'square_loss',epsilon=0.55,max_iter=100,tol=1e-4)
+    return procrustes(np.dot(M,x_tgt), x_src)
+
+
+def align(EMB, TRANS, lglist, args):
+    nmax, l = args.maxload, len(lglist)
+    # create a list of language pairs to sample from
+    # (default == higher probability to pick a language pair contianing the pivot)
+    # if --uniform: uniform probability of picking a language pair
+    samples = []
+    for i in range(l):
+        for j in range(l):
+            if j == i :
+                continue
+            if j > 0 and args.uniform == False:
+                samples.append((0,j))
+            if i > 0 and args.uniform == False:
+                samples.append((i,0))
+            samples.append((i,j))
+
+    # optimization of the l2 loss
+    print('start optimizing L2 loss')
+    lr0, bsz, nepoch, niter = args.lr, args.bsz, args.epoch, args.niter
+    for epoch in range(nepoch):
+        print("start epoch %d / %d"%(epoch+1, nepoch))
+        ones = np.ones(bsz)
+        f, fold, nb, lr = 0.0, 0.0, 0.0, lr0
+        for it in range(niter):
+            if it > 1 and f > fold + 1e-3:
+                lr /= 2
+            if lr < .05:
+                break
+            fold = f
+            f, nb = 0.0, 0.0
+            for k in range(100 *  (l-1)):
+                (i,j) = random.choice(samples)
+                embi = EMB[i][np.random.permutation(nmax)[:bsz], :]
+                embj = EMB[j][np.random.permutation(nmax)[:bsz], :]
+                perm = ot.sinkhorn(ones, ones, np.linalg.multi_dot([embi, -TRANS[i], TRANS[j].T,embj.T]), reg = 0.025, stopThr = 1e-3)
+                grad = np.linalg.multi_dot([embi.T, perm, embj])
+                f -= np.trace(np.linalg.multi_dot([TRANS[i].T, grad, TRANS[j]])) / embi.shape[0]
+                nb += 1
+                if i > 0:
+                    TRANS[i] = proj_ortho(TRANS[i] + lr * np.dot(grad, TRANS[j]))
+                if j > 0:
+                    TRANS[j] = proj_ortho(TRANS[j] + lr * np.dot(grad.transpose(), TRANS[i]))
+            print("iter %d / %d - epoch %d - loss: %.5f  lr: %.4f" % (it, niter, epoch+1, f / nb , lr))
+        print("end of epoch %d - loss: %.5f - lr: %.4f" % (epoch+1, f / max(nb,1), lr))
+        niter, bsz = max(int(niter/2),2), min(1000, bsz * 2)
+    #end for epoch in range(nepoch):
+
+    # optimization of the RCSLS loss
+    print('start optimizing RCSLS loss')
+    f, fold, nb, lr = 0.0, 0.0, 0.0, args.altlr
+    for epoch in range(args.altepoch):
+        if epoch > 1  and f-fold > -1e-4 * abs(fold):
+            lr/= 2
+        if lr < 1e-1:
+            break
+        fold = f
+        f, nb = 0.0, 0.0
+        for k in range(round(nmax / args.altbsz) * 10 * (l-1)):
+            (i,j) = random.choice(samples)
+            sgdidx = np.random.choice(nmax, size=args.altbsz, replace=False)
+            embi = EMB[i][sgdidx, :]
+            embj = EMB[j][:nmax, :]
+            # crude alignment approximation:
+            T = np.dot(TRANS[i], TRANS[j].T)
+            scores = np.linalg.multi_dot([embi, T, embj.T])
+            perm = np.zeros_like(scores)
+            perm[np.arange(len(scores)), scores.argmax(1)] = 1
+            embj = np.dot(perm, embj)
+            # normalization over a subset of embeddings for speed up
+            fi, grad = rcsls(embi, embj, embi, embj, T.T)
+            f += fi
+            nb += 1
+            if i > 0:
+                TRANS[i] = proj_ortho(TRANS[i] - lr * np.dot(grad, TRANS[j]))
+            if j > 0:
+                TRANS[j] = proj_ortho(TRANS[j] - lr * np.dot(grad.transpose(), TRANS[i]))
+        print("epoch %d - loss: %.5f - lr: %.4f" % (epoch+1, f / max(nb,1), lr))
+    #end for epoch in range(args.altepoch):
+    return TRANS
+
+def convex_init(X, Y, niter=100, reg=0.05, apply_sqrt=False):
+    n, d = X.shape
+    K_X, K_Y = np.dot(X, X.T), np.dot(Y, Y.T)
+    K_Y *= np.linalg.norm(K_X) / np.linalg.norm(K_Y)
+    K2_X, K2_Y = np.dot(K_X, K_X), np.dot(K_Y, K_Y)
+    P = np.ones([n, n]) / float(n)
+    for it in range(1, niter + 1):
+        G = np.dot(P, K2_X) + np.dot(K2_Y, P) - 2 * np.dot(K_Y, np.dot(P, K_X))
+        q = ot.sinkhorn(np.ones(n), np.ones(n), G, reg, stopThr=1e-3)
+        alpha = 2.0 / float(2.0 + it)
+        P = alpha * q + (1.0 - alpha) * P
+    return procrustes(np.dot(P, X), Y).T
+
+
+###### MAIN ######
+
+lglist = args.lglist.split('-')
+l = len(lglist)
+
+# embs:
+EMB = {}
+for i in range(l):
+    fn = args.embdir + '/wiki.' + lglist[i] + '.vec'
+    _, vecs = load_vectors(fn, maxload=args.maxload)
+    EMB[i] = vecs
+
+#init
+print("Computing initial bilingual apping with Gromov-Wasserstein...")
+TRANS={}
+maxinit = 2000
+emb0 = EMB[0][:maxinit,:]
+C0 = GWmatrix(emb0)
+TRANS[0] = np.eye(300)
+for i in range(1, l):
+    print("init "+lglist[i])
+    embi = EMB[i][:maxinit,:]
+    TRANS[i] = gromov_wasserstein(embi, emb0, C0)
+
+# align
+align(EMB, TRANS, lglist, args)
+
+print('saving matrices in ' + args.outdir)
+languages=''.join(lglist)
+for i in range(l):
+    save_matrix(args.outdir + '/W-' + languages + '-' + lglist[i], TRANS[i])