I manually implemented the hierarchical softmax, since I did not find its implementation. I implemented my model as follows. The model is simple word2vec model, but instead of using negative sampling, I want to use hierarchical softmax. In hierarchical softmax, there is no output word representations like the ones used in vanila softmax, or negative sampling. In hierarchical softmax, we have a target word, and a binary tree with each leaf node corresponding to words in the vocabulary. I have one question and wanted feedback on the correctness of my approach. Thank you in advance. Following snippet correspond to implementation of binary tree for hierarchical softmax. I use huffman tree as a binary tree for softmax.
import numpy as np
import torch.nn as nn
import torch
import heapq
class Node:
def __init__(self, token, freq):
self.vec = torch.randn(300, requires_grad=True, dtype=torch.float)
self.token = token
self.freq = freq
self.left = None
self.right = None
def __lt__(self, other):
return self.freq < other.freq
def __gt__(self, other):
return self.freq > other.freq
def __eq__(self, other):
if(other == None):
return False
if(not isinstance(other, Node)):
return False
return self.freq == other.freq
class HuffmanTree:
def __init__(self):
self.heap = []
self.codes = {}
self.reverse_mapping = {}
self.root = None
def make_heap(self, frequency):
for key in frequency:
node = Node(key, frequency[key])
heapq.heappush(self.heap, node)
def merge_nodes(self):
while(len(self.heap)>1):
node1 = heapq.heappop(self.heap)
node2 = heapq.heappop(self.heap)
merged = Node(None, node1.freq + node2.freq)
merged.left = node1
merged.right = node2
heapq.heappush(self.heap, merged)
def make_codes_helper(self, root, current_code):
if(root==None):
return
if(root.token != None):
self.codes[root.token] = current_code
self.reverse_mapping[current_code] = root.token
return
self.make_codes_helper(root.left, current_code + "0")
self.make_codes_helper(root.right, current_code + "1")
def make_codes(self):
root = heapq.heappop(self.heap)
self.root = root
current_code = ""
self.make_codes_helper(root, current_code)
Above snippet is to construct the binary hauffman tree for softmax. Then I initialise the leaves of the binary tree as words in my vocabulary as follows.
import pickle
d = pickle.load(open("token_freqs.p", "rb"))
# d1 is a dictionary with tokens as keys and their frequencies as values
d1 = {k:v for k,v in sorted(d.items(), key=lambda x: x[1], reverse=True)}
# ht corresponds to one instance of a huffman tree for softmax
ht = HuffmanTree()
ht.make_heap(d1)
ht.merge_nodes()
ht.make_codes()
Now I define the calculation of loss function for hierarchical softmax as follows :
def cal_loss(h, target):
path_to_word = ht.codes[target]
loss = torch.zeros(1, requires_grad=True, dtype=torch.float)
root=ht.root
for i in path_to_word:
if(i=='0'):
loss = loss + torch.log(torch.sigmoid(torch.dot(root.vec, h)))
root = root.left
else:
loss = loss + torch.log(torch.sigmoid(-1*torch.dot(root.vec, h)))
root = root.right
loss = loss*-1
return loss
And I define my word2vec model class as follows :
vocab_size=30000
class Word2Vec(nn.Module):
def __init__(self):
super().__init__()
self.embedding_layer = nn.Linear(vocab_size, 300, bias=False)
def forward(self, inp):
hidd = self.embedding_layer(inp)
return hidd
model = Word2Vec()
optimizer = torch.optim.SGD(model.parameters(), lr=1e-4, momentum=0.9)
for word, target in tranning_set:
hidd = model(word)
loss = cal_loss(hidd, target)
loss.backward()
optimizer.step()
optimizer.zero_grad()
Here my two questions are :
(1) As we can see, I’ve added model.parameters() in the optimizer I defined. I doubt that this just includes the parameters of the model defined in Word2Vec class. How do I ensure that the parameters corresponding to the binary tree built for softmax approximation also get updated. I know that we can easily get the model parameters using model.parameters(), but for the binary tree class I implemented, we can not do anything like that. How do I ensure that my model updates all the parameters in the huffman tree as well during the training.
(2) My second question is about the correctness of my approach. Is my approach to train a Word2Vec model like this a correct one?
Also since we do not have any implementation of hierarchical softmax in the pytorch library, I’d like to contribute a tested version of this technique. Can anyone guide me regarding what should I do next to make this contribution?