Spark-nlp: Sentence Embeddings Approaches

[SoundBot]

Is there a way to extend sparknlp and create my custom embedder similar to BertEmbeddings? There are some interesting models on TF Hub which I would like to try.

[maziyarpanahi]

Currently, the answer is no. The WordEmbeddings annotator has this flexibility which regardless of how the embeddings were produced, as long as it follows the format it can be loaded such as GloVe, FastText, Word2Vec, etc.
However, BertEmbeddings annotator was created based on BERT itself. If you have an embedding model that is in the same format (ex: fine-tuned BERT embeddings) then you can just use the notebook and convert it to Spark NLP BertEmbeddingsModel.

That's said if you can provide some examples we can take a look at them to see if it's worth implementing an annotator for those models in TF Hub (considering accuracy, performance, etc.)

[SoundBot]

@maziyarpanahi Thanks for the prompt reply! Currently I'm interested in sentence-level embeddings, which can be generated using Universal Sentence Encoder or BERT.
If I read the docs correctly, sparknlp generates word-level embeddings, which can be fed to SentenceEmbeddings annotator to average/sum word-level results. I believe this approach looses a lot of contextual information from a sentence.

[maziyarpanahi]

Actually, one of the authors of BERT embeddings suggested averaging the vectors:

It should be noted that although the* “[CLS]”* acts as an “aggregate representation” for classification tasks, this is not the best choice for a high quality sentence embedding vector. According to BERT author Jacob Devlin (https://github.com/google-research/bert/issues/164)

Original comment:

I'm not sure what these vectors are, since BERT does not generate meaningful sentence vectors. It seems that this is is doing average pooling over the word tokens to get a sentence vector, but we never suggested that this will generate meaningful sentence representations. And even if they are decent representations when fed into a DNN trained for a downstream task, it doesn't mean that they will be meaningful in terms of cosine distance. (Since cosine distance is a linear space where all dimensions are weighted equally).

The USE is a whole another approach and I do agree simply averaging may not be the best way especially with contextualized embeddings. I am working on introducing other pooling strategies for BERT to average the last 4 layers instead of just having 1 layer at the time, and also extend the SentenceEmbeddings to do more such as weighted-average, including TF-IDF as a weight factor, and SIF (Smooth Inverse Frequency).

PS: I would really like to continue this discussion here for the further developments of the Spark NLP's embeddings toolkit as I am using those myself for the similarity engine I am working on. I believe this conversation would help developing better approaches towards sentence and document embeddings.

[SoundBot]

Thanks, that sounds useful.

The reason why I believe we need generic "plug your TF/PyTorch model" annotator is because there is no "one size fits all" model, since there are different SOTA models for text summarization, text similarity, question answering and other tasks.

Embedding models are also very resource hungry. I was getting OOMs on nodes with 400GB RAM using fairly small documents (< 2M characters), so efficient memory management is also something to think about.

[maziyarpanahi]

I agree. If I can manage TF Hub I will make it a bit more generic to have flexibility in which to use.

The new version we released addresses the memory issue for BertEmbeddings. There was a bad memory leak not visible to dataset less than 1 million sentences. I used the new release myself to test on 18 million sentences. (With 100K it manages to stay right on 8G we gave to Spark session locally)

PS: give the new release a try please and let me know how it goes.