Add corrector regularization to training

[mcgibbon]

Add an optional training-time penalty on the magnitude of corrector adjustments, enabling L1/L2 regularization that discourages the model from relying on the corrector. Supported for both single-module and coupled training; when configured, the specified loss is computed between each corrector's per-variable correction (post-corrector minus pre-corrector) and a zero baseline in the loss-normalized space, then accumulated into the training loss.

Changes:

• fme.core.step.step.StepResult: new dataclass replacing the bare TensorDict return type of StepABC.step; carries the denormalized output plus optional per-variable corrections. All StepABC implementations and consumers updated.

• fme.core.step.single_module.step_with_adjustments: records corrections = post − pre when a corrector runs; MultiCallStep forwards corrections from its wrapped step.

• fme.core.loss.CorrectorRegularizationConfig: new LossConfig + weight wrapper, with validation rejecting EnsembleLoss/NaN/global_mean_type.

• fme.ace.stepper.single_module.TrainStepperConfig.corrector_regularization: optional field; built via Stepper.build_corrector_regularization_loss and applied per optimized step in TrainStepper._accumulate_loss. Per-step (`corrector_regularization_step_{i}`) and epoch-aggregated (`corrector_regularization`) metrics.

• fme.coupled.stepper.ComponentTrainingConfig.corrector_regularization: per-realm version. `CoupledStepperTrainLoss.compute_corrector_regularization` wired into `CoupledTrainStepper._accumulate_step_loss`; emits `loss/{realm}corrector_regularization_step{i}`.

• `fme.ace` re-exports `LossConfig` and `CorrectorRegularizationConfig` so nested-dataclass symbol checks see them.

• Tests added

• If dependencies changed, "deps only" image rebuilt and "latest_deps_only_image.txt" file updated

[jpdunc23]

What loss / weight do you suggest I try with this feature? Thinking about launching two experiments, one with an L1 penalty and another with an L2 penalty, but lmk if you have a particular config in mind.

[mcgibbon]

What loss / weight do you suggest I try with this feature? Thinking about launching two experiments, one with an L1 penalty and another with an L2 penalty, but lmk if you have a particular config in mind.

Yeah that matches what I was thinking.

The normalization is the same as for the main loss term, which helps reason about weighting. L2 is probably the better match for the behavior you want, and also allows the model to keep slightly-negative zero-values with little penalty, so I'd focus on that if you want to do a weight sweep. 1.0 is a reasonable starting point for weighting, which means the model puts equal importance to obeying the corrector's goals and to having skill. Even a smaller value like 0.05-0.2 should significantly reduce the divergence between pre-and post-corrector values, because you'll only get divergence to the extent that this helps the model predict better. If you got large divergence even with 0.2 weight, it would imply the divergence is fairly significantly helping the model's skill, which I don't expect.

Maybe try 0.2, 1.0, and 5.0 on L2 as a first pass if you want to sweep, or 1.0 on L2 if you're limiting to one config?

[jpdunc23]

@mcgibbon please merge #1223 when you can (fixed on my exper branch after failing job)

[jpdunc23]

Experiment here: https://wandb.ai/ai2cm/ace-samudra-cm4/runs/xzmbufek

[mcgibbon]

Experiment here: https://wandb.ai/ai2cm/ace-samudra-cm4/runs/xzmbufek

Looks like so far it's improving the thetao evolution as much as the precorrector optimization. To be seen if his continues, or if the regularization approach bottoms out earlier on how much improvement it gives.

[jpdunc23]

Could you please add a feature to exclude the penalty for specific variables? Excluding sea ice in pre-corrector optim removed the small positive sea ice biases in mid latitudes, so I think it should do the same here.

Here's an example from my PR of configuring this using the NameAndPrefixMatcher class:

ace/fme/ace/stepper/single_module.py

Lines 1462 to 1515 in 046a3f8

	@dataclasses.dataclass
	class PreCorrectorOptimizationConfig:
	"""
	Configuration enabling pre-corrector optimization of the training loss.
	When this config is present, the training loss for corrector-modified
	variables is computed against the model's pre-correction (uncorrected)
	output rather than its corrected output, except for variables matched by
	``exclude_names_and_prefixes``, which continue to be optimized against their
	corrected (post-adjustment) values. The presence of this object is itself
	the on/off switch; there is no separate ``enabled`` flag.
	The rollout state and returned predictions always use the fully corrected
	outputs regardless of this config; only the training loss target is affected.
	Parameters:
	exclude_names_and_prefixes: Names and prefixes of variables to exclude
	from pre-corrector optimization. Matching follows the same
	name-and-prefix convention as spatial masking: a bare name matches
	the 2D variable and all of its 3D levels, a trailing-underscore
	prefix (e.g. ``thetao_``) matches all levels, and an explicit
	``name_<level>`` matches exactly.
	"""
	exclude_names_and_prefixes: list[str] = dataclasses.field(default_factory=list)
	def __post_init__(self):
	self._matcher = NameAndPrefixMatcher(self.exclude_names_and_prefixes)
	def select_precorrected(self, uncorrected: TensorMapping) -> TensorDict:
	"""Return the subset of ``uncorrected`` to override into the loss target.
	Excluded variables are dropped so that they retain their corrected value
	in the loss; all other corrector-modified variables are returned so they
	override their corrected value with their pre-correction value.
	"""
	return {
	name: value
	for name, value in uncorrected.items()
	if not self._matcher.matches(name)
	}
	def unmatched_exclusions(self, names: Iterable[str]) -> list[str]:
	"""Return exclusion entries that match none of ``names``.
	Used for warn-once validation against the actual set of
	corrector-modified variables, which is only known at runtime.
	"""
	names = list(names)
	return [
	entry
	for entry in self.exclude_names_and_prefixes
	if not any(NameAndPrefixMatcher([entry]).matches(name) for name in names)
	]