feat: implement hardware-adaptive compute bounding and dynamic entropy routing (Eqs. 3-4)

train.py

@@ -158,16 +158,40 @@ class StackedLDMHeads(nn.Module):
         logits = self.head(forecast)
         return logits
 
-class LogitUncertaintyFilter(nn.Module):
+class HardwareAdaptiveRouter(nn.Module):
+    def __init__(self, c_base=1.0, c_bf16=5.0):
+        super().__init__()
+        # c_base: FLOP cost proxy for the LDM block projection
+        # c_bf16: FLOP cost proxy for a single token bfloat16 refinement
+        self.c_base = c_base
+        self.c_bf16 = c_bf16
+
     def compute_entropy(self, logits: torch.Tensor) -> torch.Tensor:
         probs = F.softmax(logits.float(), dim=-1)
         entropy = -torch.sum(probs * torch.log(probs + 1e-9), dim=-1)
         return entropy
 
-    def forward(self, logits: torch.Tensor, threshold: float):
+    def forward(self, logits: torch.Tensor, c_target: float):
         entropy = self.compute_entropy(logits)
-        mask = entropy >= threshold
-        return mask, entropy
+        
+        # Equation (3): C_step = C_base + sum(M_i) * C_BF16 <= c_target
+        # Calculate the maximum number of bfloat16 refinements we can afford
+        max_refinements = max(0, int((c_target - self.c_base) / self.c_bf16))
+        L = entropy.shape[-1]
+        allowed_refinements = min(max_refinements, L)
+        
+        mask = torch.zeros_like(entropy, dtype=torch.bool)
+        dynamic_tau = float('inf')
+        
+        if allowed_refinements > 0:
+            # Equation (4): tau = inf { t | C_step(t) <= C_target }
+            sorted_entropy, indices = torch.sort(entropy, descending=True)
+            dynamic_tau = sorted_entropy[allowed_refinements - 1].item()
+            
+            # Apply the computed hardware-bounded threshold
+            mask[indices[:allowed_refinements]] = True
+            
+        return mask, entropy, dynamic_tau
 
 class ActorCriticPruner:
     def __init__(self, lm_head, lambda_reg=0.1):
@@ -216,12 +240,12 @@ class ActorCriticPruner:
 
 
 class ADAPTDIFFPipeline(nn.Module):
-    def __init__(self, base_lm_model, block_size=12, entropy_threshold=1.5):
+    def __init__(self, base_lm_model, block_size=12, target_budget=15.0):
         super().__init__()
         self.base_model = base_lm_model.model
         self.lm_head = base_lm_model.lm_head
         self.block_size = block_size
-        self.entropy_threshold = entropy_threshold
+        self.target_budget = target_budget
 
         self.ldm_heads = StackedLDMHeads(
             hidden_size=base_lm_model.config.hidden_size,
@@ -229,10 +253,13 @@ class ADAPTDIFFPipeline(nn.Module):
             block_size=block_size
         ).to(DEVICE)
         
-        self.router = LogitUncertaintyFilter()
+        self.router = HardwareAdaptiveRouter(c_base=1.0, c_bf16=5.0)
         self.pruner = ActorCriticPruner(self.lm_head)
 
-    def generate_adapt_diff(self, input_ids, max_new_tokens=128):
+    def generate_adapt_diff(self, input_ids, max_new_tokens=128, c_target=None):
+        if c_target is None:
+            c_target = self.target_budget
+            
         current_seq = input_ids.clone()
         generated_count = 0
         total_full_transformer_evals = 0
@@ -245,7 +272,7 @@ class ADAPTDIFFPipeline(nn.Module):
             block_logits = self.ldm_heads(last_hidden).squeeze(0).squeeze(0)
             draft_tokens = torch.argmax(block_logits, dim=-1)
 
-            mask, entropy = self.router(block_logits, self.entropy_threshold)
+            mask, entropy, dynamic_tau = self.router(block_logits, c_target)
 
             if not mask.any():
                 final_block = draft_tokens
@@ -277,7 +304,7 @@ a2d_model = AutoModelForCausalLM.from_pretrained(
     device_map=DEVICE
 )
 
-pipeline = ADAPTDIFFPipeline(a2d_model, block_size=12, entropy_threshold=1.5)
+pipeline = ADAPTDIFFPipeline(a2d_model, block_size=12, target_budget=15.0)
 print("Downloading LDM head projection weights for calibration baseline...")
 ldm_weights_path = hf_hub_download(repo_id=ADAPT_DIFF_ID, filename="ldm_heads.pt")
 pipeline.ldm_heads.load_state_dict(torch.load(ldm_weights_path, map_location=DEVICE))