cleanup

Author: lucidrains

HRM/hrm.py

@@ -79,7 +79,8 @@ def __init__(
         min_reasoning_steps_epsilon_prob = 0.5,            # they stochastically choose the minimum segment from 2 .. max with this probability, and 1 step the rest of the time
         max_reasoning_steps = 10,
         act_loss_weight = 1.,
-        ignore_index = -1
+        discount_factor = 1.,
+        ignore_index = -1,
     ):
         super().__init__()
 
@@ -134,6 +135,8 @@ def __init__(
 
         # Q(continue|halt) for their adaptive computation time setup
 
+        self.discount_factor = discount_factor
+
         self.act_loss_weight = act_loss_weight
 
         self.min_reasoning_steps_epsilon_prob = min_reasoning_steps_epsilon_prob
@@ -255,15 +258,26 @@ def evaluate_pred():
                     should_halt = q_halt > q_continue
 
                     if return_loss:
+
+                        # Q_halt
+
                         with torch.no_grad():
                             is_correct = (evaluate_pred().argmax(dim = -1) == labels).all(dim = -1)
 
-                        halt_target_loss = F.binary_cross_entropy(q_halt, is_correct.float())
+                        halt_target_loss = F.binary_cross_entropy(
+                            q_halt,
+                            is_correct.float()
+                        )
 
                         act_losses.append(halt_target_loss)
 
+                        # Q_continue
+
                         if exists(prev_q_continue):
-                            continue_target_loss = F.binary_cross_entropy(prev_q_continue, torch.maximum(q_continue, q_halt))
+                            continue_target_loss = F.binary_cross_entropy(
+                                prev_q_continue,
+                                torch.maximum(q_continue, q_halt) * self.discount_factor
+                            )
 
                             act_losses.append(continue_target_loss)
 

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "HRM-pytorch"
-version = "0.0.8"
+version = "0.0.9"
 description = "The proposal from a Singaporean AGI company"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }