Update modeling_quiet.py

modeling_quiet.py

@@ -1262,6 +1262,7 @@ class QuietForCausalLM(QuietPreTrainedModel):
     
         # 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:
@@ -1283,16 +1284,6 @@ class QuietForCausalLM(QuietPreTrainedModel):
                     nn.Linear(talk_input_dim, talk_output_dim, bias=False)
                 )])
 
-        # Add batch normalization to the model
-        self.bn_lm_head = nn.BatchNorm1d(config.vocab_size)
-        self.bn_talk_head = nn.BatchNorm1d(talk_output_dim)
-
-        # Initialize weights using Xavier initialization
-        self.apply(self._init_weights)
-
-        # Add dropout regularization
-        self.dropout = nn.Dropout(config.hidden_dropout_prob)
-
         # Initialize weights and apply final processing
         self.post_init()
 
@@ -1313,14 +1304,6 @@ class QuietForCausalLM(QuietPreTrainedModel):
 
     def get_decoder(self):
         return self.model
-    
-    def _init_weights(self, module):
-        if isinstance(module, nn.Linear):
-            nn.init.xavier_uniform_(module.weight)
-            if module.bias is not None:
-                nn.init.constant_(module.bias, 0)
-        elif isinstance(module, nn.Embedding):
-            nn.init.xavier_uniform_(module.weight)
 
     @torch.no_grad()
     def infer(
@@ -1719,11 +1702,10 @@ class QuietForCausalLM(QuietPreTrainedModel):
             hidden_states = outputs[0]
             prev_rm_logits = rm_logits  # for policy gradient
             prev_rm_tokens = cur_rm_tokens  # for policy gradient
+
             if ahead_idx == 0:
                 hidden_states_lm = hidden_states
                 logits = self.lm_head(hidden_states_lm)
-                logits = self.bn_lm_head(logits.transpose(1, 2)).transpose(1, 2)
-                logits = self.dropout(logits)
                 base_hidden_states = hidden_states.clone()
                 initial_loss_logits = logits.clone()
                 if self.optimize_lm_head_only_at_start or self.optimize_model_only_at_start:
@@ -1754,12 +1736,8 @@ class QuietForCausalLM(QuietPreTrainedModel):
                         head_input_hidden_states = talk_hidden_states
 
                     residual_logits = self.talk_head[0](head_input_hidden_states)
-                    residual_logits = self.bn_talk_head(residual_logits.transpose(1, 2)).transpose(1, 2)
-                    residual_logits = self.dropout(residual_logits)
                     if self.use_shallow_talk:
-                        residual_logits = apply_head(self.lm_head, residual_logits, detach=self.optimize_lm_head_only_at_start) 
-                        residual_logits = self.bn_lm_head(residual_logits.transpose(1, 2)).transpose(1, 2)   
-                        residual_logits = self.dropout(residual_logits)                    
+                        residual_logits = apply_head(self.lm_head, residual_logits, detach=self.optimize_lm_head_only_at_start)                        
                     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