Add Monte Carlo conditioning search after gradient ascent

search() generates 16 random unit-direction perturbations (2.5% of
conditioning norm each), scores each with the VF, returns the best.
Runs after ascend() in _v2_ascend_conditioning — ascent exploits the
known gradient direction, search explores the neighbourhood around it
for any better nearby region. Total cost: 16 VF forward passes ≈ free.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
digital-garbage
2026-05-30 15:47:28 +03:00
co-authored by Claude Sonnet 4.6
parent 1ec69f68ba
commit 4ca896c73a
2 changed files with 28 additions and 2 deletions
+22
View File
@@ -117,6 +117,28 @@ class OnlineValueFunction(nn.Module):
prev_reward = reward_val
return c.detach().to(orig_device, orig_dtype)
def search(self, conditioning, n=16):
"""Monte Carlo conditioning search: score N random perturbations, return the best."""
if not self.is_ready():
return conditioning
mlp_device = next(self.parameters()).device
orig_device = conditioning.device
orig_dtype = conditioning.dtype
eps = conditioning.float().norm().item() * 0.025
with torch.inference_mode(False), torch.no_grad():
c_base = torch.empty(conditioning.shape, dtype=torch.float32, device=mlp_device)
c_base.copy_(conditioning)
best = c_base
best_score = float(self.forward(self.compress(c_base).unsqueeze(0)).item())
for _ in range(n):
noise = F.normalize(torch.randn_like(c_base).flatten(), dim=0).reshape(c_base.shape)
candidate = c_base + eps * noise
score = float(self.forward(self.compress(candidate).unsqueeze(0)).item())
if score > best_score:
best_score = score
best = candidate
return best.detach().to(orig_device, orig_dtype)
def is_ready(self):
return len(self.buffer_c) >= self.MIN_SAMPLES