reducing word vectors

ReduceVectors.java

package io.anserini.embeddings.nn;

import org.deeplearning4j.models.embeddings.loader.WordVectorSerializer;
import org.deeplearning4j.models.embeddings.wordvectors.WordVectors;
import org.deeplearning4j.models.word2vec.wordstore.VocabCache;
import org.nd4j.linalg.api.ndarray.INDArray;
import org.nd4j.linalg.dimensionalityreduction.PCA;

import java.io.IOException;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.concurrent.atomic.AtomicInteger;

public class ReduceVectors {

    public static void main(String[] args) throws IOException {
        Path model = Paths.get(args[0]);
        WordVectors wordVectors = WordVectorSerializer.readWord2VecModel(model.toFile());

        int dim = 8;
        // see https://arxiv.org/abs/1702.01417
        INDArray x = postProcess(wordVectors.lookupTable().getWeights(), dim);

        // see https://arxiv.org/abs/1708.03629#
        INDArray pcaX = PCA.pca(x, dim, true);
        INDArray reduced = postProcess(pcaX, dim);

        wordVectors.lookupTable().resetWeights();
        VocabCache vocab = wordVectors.vocab();
        AtomicInteger inc = new AtomicInteger();
        vocab.words().forEach(obj -> {
            String word = (String) obj;
            wordVectors.lookupTable().putVector(word, reduced.getRow(inc.get()));
            inc.getAndIncrement();
        });
        WordVectorSerializer.writeWordVectors(wordVectors.lookupTable(), "reduced-" + model.getFileName());

    }

    private static INDArray postProcess(INDArray weights, int d) {
        INDArray meanWeights = weights.sub(weights.meanNumber());
        INDArray pca = PCA.pca(meanWeights, d, true);
        for (int j = 0; j < weights.rows(); j++) {
            INDArray v = meanWeights.getRow(j);
            for (int s = 0; s < d; s++) {
                INDArray u = pca.getColumn(s);
                INDArray mul = u.mmul(v).transpose().mmul(u);
                v.subi(mul.transpose());
            }
        }
        return weights;
    }
}