Feedback on manually implemented hierarchical softmax

ruchitpatel: <p>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.</p>
<pre><code class="lang-auto">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 &lt; other.freq
    
    def __gt__(self, other):
        return self.freq &gt; 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)&gt;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)
</code></pre>
<p>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.</p>
<pre><code class="lang-auto">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()
</code></pre>
<p>Now I define the calculation of loss function for hierarchical softmax as follows :</p>
<pre><code class="lang-auto">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
</code></pre>
<p>And I define my word2vec model class as follows :</p>
<pre><code class="lang-auto">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()
</code></pre>
<p>Here my two questions are :<br>
(1) As we can see, I’ve added <code>model.parameters()</code> 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.</p>
<p>(2) My second question is about the correctness of my approach. Is my approach to train a Word2Vec model like this a correct one?</p>
<p>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?</p>

Amir_Raza: <p>I searched the internet for some references, but I only found this as an example pytorch implemention. Maybe you can cross check your code with this. <a href="https://github.com/leimao/Two_Layer_Hierarchical_Softmax_PyTorch/blob/master/utils.py#L98" rel="nofollow noopener">https://github.com/leimao/Two_Layer_Hierarchical_Softmax_PyTorch/blob/master/utils.py#L98</a></p>