diff --git a/src/search/classic/node.h b/src/search/classic/node.h index 8a4e598fdb..e0edf38a73 100644 --- a/src/search/classic/node.h +++ b/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(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; diff --git a/src/search/classic/params.cc b/src/search/classic/params.cc index e61b0f9c88..5c7637ffd6 100644 --- a/src/search/classic/params.cc +++ b/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(kUCIRatingAdvId, -10000.0f, 10000.0f) = 0.0f; options->Add(kSearchSpinBackoffId) = false; options->Add(kGarbageCollectionDelayId, 0.0f, 100.0f) = 10.0f; + options->Add(kPolicyDecayScaleId, 1.0f, 1000000.0f) = 100.0f; } void SearchParams::Populate(OptionsParser* options) { @@ -725,7 +732,8 @@ BaseSearchParams::BaseSearchParams(const OptionsDict& options) kMaxCollisionVisitsScalingPower( options.Get(kMaxCollisionVisitsScalingPowerId)), kSearchSpinBackoff(options_.Get(kSearchSpinBackoffId)), - kGarbageCollectionDelay(options_.Get(kGarbageCollectionDelayId)) {} + kGarbageCollectionDelay(options_.Get(kGarbageCollectionDelayId)), + kPolicyDecayScale(options_.Get(kPolicyDecayScaleId)) {} SearchParams::SearchParams(const OptionsDict& options) : BaseSearchParams(options), diff --git a/src/search/classic/params.h b/src/search/classic/params.h index d84dbad5d8..70e224948b 100644 --- a/src/search/classic/params.h +++ b/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 { diff --git a/src/search/classic/search.cc b/src/search/classic/search.cc index 101d62e941..9637fd85e8 100644 --- a/src/search/classic/search.cc +++ b/src/search/classic/search.cc @@ -465,11 +465,28 @@ std::vector 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(edge.GetN()), policy_decay_scale, + node->GetNumEdges()); + } + } + if (policy_decay_sum == 0.0f) policy_decay_sum = 1.0f; + } + std::vector> 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 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(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(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(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(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); } diff --git a/src/utils/fastmath.h b/src/utils/fastmath.h index 6aa49412aa..7846ff5aef 100644 --- a/src/utils/fastmath.h +++ b/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