Потеря грамматики word2vec не уменьшается

Я работаю над реализацией архитектуры word2vec с нуля. Но моя модель не сходится.

class SkipGramBatcher:
  def __init__(self, text):
    self.text = text.results

  def get_batches(self, batch_size):
    n_batches = len(self.text)//batch_size
    pairs = []


    for idx in range(0, len(self.text)):
      window_size = 5
      idx_neighbors = self._get_neighbors(self.text, idx, window_size)
      #one_hot_idx = self._to_one_hot(idx)
      #idx_pairs = [(one_hot_idx, self._to_one_hot(idx_neighbor)) for idx_neighbor in idx_neighbors]
      idx_pairs = [(idx,idx_neighbor) for idx_neighbor in idx_neighbors]
      pairs.extend(idx_pairs)


    for idx in range(0, len(pairs), batch_size):
      X = [pair[0] for pair in pairs[idx:idx+batch_size]]
      Y = [pair[1] for pair in pairs[idx:idx+batch_size]]
      yield X,Y

  def _get_neighbors(self, text, idx, window_size):
    text_length = len(text)
    start = max(idx-window_size,0)
    end = min(idx+window_size+1,text_length)
    neighbors_words = set(text[start:end])

    return list(neighbors_words)

  def _to_one_hot(self, indexes):
    n_values = np.max(indexes) + 1
    return np.eye(n_values)[indexes]

Я использую корпус text8 и применяю методы предварительной обработки, такие как стемминг, лемматизация и субдискретизация. Также я исключил английские стоп-слова и ограниченный словарный запас.

vocab_size = 20000
text_len = len(text)
test_text_len = int(text_len*0.15)
preprocessed_text = PreprocessedText(text,vocab_size)

Я использую тензорный поток для вычисления графа

train_graph = tf.Graph()
with train_graph.as_default():
  inputs = tf.placeholder(tf.int32, [None], name='inputs')
  labels = tf.placeholder(tf.int32, [None, None], name='labels')

n_embedding =  300
with train_graph.as_default():
  embedding = tf.Variable(tf.random_uniform((vocab_size, n_embedding), -1, 1))
  embed = tf.nn.embedding_lookup(embedding, inputs)

И применить отрицательную выборку

# Number of negative labels to sample
n_sampled = 100
with train_graph.as_default():
  softmax_w = tf.Variable(tf.truncated_normal((vocab_size, n_embedding))) # create softmax weight matrix here
  softmax_b = tf.Variable(tf.zeros(vocab_size), name = "softmax_bias") # create softmax biases here

  # Calculate the loss using negative sampling

  loss = tf.nn.sampled_softmax_loss(
      weights=softmax_w,
      biases=softmax_b,
      labels=labels,
      inputs=embed,
      num_sampled=n_sampled,
      num_classes=vocab_size)

  cost = tf.reduce_mean(loss)
  optimizer = tf.train.AdamOptimizer().minimize(cost)

Наконец я тренирую свою модель

epochs = 10
batch_size = 64
avg_loss = []

with train_graph.as_default():
    saver = tf.train.Saver()

with tf.Session(graph=train_graph) as sess:
  iteration = 1
  loss = 0
  sess.run(tf.global_variables_initializer())

  for e in range(1, epochs+1):
    batches = skip_gram_batcher.get_batches(batch_size)
    start = time.time()
    for  batch_x,batch_y  in batches:

      feed = {inputs: batch_x, 
             labels: np.array(batch_y)[:, None]}
      train_loss, _ = sess.run([cost, optimizer], feed_dict=feed)
      loss += train_loss

      if iteration % 100 == 0: 
        end = time.time()
        print("Epoch {}/{}".format(e, epochs),
                "Iteration: {}".format(iteration),
                "Avg. Batch loss: {:.4f}".format(loss/iteration),
                "{:.4f} sec/batch".format((end-start)/100))
        #loss = 0
        avg_loss.append(loss/iteration)
        start = time.time()

      iteration += 1

  save_path = saver.save(sess, "checkpoints/text8.ckpt")

Но после запуска этой модели мои средние потери партий резко не уменьшаются.

Думаю, я должен был где-то ошибиться. Любая помощь приветствуется