Add combined infill-prediction multi-task training

[mcgibbon]

Adds a new InfillPredictionStep that trains on multiple tasks per batch (auto-encoding, infill, prediction, infill-prediction, combined-all) by randomly masking input/output variables. At inference time it behaves like SingleModuleStep. Per gradient update, the normal prediction loss for all forward steps is computed, then one additional forward pass with task masking adds a task loss.

Task sampling is generic and works with any step type — variable pools are derived from StepConfigABC.accepted_input_names and loss_names, with config-time validation that raises if enabled tasks are infeasible for the step's variable routing.

Changes:

• fme.core.step.task_step: New TaskConfig, TaskWeights, TaskSamplingConfig, TaskSampler, SampledTasks, InferenceSchemeConfig, InfillPredictionStepConfig, InfillPredictionStep

  • TaskSampler samples per-batch-element task types and variable assignments respecting disjointness and minimum-variable constraints
  • InfillPredictionStep uses all_names for input channels with channel mask indicators and non_forcing_names for output channels; fills missing variables with zeros at inference

• fme.core.step.step.StepConfigABC: Add accepted_input_names, accepted_next_step_input_names, and all_training_names properties so step configs can declare training-time variable capabilities distinct from inference-time routing

• fme.ace.stepper.single_module.StepperConfig.all_names: Check all_training_names first, fall back to set(input_names + output_names)

• fme.ace.stepper.single_module.TrainStepperConfig: Add optional task_sampling: TaskSamplingConfig field

• fme.ace.stepper.single_module.TrainStepper: When task_sampling is set, build TaskSampler from generic accepted_input_names/loss_names; after the prediction loop, run one additional task step and accumulate its loss

  • Handles n_ensemble > 1 by repeat-interleaving per-IC masks/ground-truth across ensemble members
  • Sources forcing from derived input_only_names (accepted inputs minus loss names)

• fme.ace.train.train_config.TrainConfig.__post_init__: Config-time validation via TaskSamplingConfig.validate_for_step

• fme.ace.stepper.loss_schedule.LossSchedule: Add is_training property

• Tests added

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

[mcgibbon]

Design document for this feature (removed from the repo in the final commit):
https://github.com/ai2cm/ace/blob/b56ae66f306279a0b397b26c69109dcf252a4398/infill-prediction-design.md

[mcgibbon]

Example configuration

Below is an example YAML snippet showing how to configure multi-task training. The new configuration lives in two places: the step type (infill_prediction under stepper.step) and the task sampling options (stepper_training.task_sampling).

stepper:
  step:
    type: infill_prediction
    config:
      builder:
        type: SphericalFourierNeuralOperatorNet
        config:
          scale_factor: 1
          embed_dim: 256
          num_layers: 8
      # All variables the model can see during training.
      # Drives data loading — every name here must be in the dataset.
      all_names:
        - air_temperature
        - specific_humidity
        - surface_pressure
        - precipitation_rate   # non-forcing, training-only (not in inference scheme)
        - insolation           # forcing
        - land_sea_mask        # forcing
      # Forcing variables are input-only: never predicted by the model.
      forcing_names:
        - insolation
        - land_sea_mask
      normalization:
        network:
          means: {air_temperature: 250.0, specific_humidity: 0.005, ...}
          stds:  {air_temperature: 30.0, specific_humidity: 0.004, ...}
      # Inference scheme defines the variable routing at inference time.
      # Must be a subset of all_names.
      inference_scheme:
        in_names:
          - air_temperature
          - specific_humidity
          - surface_pressure
          - insolation
          - land_sea_mask
        out_names:
          - air_temperature
          - specific_humidity
          - surface_pressure
        next_step_forcing_names: []
        prescribed_prognostic_names: []
      # Must be True (the network receives 2x channels: values + mask indicators).
      include_channel_mask_inputs: true
      # Optional, same as SingleModuleStepConfig:
      # ocean: ...
      # corrector: ...
      # secondary_decoder: ...

stepper_training:
  loss:
    type: MSE
  task_sampling:
    tasks:
      # Each task has a probability (relative sampling weight, normalized to
      # sum to 1) and loss_scaling (multiplier on the output mask).
      # Set probability to 0 to disable a task.
      auto_encode:
        probability: 1.0
        loss_scaling: 1.0
      infill:
        probability: 1.0
        loss_scaling: 1.0
      prediction:
        probability: 1.0
        loss_scaling: 1.0
      infill_prediction:
        probability: 1.0
        loss_scaling: 1.0
      combined_all:
        probability: 1.0
        loss_scaling: 1.0
    # Minimum number of input/output variables selected per sample.
    # Must be >= 1.
    min_input_variables: 1
    min_output_variables: 1

What you can configure

• Which tasks are active — set any task's probability to 0 to disable it; the remaining probabilities are normalized.
• Relative task frequency — e.g. prediction: {probability: 3.0} and infill: {probability: 1.0} means prediction is sampled 3x more often.
• Per-task loss scaling — loss_scaling multiplies the output mask so the loss is up-weighted for that task's samples, independently of sampling frequency.
• Minimum variable counts — min_input_variables and min_output_variables set floors on how many variables are selected per sample. The sampler will raise at init if a task can't satisfy these constraints given the available variables.
• Training-only variables — all_names can include variables not in inference_scheme.in_names/out_names. These are loaded for training tasks but not required at inference. The model's input channel count is 2 * len(all_names) (values + channel masks), and output channel count is len(all_names) - len(forcing_names).
• Inference behavior — inference_scheme controls inference-time variable routing identically to SingleModuleStepConfig (in_names, out_names, next_step_forcing_names, prescribed_prognostic_names). The step fills missing training-only variables with zeros and marks them absent via the channel mask.
• Standard step options — ocean, corrector, and secondary_decoder are supported as on SingleModuleStepConfig.
• Any step type — task_sampling works with any step type, not just infill_prediction. Variable pools are derived from the step's accepted_input_names and loss_names. Config validation raises if enabled tasks are infeasible for the step's variable routing.

What you can't configure

• Task definitions themselves — the five task types (auto_encode, infill, prediction, infill_prediction, combined_all) and their variable selection constraints are fixed. You can't add custom tasks or change the disjointness rules.
• Variable selection distribution — the algorithm uses uniform random subsets with a symmetric input/output split for disjoint tasks. There's no way to bias selection toward specific variables or change the distribution shape.
• Per-variable weights — you can't assign different importance to individual variables within the task sampling (though the standard loss config's per-channel weights still apply to the task loss).
• Task step placement — the task step always runs once after the final prediction step, using the last generated state as "previous step" data. You can't run it at intermediate steps or run multiple task steps per update.
• Separate task loss config — the task loss reuses the same StepLossConfig as the prediction loss. There's no way to use a different loss type (e.g. MSE for prediction, L1 for tasks).
• include_channel_mask_inputs: false — this must be True for InfillPredictionStep; the network needs the mask indicators to know which inputs are real vs. masked zeros.
• Residual prediction — not supported with InfillPredictionStep (always disabled).
• n_ic_timesteps > 1 — always 1; multi-timestep initial conditions are not supported.

[mcgibbon]

This file should probably be named something more specific to the contents of the file than the purpose of our experiments right now, like this seems more related to "tasks" than infill and prediction specifically.

[mcgibbon]

Renamed to task_step.py (and the test to test_task_step.py).

[mcgibbon]

Instead of having two properties for each task, can we have a sub-configuration for each task that contains these two properties? Something like "probability" and "loss_scaling"?

[mcgibbon]

Restructured to use a TaskConfig sub-dataclass with probability and loss_scaling fields per task.

[mcgibbon]

This has weights and probabilities, and controls what tasks we're using, so let's just call it "tasks:".

[mcgibbon]

Renamed to tasks.

[mcgibbon]

nit: Let's not discuss future plans in a docstring.

[mcgibbon]

Removed.

[mcgibbon]

Is this an enforced assumption? if not, we should raise when it's False.

Ideally we should raise at configuration validation time, if this is required.

[mcgibbon]

The validation already happens at TrainStepper.__init__ time (raises ValueError if step config isn't InfillPredictionStepConfig). Replaced the cast with a stored typed attribute (self._task_step_config) set during init, so there's no runtime cast needed.

[mcgibbon]

Great. Is there also a __post_init__ configuration check? When we run fme.ace.validate_config, I want an exception raised if it's invalid.

[mcgibbon]

Pre-review agent: Addressed in 7494484. Added TaskSamplingConfig.validate_for_step(accepted_input_names, loss_names) which is called from TrainConfig.__post_init__, so fme.ace.validate_config will raise at config-validation time. Additionally, this commit generalizes task sampling to work with any step type by using accepted_input_names/loss_names from StepConfigABC instead of requiring InfillPredictionStepConfig.

[mcgibbon]

nit: Remove comment about how this was committed.

[mcgibbon]

Removed.

[mcgibbon]

I'm working on updating this so that it works seamlessly with all Step types so long as their input/output names are appropriate to the task. The InfillPredictionStep will still be needed for most tasks since we don't have other steps that accept output-step values as inputs, but we could use e.g. a task configuration using only prediction along with the existing SingleModuleStep.

[mcgibbon]

This PR should add test_train.py-style tests that cover at least SingleModuleStep in prediction-only mode, and the infill step in a more complex task mode. It would be nice if they can run as fast as possible.

[mcgibbon]

OK sorry my last comment was not fully right. The infill prediction step should exist in its own file, appropriately named. The task sampling logic probably doesn't belong in fme/core/step at all. It might even be best suited for fme/ace/stepper/task.py.

[mcgibbon]

Claude: Done. Moved task sampling logic (TaskConfig, TaskWeights, TaskSamplingConfig, TaskSampler, SampledTasks, TaskType) to fme/ace/stepper/task.py. InfillPredictionStep and InfillPredictionStepConfig stay in fme/core/step/infill_prediction.py. Tests moved to fme/ace/stepper/test_task.py.

[mcgibbon]

I think we don't want the metric name to change when we're doing task-based training. Use the existing value (is it batch_loss? or just loss at this level?).

[mcgibbon]

Claude: Removed the separate loss_task metric. The task loss is already accumulated into the optimizer before metrics["loss"] is recorded, so it's included in the total loss metric.

[mcgibbon]

Question: Doesn't the task sampler need to know about input names that are accepted at the first and at the second step, separately?

[mcgibbon]

Claude: The task sampler's "previous step" and "current step" masks both control what goes into args.input — they differ in data source (model output vs ground truth), not which step input slot they target. The next_step_input_data is built independently from forcing and is not task-sampled.

So accepted_input_names is the right single pool: it covers everything the step can receive as args.input, regardless of whether the data came from the previous step's output or current-step ground truth.

For InfillPredictionStep, accepted_input_names = all_names (broader than inference input_names). For SingleModuleStep, accepted_input_names = input_names (same as inference). In both cases, the same set applies to both previous-step and current-step sourcing.

[mcgibbon]

Question: Why is this pattern different from above, self._loss_obj = self._stepper.build_loss(config.loss)?

[mcgibbon]

Claude: Good catch — they should use the same pattern. Refactored to add an out_names parameter to Stepper.build_loss(), so both losses are now built through the same method:

self._loss_obj = self._stepper.build_loss(config.loss)
# ...
self._task_loss_obj = self._stepper.build_loss(config.loss, out_names=non_forcing)

The only difference is out_names: the regular loss uses loss_names (the default), while the task loss uses the non-forcing subset of accepted input names.

[mcgibbon]

Claude: Added two integration tests in test_train.py (commit df07464):

1. test_train_task_sampling_single_module_step — SingleModuleStep with prediction-only task sampling, runs full train_main pipeline
2. test_train_task_sampling_infill_prediction_step — InfillPredictionStep with all 5 task types enabled, runs full train_main pipeline

Both run in ~1.5s each (3s total) using minimal data and a tiny SFNO model.

[yyexela]

One thing that I think would make the config easier to parse:

• The example config has precipitation_rate in all_names but the reader has to check inference_scheme and forcing_names to see it's a training-only non-forcing variable. I recommend moving it to another place like training_only_names to make it more readable. Then, you can remove all_names from the config and build it in the code from inference_scheme inputs ∪ outputs ∪ forcing_names ∪ training_only_names

[yyexela]

Otherwise I think this looks good, but it's a lot to add to main at once, might be good to run some tests on each task to see what is most useful

[mcgibbon]

Otherwise I think this looks good, but it's a lot to add to main at once, might be good to run some tests on each task to see what is most useful

Yeah, perhaps @Arcomano1234 can give you some multi-task training assignments. For now I'll work on another PR to merge just the non-task portions of this, to reduce the chance of merge conflicts. FWIW because we don't support training backwards compatibility, it's easier to remove or modify tasks after the fact if needed, compared to removing anything that's on the module config.

[mcgibbon]

I drafted a PR that attempts to split part of this out, but it doesn't seem to make much sense on its own. Closing it.