scale by num_senses and update weights

backpack_model.py

@@ -159,6 +159,10 @@ class BackpackGPT2Model(BackpackGPT2PreTrainedModel):
 
         # Compute resulting outputs
         hidden_states = torch.sum(contextualization @ senses, dim=1) # (bs, nv, s, d) -> (bs, s, d)
+
+        # divide hidden_states by 1 / num_senses
+        hidden_states = hidden_states / self.num_senses
+
         return BackpackGPT2BaseModelOutput(
             hidden_states=hidden_states,
             contextualization=contextualization,
@@ -187,31 +191,24 @@ class BackpackGPT2LMHeadModel(BackpackGPT2PreTrainedModel):
   def __init__(self, config):
     super().__init__(config)
     self.backpack = BackpackGPT2Model(config)
-    self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
 
     # Model parallel
     self.model_parallel = False
     self.device_map = None
 
-    self.tie_weights()
-
-  def tie_weights(self):
-      self.lm_head.weight = self.backpack.word_embeddings.weight # also tied with the underlying underlying transf
-
   def get_lm_head(self):
       return self.lm_head
 
   def forward(self, input_ids, position_ids=None):
       outputs = self.backpack(input_ids, position_ids=position_ids)
       hidden_states, contextualization = outputs.hidden_states, outputs.contextualization
-      lm_logits = self.lm_head(hidden_states) # (bs, s, V)
+      # unembed the hidden_states
+      lm_logits = torch.einsum('bsd,nd->bsn', hidden_states, self.backpack.word_embeddings.weight)    
       return BackpackGPT2LMHeadModelOutput(
             logits=lm_logits,
             contextualization=contextualization,
         )
-    #   CausalLMOutput = namedtuple('CausalLMOutput', ['logits'])
-    #   return CausalLMOutput(logits=lm_logits) 
-
+    
   def run_with_custom_contextualization(self, input_ids, contextualization):
       outputs = self.backpack.run_with_custom_contextualization(input_ids, contextualization)
       hidden_states, contextualization = outputs.hidden_states, outputs.contextualization

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:844eb078f8af73181515736354aedcd84d99b1dd21e1218da5e7d4454df46463
-size 836334888
+oid sha256:912f6a4da875f80a90a238d8714e1a27f79f257693509a9dfd3dd7e2a39165e7
+size 990745984