Update modeling_quiet.py

modeling_quiet.py

@@ -1252,7 +1252,7 @@ class QuietForCausalLM(QuietPreTrainedModel, GenerationMixin):
     
         # For visualization
         self.eval_mode = False
-        
+
         num_talk = 1
         talk_input_dim = config.hidden_size if not self.use_concat_talk_head else config.hidden_size * 2
         if self.use_weighted_talk_head:
@@ -1273,8 +1273,6 @@ class QuietForCausalLM(QuietPreTrainedModel, GenerationMixin):
                 self.talk_head = nn.ModuleList([nn.Sequential(
                     nn.Linear(talk_input_dim, talk_output_dim, bias=False)
                 )])
-            
-            self.mixing_head = nn.Linear(config.hidden_size * 2, 1)
                 
         self.apply(self._init_weights)
 
@@ -1668,10 +1666,12 @@ class QuietForCausalLM(QuietPreTrainedModel, GenerationMixin):
                     residual_logits = self.talk_head[0](head_input_hidden_states)
                     if self.use_shallow_talk:
                         residual_logits = apply_head(self.lm_head, residual_logits, detach=self.optimize_lm_head_only_at_start)                        
-                        residual_logits = residual_logits.to(logits.device)
-                        mixing_weights = self.mixing_head(torch.cat([cur_base_hidden, talk_hidden_states], dim=-1))
-                        mixing_weights = torch.sigmoid(mixing_weights)
-                        logits = base_logits * (1 - mixing_weights) + residual_logits * mixing_weights
+                    residual_logits = residual_logits.to(logits.device)
+                    if self.use_weighted_talk_head:
+                        # combine the cur_base_hidden with the talk_hidden_states according to the weighted head
+                        residual_logits = cur_base_hidden * (1 - residual_logits) + talk_hidden_states * residual_logits
+                        residual_logits = apply_head(self.lm_head, residual_logits, detach=self.optimize_lm_head_only_at_start)
+
                 assert sum([self.cumulative_residual, self.clever_residual, self.skip_residual, self.no_residual]) == 1
                 if self.clever_residual:
                     if ahead_idx >= self.n_ahead - 1: