[patch] fix bug in solver constraint handling (#16)

* small bug fixes

* lambo tweaks

* revert coord selection changes

* clean up commented code

* fix normalization term

Author: samuelstanton

cortex/acquisition/_graph_nei.py

@@ -126,11 +126,15 @@ def tree_output_to_dict(
             )
 
     if constraints is not None:
-        for constraint in constraints:
-            constraint_values = tree_output.fetch_task_outputs(constraint)["logits"]
-            constraint_values = constraint_values.softmax(dim=-1)[..., 1]
+        for c_list in constraints.values():
+            for constraint in c_list:
+                if constraint in result:
+                    continue
 
-            result[constraint] = constraint_values
+                constraint_values = tree_output.fetch_task_outputs(constraint)["logits"]
+                constraint_values = constraint_values.softmax(dim=-1)[..., 1]
+
+                result[constraint] = constraint_values
 
     return result
 
@@ -163,6 +167,7 @@ def get_graph_nei_runtime_kwargs(
         "f_ref": f_ref,
         "f_baseline": f_baseline,
     }
+    print(f"[INFO][LaMBO-2] Baseline value: {f_baseline.mean(0).max().item():.4f}")
     return res
 
 

cortex/model/leaf/_classifier_leaf.py

@@ -22,7 +22,7 @@ def check_probs(probs: torch.Tensor, dim: int = -1) -> bool:
     if torch.any(probs < 0) or torch.any(probs > 1):
         raise ValueError("Probabilities must be between 0 and 1")
 
-    if not torch.allclose(probs.sum(dim=dim), torch.ones(probs.shape[:-1])):
+    if not torch.allclose(probs.sum(dim=dim), torch.ones(probs.shape[:-1], device=probs.device)):
         raise ValueError("Probabilities must sum to 1")
 
     return True

cortex/optim/generative/_lambo.py

@@ -1,4 +1,5 @@
 import math
+import pprint
 import warnings
 from typing import Callable, Optional
 
@@ -139,7 +140,7 @@ def step(self) -> None:
         activations, trunk_outputs = self._get_latent_variables(generation_inputs)
 
         delta = torch.nn.Parameter(torch.zeros_like(activations))
-        optimizer = torch.optim.Adam([delta], lr=self.guidance_step_size)
+        optimizer = torch.optim.Adam([delta], lr=self.guidance_step_size, betas=(0.09, 0.0999))
         metrics = {"step": self._step_count}
 
         # get initial solution before guidance
@@ -207,13 +208,15 @@ def step(self) -> None:
             grad_norm = feature_grad.norm(dim=(-2, -1), keepdim=True)
             metrics.update(
                 {
+                    "act_obj_val": tgt_obj_vals.mean().item(),
                     "masked_design_loss": design_loss.item(),
                     "masked_design_loss_grad_norm": grad_norm.mean().item(),
                     "masked_token_loss": kl_div.item(),
                     "masked_obj_loss": obj_loss.item(),
                     "token_entropy": entropy.item(),
                 }
             )
+            pprint.pp(metrics)
 
         self._step_count += 1
 
@@ -257,7 +260,13 @@ def coord_score(tok_embeddings):
                 null_embedding,
                 is_excluded=~pos_is_feasible,
             )
-            position_probs = (position_scores * self.feature_attr_temp).softmax(-1)
+            denom = torch.where(position_scores > float("-inf"), position_scores, 0.0).abs().sum(-1, keepdim=True)
+            position_scores = position_scores / (denom + 1e-6)
+
+            position_probs = (position_scores / self.feature_attr_temp).softmax(-1)
+            hand_tuned_entropy = torch.distributions.Categorical(probs=position_probs).entropy().median()
+            print(f"[INFO][LaMBO-2]: Hand-tuned entropy = {hand_tuned_entropy}")
+
             edit_idxs = torch.multinomial(position_probs, self.num_mutations_per_step, replacement=False)
             edit_idxs = edit_idxs.sort(dim=-1).values