Implement positive policy decay

src/search/classic/node.h

@@ -39,6 +39,7 @@
 #include "chess/position.h"
 #include "neural/encoder.h"
 #include "proto/net.pb.h"
+#include "utils/fastmath.h"
 #include "utils/mutex.h"
 
 namespace lczero {
@@ -403,10 +404,24 @@ class EdgeAndNode {
     return edge_ ? edge_->GetMove(flip) : Move();
   }
 
-  // Returns U = numerator * p / N.
+  // Returns U = numerator * p_eff / N.
   // Passed numerator is expected to be equal to (cpuct * sqrt(N[parent])).
-  float GetU(float numerator) const {
-    return numerator * GetP() / (1 + GetNStarted());
+  // policy_decay_scale_per_move controls positive policy decay with fixed sqrt decay.
+  // When scale=0, no decay (p_eff = p).
+  // num_edges is the number of legal moves (used to scale the decay parameter).
+  // policy_decay_sum is the sum of raw P_eff for all edges (for sum normalization).
+  float GetU(float numerator, int num_edges,
+             float policy_decay_scale_per_move = 0.0f,
+             float policy_decay_sum = 1.0f) const {
+    float p = GetP();
+    if (policy_decay_scale_per_move > 0.0f && p > 0.0f) {
+      // Calculate raw (unnormalized) P_eff with fixed sqrt decay
+      p = ApplyPolicyDecay(p, static_cast<float>(GetN()),
+                           policy_decay_scale_per_move, num_edges);
+      // Normalize by sum to ensure sum(P_eff) = 1.0
+      p /= policy_decay_sum;
+    }
+    return numerator * p / (1 + GetNStarted());
   }
 
   std::string DebugString() const;

src/search/classic/params.cc

@@ -529,6 +529,12 @@ const OptionId BaseSearchParams::kGarbageCollectionDelayId{
     "garbage-collection-delay", "GarbageCollectionDelay",
     "The percentage of expected move time until garbage collection start. "
     "Delay lets search find transpositions to freed search tree branches."};
+const OptionId BaseSearchParams::kPolicyDecayScaleId{
+    "policy-decay-scale-per-move", "PolicyDecayScalePerMove",
+    "Scale parameter for positive policy decay, specified per legal move. "
+    "Effective scale = scale_per_move * num_legal_moves. Controls how many "
+    "visits per available move before decay becomes significant. Uses fixed "
+    "sqrt decay (exponent=0.5). Range: 1-1000000, default 100."};
 
 const OptionId SearchParams::kMaxPrefetchBatchId{
     "max-prefetch", "MaxPrefetch",
@@ -631,6 +637,7 @@ void BaseSearchParams::Populate(OptionsParser* options) {
   options->Add<FloatOption>(kUCIRatingAdvId, -10000.0f, 10000.0f) = 0.0f;
   options->Add<BoolOption>(kSearchSpinBackoffId) = false;
   options->Add<FloatOption>(kGarbageCollectionDelayId, 0.0f, 100.0f) = 10.0f;
+  options->Add<FloatOption>(kPolicyDecayScaleId, 1.0f, 1000000.0f) = 100.0f;
 }
 
 void SearchParams::Populate(OptionsParser* options) {
@@ -725,7 +732,8 @@ BaseSearchParams::BaseSearchParams(const OptionsDict& options)
       kMaxCollisionVisitsScalingPower(
           options.Get<float>(kMaxCollisionVisitsScalingPowerId)),
       kSearchSpinBackoff(options_.Get<bool>(kSearchSpinBackoffId)),
-      kGarbageCollectionDelay(options_.Get<float>(kGarbageCollectionDelayId)) {}
+      kGarbageCollectionDelay(options_.Get<float>(kGarbageCollectionDelayId)),
+      kPolicyDecayScale(options_.Get<float>(kPolicyDecayScaleId)) {}
 
 SearchParams::SearchParams(const OptionsDict& options)
     : BaseSearchParams(options),

src/search/classic/params.h

@@ -130,6 +130,7 @@ class BaseSearchParams {
     return kMaxOutOfOrderEvalsFactor;
   }
   float GetNpsLimit() const { return kNpsLimit; }
+  float GetPolicyDecayScale() const { return kPolicyDecayScale; }
 
   int GetTaskWorkersPerSearchWorker() const {
     return kTaskWorkersPerSearchWorker;
@@ -231,6 +232,7 @@ class BaseSearchParams {
   static const OptionId kUCIRatingAdvId;
   static const OptionId kSearchSpinBackoffId;
   static const OptionId kGarbageCollectionDelayId;
+  static const OptionId kPolicyDecayScaleId;
 
  protected:
   const OptionsDict& options_;
@@ -290,6 +292,7 @@ class BaseSearchParams {
   const float kMaxCollisionVisitsScalingPower;
   const bool kSearchSpinBackoff;
   const float kGarbageCollectionDelay;
+  const float kPolicyDecayScale;
 };
 
 class SearchParams : public BaseSearchParams {

src/search/classic/search.cc

@@ -465,11 +465,28 @@ std::vector<std::string> Search::GetVerboseStats(const Node* node) const {
   const float cpuct = ComputeCpuct(params_, node->GetN(), is_root);
   const float U_coeff =
       cpuct * std::sqrt(std::max(node->GetChildrenVisits(), 1u));
+  const float policy_decay_scale = params_.GetPolicyDecayScale();
+
+  // Calculate sum of raw P_eff for normalization
+  float policy_decay_sum = 1.0f;
+  if (policy_decay_scale > 0.0f) {
+    policy_decay_sum = 0.0f;
+    for (const auto& edge : node->Edges()) {
+      float p = edge.GetP();
+      if (p > 0.0f) {
+        policy_decay_sum += ApplyPolicyDecay(
+            p, static_cast<float>(edge.GetN()), policy_decay_scale,
+            node->GetNumEdges());
+      }
+    }
+    if (policy_decay_sum == 0.0f) policy_decay_sum = 1.0f;
+  }
+
   std::vector<std::tuple<uint32_t, float, EdgeAndNode>> edges;
   edges.reserve(node->GetNumEdges());
   for (const auto& edge : node->Edges()) {
     edges.emplace_back(edge.GetN(),
-                       edge.GetQ(fpu, draw_score) + edge.GetU(U_coeff),
+                       edge.GetQ(fpu, draw_score) + edge.GetU(U_coeff, node->GetNumEdges(), policy_decay_scale, policy_decay_sum),
                        edge);
   }
   std::sort(edges.begin(), edges.end());
@@ -548,15 +565,22 @@ std::vector<std::string> Search::GetVerboseStats(const Node* node) const {
     const auto& edge = std::get<2>(edge_tuple);
     float Q = edge.GetQ(fpu, draw_score);
     float M = m_evaluator.GetMUtility(edge, Q);
+    // Calculate P_eff for display (keeping label as "P" for compatibility)
+    float p_eff = edge.GetP();
+    if (policy_decay_scale > 0.0f && p_eff > 0.0f) {
+      p_eff = ApplyPolicyDecay(p_eff, static_cast<float>(edge.GetN()),
+                               policy_decay_scale, node->GetNumEdges());
+      p_eff /= policy_decay_sum;
+    }
     std::ostringstream oss;
     oss << std::left;
     // TODO: should this be displaying transformed index?
     print_head(&oss, edge.GetMove(is_black_to_move).ToString(true),
                MoveToNNIndex(edge.GetMove(), 0), edge.GetN(),
-               edge.GetNInFlight(), edge.GetP());
+               edge.GetNInFlight(), p_eff);
     print_stats(&oss, edge.node());
-    print(&oss, "(U: ", edge.GetU(U_coeff), ") ", 6, 5);
-    print(&oss, "(S: ", Q + edge.GetU(U_coeff) + M, ") ", 8, 5);
+    print(&oss, "(U: ", edge.GetU(U_coeff, node->GetNumEdges(), policy_decay_scale, policy_decay_sum), ") ", 6, 5);
+    print(&oss, "(S: ", Q + edge.GetU(U_coeff, node->GetNumEdges(), policy_decay_scale, policy_decay_sum) + M, ") ", 8, 5);
     print_tail(&oss, edge.node());
     infos.emplace_back(oss.str());
   }
@@ -1710,6 +1734,24 @@ void SearchWorker::PickNodesToExtendTask(
       const float cpuct = ComputeCpuct(params_, node->GetN(), is_root_node);
       const float puct_mult =
           cpuct * std::sqrt(std::max(node->GetChildrenVisits(), 1u));
+      const float policy_decay_scale = params_.GetPolicyDecayScale();
+
+      // Calculate sum of raw P_eff for normalization
+      float policy_decay_sum = 1.0f;
+      if (policy_decay_scale > 0.0f) {
+        policy_decay_sum = 0.0f;
+        auto edge_iter = node->Edges();
+        for (int i = 0; i < max_needed; ++i, ++edge_iter) {
+          float p = current_pol[i];
+          if (p > 0.0f) {
+            policy_decay_sum += ApplyPolicyDecay(
+                p, static_cast<float>(edge_iter.GetN()), policy_decay_scale,
+                node->GetNumEdges());
+          }
+        }
+        if (policy_decay_sum == 0.0f) policy_decay_sum = 1.0f;
+      }
+
       int cache_filled_idx = -1;
       while (cur_limit > 0) {
         // Perform UCT for current node.
@@ -1732,8 +1774,14 @@ void SearchWorker::PickNodesToExtendTask(
           int nstarted = current_nstarted[idx];
           const float util = current_util[idx];
           if (idx > cache_filled_idx) {
-            current_score[idx] =
-                current_pol[idx] * puct_mult / (1 + nstarted) + util;
+            float p = current_pol[idx];
+            if (policy_decay_scale > 0.0f && p > 0.0f) {
+              p = ApplyPolicyDecay(p, static_cast<float>(cur_iters[idx].GetN()),
+                                   policy_decay_scale, node->GetNumEdges());
+              // Normalize by sum to ensure sum(P_eff) = 1.0
+              p /= policy_decay_sum;
+            }
+            current_score[idx] = p * puct_mult / (1 + nstarted) + util;
             cache_filled_idx++;
           }
           if (is_root_node) {
@@ -2074,11 +2122,28 @@ int SearchWorker::PrefetchIntoCache(Node* node, int budget, bool is_odd_depth) {
       cpuct * std::sqrt(std::max(node->GetChildrenVisits(), 1u));
   const float fpu =
       GetFpu(params_, node, node == search_->root_node_, draw_score);
+  const float policy_decay_scale = params_.GetPolicyDecayScale();
+
+  // Calculate sum of raw P_eff for normalization
+  float policy_decay_sum = 1.0f;
+  if (policy_decay_scale > 0.0f) {
+    policy_decay_sum = 0.0f;
+    for (const auto& edge : node->Edges()) {
+      float p = edge.GetP();
+      if (p > 0.0f) {
+        policy_decay_sum += ApplyPolicyDecay(
+            p, static_cast<float>(edge.GetN()), policy_decay_scale,
+            node->GetNumEdges());
+      }
+    }
+    if (policy_decay_sum == 0.0f) policy_decay_sum = 1.0f;
+  }
+
   for (auto& edge : node->Edges()) {
     if (edge.GetP() == 0.0f) continue;
     // Flip the sign of a score to be able to easily sort.
     // TODO: should this use logit_q if set??
-    scores.emplace_back(-edge.GetU(puct_mult) - edge.GetQ(fpu, draw_score),
+    scores.emplace_back(-edge.GetU(puct_mult, node->GetNumEdges(), policy_decay_scale, policy_decay_sum) - edge.GetQ(fpu, draw_score),
                         edge);
   }
 

src/utils/fastmath.h

@@ -104,4 +104,41 @@ inline float FastSign(const float a) {
 #endif
 }
 
+// Fast approximate 1/sqrt(x) using bit manipulation.
+// Based on the classic Quake III algorithm.
+// Expects positive input values. Does no range checking; produces undefined
+// behavior for zero, negative, or special values (NaN, infinity).
+inline float FastInvSqrt(const float a) {
+  float halfx = 0.5f * a;
+  uint32_t i;
+  std::memcpy(&i, &a, sizeof(float));
+  i = 0x5f3759df - (i >> 1);  // Magic constant
+  float y;
+  std::memcpy(&y, &i, sizeof(float));
+  y = y * (1.5f - halfx * y * y);  // Newton iteration
+  return y;
+}
+
+// Apply positive policy decay transformation with fixed sqrt decay (exponent=0.5).
+// Returns raw (unnormalized) P_eff = 1 / (1 + odds * power_term) where:
+//   odds = 1/P - 1
+//   effective_scale = scale_per_move * num_legal_moves
+//   power_term = 1/sqrt(1 + N/effective_scale)
+// When P=0 or num_legal_moves<=0, returns P unchanged.
+// NOTE: Caller must normalize by sum of all raw_P_eff values to ensure sum(P_eff) = 1.
+// Uses FastInvSqrt for optimal performance.
+inline float ApplyPolicyDecay(float p, float n_child, float scale_per_move,
+                               int num_legal_moves) {
+  if (p == 0.0f || num_legal_moves <= 0) return p;
+
+  float effective_scale = scale_per_move * num_legal_moves;
+  float base = 1.0f + n_child / effective_scale;
+
+  // Sqrt decay: (1 + N/scale)^(-0.5) = 1/sqrt(1 + N/scale)
+  float power_term = FastInvSqrt(base);
+
+  float odds = 1.0f / p - 1.0f;
+  return 1.0f / (1.0f + odds * power_term);
+}
+
 }  // namespace lczero