Collaborative-Fraud-Intelligence-Simulator / backend /app /application /services /fl_hyperparameter_optimizer.py
| """Automated FL Hyperparameter Optimizer (Optuna / Bayesian TPE). | |
| Executes Bayesian optimization over federated learning hyperparameters (learning rate, local epochs, | |
| DP clip norm, noise multiplier, staleness decay, FedProx penalty mu) tailored to Non-IID bank data distributions. | |
| """ | |
| from __future__ import annotations | |
| import logging | |
| import time | |
| from typing import Any | |
| import numpy as np | |
| import optuna | |
| from app.application.services.fl_dirichlet_partitioner import DirichletPartitioner | |
| from app.application.services.fl_engine import FederatedLearningEngine as FLEngine | |
| from app.domain.value_objects import ModelWeights | |
| logger = logging.getLogger(__name__) | |
| # Suppress noisy Optuna verbosity by default | |
| optuna.logging.set_verbosity(optuna.logging.WARNING) | |
| class FLHyperparameterOptimizer: | |
| """Automated Optuna Bayesian TPE Hyperparameter Tuning Engine for FL.""" | |
| def __init__( | |
| self, | |
| study_name: str = "fl_hpo_study", | |
| dirichlet_alpha: float = 0.5, | |
| num_clients: int = 3, | |
| num_rounds: int = 5, | |
| seed: int = 42, | |
| ) -> None: | |
| self.study_name = study_name | |
| self.dirichlet_alpha = dirichlet_alpha | |
| self.num_clients = num_clients | |
| self.num_rounds = num_rounds | |
| self.seed = seed | |
| # Create Optuna TPE sampler and MedianPruner | |
| self.sampler = optuna.samplers.TPESampler(seed=seed) | |
| self.pruner = optuna.pruners.MedianPruner(n_startup_trials=2, n_warmup_steps=1) | |
| self.study = optuna.create_study( | |
| study_name=study_name, | |
| direction="maximize", | |
| sampler=self.sampler, | |
| pruner=self.pruner, | |
| ) | |
| def _generate_synthetic_bank_data(self) -> tuple[np.ndarray, np.ndarray]: | |
| """Generate reproducible synthetic fraud transaction features and labels.""" | |
| rng = np.random.default_rng(self.seed) | |
| n_samples = 600 | |
| n_features = 10 | |
| # Class 0: Legitimate, Class 1: Fraud | |
| X0 = rng.normal(loc=0.0, scale=1.0, size=(int(n_samples * 0.9), n_features)) | |
| y0 = np.zeros(int(n_samples * 0.9)) | |
| X1 = rng.normal(loc=2.5, scale=1.2, size=(int(n_samples * 0.1), n_features)) | |
| y1 = np.ones(int(n_samples * 0.1)) | |
| X = np.vstack([X0, X1]) | |
| y = np.concatenate([y0, y1]) | |
| perm = rng.permutation(len(y)) | |
| return X[perm], y[perm] | |
| def objective(self, trial: optuna.Trial) -> float: | |
| """Evaluate a single trial set of candidate FL hyperparameters using real PyTorch training.""" | |
| # 1. Suggest hyperparameters | |
| lr = trial.suggest_float("learning_rate", 1e-4, 1e-1, log=True) | |
| local_epochs = trial.suggest_int("local_epochs", 1, 5) | |
| dp_clip_norm = trial.suggest_float("dp_clip_norm", 0.1, 5.0) | |
| dp_noise_multiplier = trial.suggest_float("dp_noise_multiplier", 0.1, 2.0) | |
| staleness_gamma = trial.suggest_float("staleness_gamma", 0.1, 3.0) | |
| fedprox_mu = trial.suggest_float("fedprox_mu", 1e-3, 1.0, log=True) | |
| batch_size = trial.suggest_categorical("batch_size", [16, 32, 64]) | |
| # 2. Generate dataset and partition across bank nodes using Dirichlet Dir(alpha) | |
| X, y = self._generate_synthetic_bank_data() | |
| split_idx = int(len(X) * 0.8) | |
| X_train, X_val = X[:split_idx], X[split_idx:] | |
| y_train, y_val = y[:split_idx], y[split_idx:] | |
| client_datasets = DirichletPartitioner.partition_dataset( | |
| features=X_train, | |
| labels=y_train, | |
| num_clients=self.num_clients, | |
| alpha=self.dirichlet_alpha, | |
| seed=self.seed + trial.number, | |
| ) | |
| from app.application.services.model_service import ModelService | |
| from app.application.services.privacy_service import PrivacyService | |
| from app.config import get_settings | |
| settings = get_settings() | |
| model_service = ModelService(settings) | |
| privacy_service = PrivacyService() | |
| fl_engine = FLEngine(settings, model_service, privacy_service) | |
| # Initialize global model with reproducible seed | |
| rng = np.random.default_rng(self.seed) | |
| global_model = model_service.create_model(dp_compatible=True) | |
| global_weights = model_service.get_parameters(global_model) | |
| final_score = 0.5 | |
| # 3. Simulate federated rounds with genuine local PyTorch training | |
| for round_idx in range(1, self.num_rounds + 1): | |
| client_updates = [] | |
| client_samples = [] | |
| # Staleness attenuation factor S(tau) = (1 + 0.1 * round)^(-gamma) | |
| staleness_weight = (1.0 + (round_idx - 1) * 0.1) ** (-staleness_gamma) | |
| for client_idx, (X_i, y_i) in enumerate(client_datasets): | |
| # Instantiate client model matching global weights | |
| local_model = model_service.create_model(dp_compatible=True) | |
| local_model = model_service.set_parameters(local_model, global_weights) | |
| # Train genuine local PyTorch model with FedProx proximal regularization | |
| local_model, _, _ = model_service.train_local( | |
| model=local_model, | |
| X_train=X_i, | |
| y_train=y_i, | |
| epochs=local_epochs, | |
| batch_size=batch_size, | |
| learning_rate=lr, | |
| global_weights=global_weights, | |
| fedprox_mu=fedprox_mu, | |
| ) | |
| updated_weights = model_service.get_parameters(local_model) | |
| # Calibrate Differential Privacy noise perturbation on local weights | |
| dp_noise = rng.normal( | |
| 0.0, | |
| (dp_noise_multiplier * dp_clip_norm) / np.sqrt(max(len(X_i), 1)), | |
| len(updated_weights.flat_weights), | |
| ) | |
| noisy_flat = (np.array(updated_weights.flat_weights) + dp_noise).tolist() | |
| noisy_weights = ModelWeights( | |
| layer_shapes=updated_weights.layer_shapes, | |
| flat_weights=noisy_flat, | |
| ) | |
| client_updates.append(noisy_weights) | |
| # Attenuate sample influence by staleness factor | |
| effective_samples = max(1, int(len(X_i) * staleness_weight)) | |
| client_samples.append(effective_samples) | |
| # Aggregate round updates via sample-weighted consensus | |
| global_weights = fl_engine.aggregate_parameters(client_updates, client_samples) | |
| global_model = model_service.set_parameters(global_model, global_weights) | |
| # Evaluate round performance on real holdout validation dataset | |
| eval_metrics = model_service.evaluate(global_model, X_val, y_val) | |
| raw_round_auc = eval_metrics.get("auc_roc") | |
| if raw_round_auc is None: | |
| logger.warning( | |
| "Trial %d produced undefined validation AUC at round %d; pruning trial explicitly", | |
| trial.number, | |
| round_idx, | |
| ) | |
| raise optuna.TrialPruned(f"Undefined validation AUC at round {round_idx}") | |
| round_auc = float(raw_round_auc) | |
| final_score = round_auc | |
| # Report step metric to Optuna for early MedianPruner decisions | |
| trial.report(round_auc, round_idx) | |
| if trial.should_prune(): | |
| logger.info( | |
| "Trial %d pruned at round %d (validation AUC: %.4f)", | |
| trial.number, | |
| round_idx, | |
| round_auc, | |
| ) | |
| raise optuna.TrialPruned() | |
| return final_score | |
| def run_optimization(self, n_trials: int = 5, timeout: float | None = 60.0) -> dict[str, Any]: | |
| """Run Optuna optimization study over n_trials with graceful pruning handling.""" | |
| start_time = time.perf_counter() | |
| logger.info( | |
| "Starting Optuna FL hyperparameter optimization study '%s' (Dirichlet alpha=%.2f, trials=%d)...", | |
| self.study_name, | |
| self.dirichlet_alpha, | |
| n_trials, | |
| ) | |
| self.study.optimize(self.objective, n_trials=n_trials, timeout=timeout) | |
| duration = (time.perf_counter() - start_time) * 1000 | |
| default_params = { | |
| "learning_rate": 0.01, | |
| "local_epochs": 2, | |
| "dp_clip_norm": 1.0, | |
| "dp_noise_multiplier": 0.5, | |
| "staleness_gamma": 0.5, | |
| "fedprox_mu": 0.01, | |
| "batch_size": 32, | |
| } | |
| # Safe best trial extraction (handles all-pruned or empty study cases) | |
| completed_trials = [ | |
| t for t in self.study.trials if t.state == optuna.trial.TrialState.COMPLETE | |
| ] | |
| if completed_trials: | |
| best_t = self.study.best_trial | |
| best_trial_number = best_t.number | |
| best_value = float(best_t.value) if best_t.value is not None else 0.5 | |
| best_params = best_t.params if best_t.params else default_params | |
| elif self.study.trials: | |
| best_t = max( | |
| self.study.trials, | |
| key=lambda t: max(t.intermediate_values.values()) if t.intermediate_values else -1.0, | |
| ) | |
| best_trial_number = best_t.number | |
| best_value = ( | |
| max(best_t.intermediate_values.values()) if best_t.intermediate_values else 0.5 | |
| ) | |
| best_params = best_t.params if best_t.params else default_params | |
| else: | |
| best_trial_number = -1 | |
| best_value = 0.5 | |
| best_params = default_params | |
| try: | |
| if len(completed_trials) >= 2: | |
| importances = optuna.importance.get_param_importances(self.study) | |
| else: | |
| importances = {} | |
| except Exception as exc: | |
| logger.debug("Could not compute Optuna parameter importances: %s", exc) | |
| importances = {} | |
| results = { | |
| "study_name": self.study_name, | |
| "dirichlet_alpha": self.dirichlet_alpha, | |
| "best_trial_number": best_trial_number, | |
| "best_value": round(best_value, 4), | |
| "best_params": best_params, | |
| "param_importances": importances, | |
| "total_trials": len(self.study.trials), | |
| "completed_trials": len(completed_trials), | |
| "pruned_trials": len( | |
| [t for t in self.study.trials if t.state == optuna.trial.TrialState.PRUNED] | |
| ), | |
| "duration_ms": round(duration, 2), | |
| } | |
| logger.info( | |
| "Completed Optuna study '%s' in %.2fms. Best trial #%d (AUC: %.4f)", | |
| self.study_name, | |
| duration, | |
| best_trial_number, | |
| best_value, | |
| ) | |
| return results | |