Skip to content

Fixes in Nn uncertainty model - #70

Open
faildeny wants to merge 8 commits into
cmd_linefrom
nn_uncertainty
Open

faildeny wants to merge 8 commits into
cmd_linefrom
nn_uncertainty

Conversation

@faildeny

Copy link
Copy Markdown
Contributor

Fix NN uncertainty evaluation and UncertaintyWeighted aggregation

Changes

  • Binary MC-dropout metrics fixed. predict_proba_mc applied softmax to a single logit, so binary models always reported entropy 0 and an accuracy equal to the share of class 0. It now uses sigmoid. evaluate() always reports the standard metrics and adds the MC-dropout ones (entropy, val_accuracy) only for classification.
  • Saved NN models load in inference.py. The model metadata now includes dropout_p and mc_samples (= T). Without them load_model refused to rebuild the model.
  • n_out is consistent across sites. load_dt4h no longer replaces the metadata class count with the number of classes present locally. A site missing a class used to get a different output layer from the other sites.
  • UncertaintyWeighted aggregates evaluation metrics. The strategy now receives evaluate_metrics_aggregation_fn, so metrics_distributed is populated.
  • Regression falls back to FedAvg. Clients only report entropy for classification, so UncertaintyWeighted with a regression task now switches to FedAvg at config time.
  • accept_failures is respected in UncertaintyWeightedFedAvg.aggregate_fit.
  • Tests: a new test_nn_uncertainty_weighted covers binary, multiclass and regression with MC dropout on. fl_runner gains server_extra_args. CI also runs on PRs to tests_and_CI.

New server arguments (server_cmd.py)

Argument Default Effect
--uw_epsilon 0.1 Smoothing term in the client weight n / (ε + entropy)^γ. Limits how much a client with near-zero entropy can dominate (max ratio ≈ (1+ε)/ε).
--uw_gamma 1.0 Strength of the uncertainty weighting. 0 is plain FedAvg; larger values give more weight to low-entropy clients.

Both only apply with --strategy UncertaintyWeighted.

@faildeny
faildeny requested a review from J-Fabila October 29, 2026 12:32
@faildeny faildeny self-assigned this Oct 29, 2026
@faildeny
faildeny changed the base branch from tests_and_CI to cmd_line October 30, 2026 07:19
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant