Add training.value_loss_weight (default 1.0) for value loss reweighting

Diagnosis: mcts/value_prediction_error stuck at 300-1000 (RMSE ~22 on score
range ±100), while loss/value stays at 0.05 because the loss divides
prediction and target by value_scale=100 (so MSE / 10000). Net effect: the
value head receives a tiny gradient relative to the policy head's
~1.75 cross-entropy loss, so it never learns to predict score scale well.

This adds a config knob to multiply the normalized value loss without
re-engineering the loss formula. value_loss_weight=50 recovers the raw
MSE magnitude (~2.5 vs policy loss ~1.75), giving the value head
comparable gradient signal.
This commit is contained in:
2026-05-11 06:34:05 +09:00
parent 169d4dcb14
commit be0c1a8d62
2 changed files with 7 additions and 1 deletions
@@ -62,6 +62,11 @@ class TrainingConfig(StrictModel):
interleave_max_batch: int = 64 interleave_max_batch: int = 64
num_workers: int = 1 num_workers: int = 1
worker_device: str = "cpu" worker_device: str = "cpu"
# Multiplier on the value-head MSE loss (already normalized by value_scale**2).
# Default 1.0 keeps current behavior; raising it (e.g. 50-100) makes the value
# head learn faster relative to policy loss. Useful when value_prediction_error
# is large but loss/value is tiny because of the normalization.
value_loss_weight: float = 1.0
@field_validator( @field_validator(
"games_per_iter", "games_per_iter",
@@ -261,7 +261,8 @@ class IsMctsTrainer:
policy_loss = -(pi * log_probs).sum(dim=-1).mean() policy_loss = -(pi * log_probs).sum(dim=-1).mean()
v_scale = float(self.network.value_scale) v_scale = float(self.network.value_scale)
value_loss = nn.functional.mse_loss(value_pred / v_scale, value_target / v_scale) value_loss = nn.functional.mse_loss(value_pred / v_scale, value_target / v_scale)
loss = policy_loss + value_loss value_weight = float(self.config.training.value_loss_weight)
loss = policy_loss + value_weight * value_loss
self.optimizer.zero_grad(set_to_none=True) self.optimizer.zero_grad(set_to_none=True)
loss.backward() loss.backward()
if self.config.optimization.grad_clip > 0: if self.config.optimization.grad_clip > 0: