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
|
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:
|
||||||
|
|||||||
Reference in New Issue
Block a user