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:
@@ -62,6 +62,11 @@ class TrainingConfig(StrictModel):
|
||||
interleave_max_batch: int = 64
|
||||
num_workers: int = 1
|
||||
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(
|
||||
"games_per_iter",
|
||||
|
||||
@@ -261,7 +261,8 @@ class IsMctsTrainer:
|
||||
policy_loss = -(pi * log_probs).sum(dim=-1).mean()
|
||||
v_scale = float(self.network.value_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)
|
||||
loss.backward()
|
||||
if self.config.optimization.grad_clip > 0:
|
||||
|
||||
Reference in New Issue
Block a user