Hey Thomas @tom  ,
great to hear from you, and awesome advice as usual 
Following your guide, I think we can use a word embedding together with the Wasserstein divergence as an criterion?
Here’s a rough guide to how it might work?
import torch
from random import randint

num_words = 10
embedding_dim = 5

# have an embedding  
# can initialize to -1 or +1
# or copy pretrained weights, see- https://discuss.pytorch.org/t/can-we-use-pre-trained-word-embeddings-for-weight-initialization-in-nn-embedding/1222
t = torch.nn.Embedding(num_words,embedding_dim)

# map from embedding to probability space
vec_to_prob = nn.Softmax()

# sample a random word to train on 
word_idx = randint(0,num_words-1)

# a batch of 1 sample of 1 index/word
word_idx = Variable(torch.LongTensor([[word_idx]]))

# vector representation
word_vec = t(word_idx)
word_vec = word_vec.squeeze(0) # drop the batch dimension

# sanity check !!! 
_ , closest_word_idx = torch.topk( torch.mv( t.weight , word_vec.squeeze(0) ) , 1 )
closest_word_idx == word_idx #true

# map to probability space, 
# could be used to calculate the Wasserstein divergence as the training objective, with a histogram from a decoder 
histogram_target = vec_to_prob(word_vec)
histogram_model = blah, blah, blah
wasserstein_loss = divergence( histogram_target , histogram_model )

# after training, histogram from decoder should be "close" to target histogram in Wasserstein space
histogram_model = histogram_target

_ , closest_word_idx = torch.topk( torch.mv( t.weight , histogram_model.squeeze(0) ) , 1 )
closest_word_idx == word_idx #true 

It seems reasonable, and simple, so I’d like to try it   Would really appreciate your opinion !
Best regards,
Ajay