flwr new @alvieupnext/fedag

FedDAG-Flower: Federated Causal Discovery with Flower

Paper: Gao, R. et al. FedDAG: Federated DAG Structure Learning. Transactions on Machine Learning Research (TMLR), 2023.

This implementation: Flower (Strategy API) + PyTorch 2, ported from the original TensorFlow 1 session-based code.

This project implements federated causal discovery: multiple clients hold local data and collaboratively learn the adjacency matrix W (d×d) of a Directed Acyclic Graph without ever sharing their raw data. They exchange only W (or W + neural network weights) with a central server that aggregates them via Federated Averaging.

---

Requirements

Declared in — pip install -e . resolves all of them:

• Python ≥ 3.11
• Flower ≥ 1.32.1 (flwr[simulation])
• PyTorch ≥ 2.3.0
• scipy ≥ 1.13, networkx ≥ 3.2.1, scikit-learn ≥ 1.4, bnlearn, openpyxl ≥ 3.1

One optional extra, needed only for hyperparameter tuning () and for bench_entrypoint.py's automatic study lookup:

pip install -e ".[tune]"     # adds optuna >= 3.0, pandas >= 1.3

Without it, tuning is unavailable and bench_entrypoint.py silently falls back to STRATEGY_DEFAULTS — nothing else breaks.

---

Execution within Flower (Flower Hub)

This algorithm is fully compatible with the modern Flower 1.8+ ServerApp / ClientApp architecture and can be deployed using standard Flower methods.

To pull and start this algorithm directly from Flower Hub into a new local project, simply run:

flwr new @alvieupnext/fedag

From there, you can use standard Flower commands (like flwr run .) to execute it in simulation or deploy it to a real federation!

---

Theory

What is FedDAG?

FedDAG learns causal Directed Acyclic Graphs (DAGs) from data distributed across multiple clients without sharing raw data. Each client holds a local dataset generated from the same (or a similar) causal structure. The server coordinates optimization of the graph structure via Federated Averaging.

Key Concepts

NOTEARS Acyclicity Constraint

The DAG constraint is enforced via a smooth function introduced by Zheng et al. (2018):

$h(W) = \mathrm\{tr\}\left(e^\{W \circ W\}\right) - d = 0 \iff W \text\{ is a DAG\}$

where $\circ$ is the element-wise product and $e^\{(\cdot)\}$ is the matrix exponential. This turns a hard combinatorial problem (finding an acyclic graph) into smooth, differentiable optimization.

Augmented Lagrangian Method

Instead of enforcing acyclicity as a hard constraint, we solve a sequence of unconstrained subproblems:

$\min_W \; \underbrace\{\frac\{1\}\{2n\}\|X - XW\|_F^2\}_\{\text\{MSE\}\} + \underbrace\{\lambda_1 \|W\|_1\}_\{\text\{sparsity\}\} + \underbrace\{\alpha \cdot h(W) + \frac\{\rho\}\{2\} h(W)^2\}_\{\text\{acyclicity penalty\}\}$

After each round, the server updates the dual variables:

\alpha \leftarrow \alpha + \rho \cdot h(W), \quad \rho \leftarrow \begin\{cases\} \rho \cdot c & \text\{if \} h(W) > \gamma \cdot h_\{\text\{prev\}\} \\ \rho & \text\{otherwise\}\end\{cases\}

where $c$ is rho-scale and $\gamma$ is h-thres (default 0.25).

Nonlinear Extension (MaskedNN)

For nonlinear SEMs, each node $j$ has its own MLP. The input is soft-masked by the graph:

$\hat\{X\}_j = \mathrm\{MLP\}_j\left(X \odot \sigma\left(\frac\{W_\{:,j\}\}\{\tau\}\right)\right)$

where $\sigma$ is the sigmoid function and $\tau$ (temperature) controls the sharpness of the mask. The graph $W$ is shared in logit space; the soft adjacency is $\sigma(W/\tau)$.

---

Getting Started

Installation

⚠️ Important — environment location

If you use venv, do not create the virtual environment inside the project folder. Flower scans the project directory at runtime and will pick up files from the venv, causing import errors and unexpected failures (e.g. flwr run crashing on packaged dependencies).

Create the venv in a sibling folder, next to the project:

parent/
├__ .venv/          ← virtual environment here
└__ FedDAG-Flower/  ← project here

If you use Miniconda / Conda, this is not a problem: conda environments live under ~/miniconda3/envs/ (or %USERPROFILE%\miniconda3\envs\ on Windows), not inside the project, so Flower won't scan them.

You can set up the environment in two ways. Miniconda is recommended if you don't already have a Python toolchain set up, as it bundles Python and isolates dependencies cleanly.

---

Option A — Miniconda (recommended)

1. Install Miniconda

If you don't already have it, download and install Miniconda from the official site. Pick the installer for your OS (Windows / macOS / Linux) and Python 3.11 or newer.

After installation, open a fresh terminal (on Windows: Anaconda Prompt or PowerShell with conda initialized) and verify:

conda --version

2. Create and activate the environment

# Create an environment named 'feddag' with Python 3.11 (the minimum this project supports)
conda create -n feddag python=3.11 -y

# Activate it
conda activate feddag

You should see (feddag) appear at the start of your prompt.

3. Install the project

# Move into the project folder
cd FedDAG-Flower

# Install the project and its dependencies in editable mode
pip install -e .

💡 We use pip install -e . inside the conda environment so that all dependencies declared in pyproject.toml (Flower, PyTorch, scipy, networkx, scikit-learn, bnlearn, openpyxl) are resolved consistently. Mixing conda install and pip install for these packages can lead to conflicts. Use pip install -e ".[tune]" instead if you also plan to run Optuna tuning.

4. Deactivating / re-entering the environment later

# Leave the environment
conda deactivate

# Re-enter it next time
conda activate feddag

5. Removing the environment (if you ever need to start over)

conda deactivate
conda env remove -n feddag

---

Option B — venv (standard library)

If you prefer Python's built-in venv, remember the warning above: create the venv outside the project folder.

From the parent folder (not from inside the project):

# Create the venv next to the project, not inside it
python -m venv .venv

# Windows:
.venv\Scripts\activate

# Linux / macOS:
source .venv/bin/activate

# Then enter the project and install
cd FedDAG-Flower
pip install -e .

Quick Demo

Run the experiment launcher in your terminal:

python scripts/run_experiments.py

You will see this menu:

   FedDAG Experiment Runner
----------------------------------------------------
  1. Run Experiments from .csv file
  2. Run an individual Experiment
  0. End

You can choose between two ways to launch experiments. There is also a third (lower-level) way using the Flower CLI directly — see More ways to run below.

---

Option 1 — Run experiments from experiments.csv

Selecting option 1 loads all experiments defined in experiments.csv and runs them in batch. Any column you leave empty falls back to the default value defined in run_experiments.py.

You will then be prompted for a range of IDs to run:

Range of IDs to run (e.g. 10-20 or a single number):

You can pick a single ID (5), a range (1-10), or a mix.

The runner shows you a summary of selected experiments before launching:

 10 Experiment(s) to run:
    ID 1: AS_FedDAG_linear /
    ID 2: AS_FedDAG_linear /
    ID 3: AS_FedDAG_linear /
    ID 4: AS_FedDAG_linear /
    ID 5: AS_FedDAG_linear / simulation
    ID 6: AS_FedDAG_linear / simulation
    ID 7: AS_FedDAG_linear / sachs
    ID 8: AS_FedDAG_linear / child
    ID 9: AS_FedDAG_linear / alarm
    ID 10: GS_FedDAG / simulation
  Run [y/n] [y]:

When you confirm, each experiment prints its full configuration before launching:

==========================================================
  Strategy : AS_FedDAG_linear
  Data       : {'num-nodes': '10', 'num-samples': '4000'}
  Overrides  : {'num-rounds': '40'}
  Auto-scale : {'rho-init': 0.006, 'rho-scale': 10}
  Configuration:
    - dataset: simulation
    - num-nodes: 10
    - sem-type: linear
    - noise-scale: 1.0
    - w-min: 0.5
    - w-max: 2.0
    - data-seed: 2022
    - rho-init: 0.006
    - rho-scale: 10
    - rho-max: 1e+16
    - h-tol: 1e-05
    - h-thres: 0.9
    - init-iter: 8
    - al-every: 1
    - alpha-init: 0
    - max-retries: 10
    - it-fl: 300
    - lambda1: 0.05
    - lr: 0.003
    - graph-thres: 0.35
    - num-rounds: 40
    - strategy: AS_FedDAG_linear
    - num-samples: 4000
==========================================================

The Auto-scale block applies the paper's recommended values for rho-init and rho-scale based on the number of nodes — different scaling rules apply for synthetic vs. real data.

That's your Flower run.

---

Option 2 — Run an individual experiment from the terminal

Selecting option 2 opens a guided menu that walks you through each configuration choice:

Data Type:
    1. Simulated
    2. Benchmark (real data)
  Option [1]: 1

  Available SEM types: linear, mlp, gp, mim, gp-add
  sem_type [linear]: linear
  num_nodes (nodes in the graph) [10]: 10
  num_clients (federated clients) [5]: 2
  num_samples (samples per client) [6000]: 4000
  Non-IID (heterogeneous data across clients) [y/n] [n]: y

  Available strategies:
    1. AS_FedDAG_linear
    2. GS_FedDAG
    3. AS_FedDAG
  Strategy [2]: 2

  Configuration for GS_FedDAG / simulation:
    [Data]
      dataset = simulation
      num-nodes = 10
      num-clients = 2
      num-samples = 4000
      sem-type = linear
      noniid = 1
      w-min = 0.5
      w-max = 2.0
      noise-scale = 1.0
      data-seed = 2022
    [FL]
      num-rounds = 150
      lr = 0.03
      it-fl = 200
    [AL]
      rho-init = 0.1
      rho-max = 1e14
      rho-scale = 5.0
      h-tol = 1e-10
      h-thres = 0.25
      init-iter = 2
      al-every = 1
      max-retries = 10
    [Network]
      hidden-size = 16
      num-hidden-layers = 4
      temperature = 0.2
      w-init = 0.0
      l1-graph-penalty = 0.0005
    [Eval]
      graph-thres = 0.3
    [Others]
      al-mode = retry
      prune-cyclic = 1
      gumbel = 0
      it-fl-cap = 0
      partition-scheme = equal
      per-client-w = 0

The menu groups parameters into five categories — Data, FL, AL, Network (the MLP), Eval — so you only have to override what matters for your run. Defaults work for everything else. Any config key that doesn't fall into one of those five groups (e.g. al-mode, prune-cyclic, gumbel, it-fl-cap, partition-scheme, per-client-w) is shown under Others for visibility, but the interactive editor (_edit_hyperparams in run_experiments.py) currently only lets you change num-rounds, lr, it-fl, the AL parameters, graph-thres, dirichletAlpha, and lambda1 (linear only). To change anything else, edit experiments.csv directly or set it via pyproject.toml / --run-config.

After confirming, you are asked whether to persist the configuration:

 Save configuration in experiments.csv (exp_id=35) [y/n] [y]:
  Saved as exp_id=35
    Run experiment now [y/n] [y]:

Saving lets you re-run or look up the results later by ID.

---

More ways to run

The two menu options above are wrappers around Flower's native CLI. You can also call it directly:

flwr run . --stream

When using the raw Flower CLI you must edit pyproject.toml to change configuration, which is more verbose. The custom menu was built to avoid this.

⚠️ Windows PowerShell — UTF-8 encoding

Flower's rich output contains Unicode characters that Windows PowerShell may fail to render, producing a UnicodeEncodeError when flwr run is invoked via conda run. To avoid this, call the flwr executable directly from the environment and set the UTF-8 flag:

$env:PYTHONUTF8 = "1"
& "$env:USERPROFILE\miniconda3\envs\feddag\Scripts\flwr.exe" run . --stream

The run_experiments.py launcher handles this automatically — the workaround is only needed when calling flwr manually from PowerShell.

🍎 macOS — gRPC fork safety

On macOS the Flower SimulationEngine can deadlock when Ray forks worker processes. is a ready-made smoke test that exports the two required overrides before launching a short 10-round run:

./run_flwr_mac_os.sh

If you invoke flwr run manually on macOS, export GRPC_ENABLE_FORK_SUPPORT=1 and OBJC_DISABLE_INITIALIZE_FORK_SAFETY=YES yourself first.

Override config on the command line

# GS_FedDAG on sachs, 150 rounds
flwr run . --stream --run-config 'strategy="GS_FedDAG" dataset="sachs" num-rounds=150'

# AS_FedDAG_linear on simulation with d=15
flwr run . --stream --run-config 'strategy="AS_FedDAG_linear" dataset="simulation" \
                     num-nodes=15 num-rounds=100'

# AS_FedDAG on the child dataset
flwr run . --stream --run-config 'strategy="AS_FedDAG" dataset="child" num-rounds=150'

Custom Centralized Datasets and True DAGs

If you want to run FedDAG-Flower on a custom dataset located on an NFS or local drive, you can provide the path to a single .csv or .npy file. The framework will automatically load it and partition it among the clients (either equally or unequally based on the partition-scheme).

You can configure this via the run-config or via OS environment variables (useful for Docker/HPC deployments):

# 1. Provide the custom data for the clients to automatically partition:
export CENTRALIZED_DATA_PATH="/mnt/nfs/my_dataset.csv"
# (or via run-config: --run-config 'centralized-data-path="/mnt/nfs/my_dataset.csv"')

# 2. Provide the Ground Truth Adjacency Matrix for the Server to evaluate metrics:
export TRUE_DAG_PATH="/mnt/nfs/my_true_dag.npy"
# (or via run-config: --run-config 'true-dag-path="/mnt/nfs/my_true_dag.npy"')

flwr run .

Note on Data Alignment: Unlike the built-in bnlearn benchmarks which automatically reorder columns to match the DAG, when using a custom CENTRALIZED_DATA_PATH, the clients load the dataset exactly as-is. You must ensure that the columns in your .csv/.npy dataset are identically ordered to the variables in the adjacency matrix you provide to TRUE_DAG_PATH.

Real-World FL Deployment and True DAGs

In a true Cross-Silo Federated Learning deployment, you do not have a centralized dataset. Instead, each client hospital or institution holds its own completely separate, private dataset on its local hardware.

To run FedDAG-Flower in this mode across a real distributed SuperLink/SuperGrid, you should use Environment Variables exclusively, because they allow each local container to specify its own paths without modifying the global run-config:

1. On Client Containers (The Silos): Set FEDERATED_DATA_PATH to point to the local institution's private dataset.

# Client 1 (Machine A)
export FEDERATED_DATA_PATH="/hospital_A/secure_records/data.csv"
flwr run .

# Client 2 (Machine B)
export FEDERATED_DATA_PATH="/hospital_B/secure_records/data.csv"
flwr run .

2. On the Server Container (The Aggregator): The Server doesn't touch the data, but it needs the true DAG to compute evaluation metrics (SHD, TPR) and a persistent location to save the final models so they aren't lost when the container terminates.

export TRUE_DAG_PATH="/mnt/server_nfs/ground_truth_dag.npy"
export OUTPUT_DIR="/mnt/server_nfs/recovered_graphs/"
flwr run .

When FEDERATED_DATA_PATH is set, the client completely bypasses the centralized dataset partitioner and directly loads its siloed file, making it perfect for real-world production setups.

---

Benchmarking-Suite integration (bench_entrypoint.py)

The sibling Bechmarking-Suite orchestrator can run FedDAG against other federated causal-discovery algorithms (e.g. FedCSL) on identical data. It never touches FedDAG's internals — it talks to one contract exposed by at the repo root:

# Bechmarking-Suite invokes exactly this (you normally don't run it by hand):
python bench_entrypoint.py run \
    --exp-json <spec.json> --data-dir <pre-generated scenario> --output-dir <per-run dir>
# → writes <output-dir>/bench_result.json  {status, shd, tpr, fdr, fpr, f1, is_dag, runtime_s, extra:{...}}

python bench_entrypoint.py resources --exp-json <spec.json> [--data-dir <dir>]
# → prints {"cpus_per_task", "mem_per_cpu", "time_limit"} — FedDAG's own SLURM estimate

What this wrapper does (all FedDAG-specific logic lives here, not in the orchestrator):

• Translates the generic spec.json into FedDAG's run-config and reuses the real scripts/run_experiments.py internals (build_run_config_list, compute_metrics_from_npy, MetricsDAG, the flwr run command builder) — so the numbers are identical to a manual FedDAG run.
• Writes the metrics into a per-run output-dir (not the shared repo root), so several concurrent runs on the same repo copy don't clobber each other's W_final.npy — this is what makes the Hydra job pool safe.
• Declares FedDAG's own SLURM resources, weighting linear vs. nonlinear strategies (GS_FedDAG/AS_FedDAG cost ~1.8× a linear run), so each algorithm asks for what it actually needs.
• Skips combinations FedDAG doesn't support (e.g. noniid for a non-GS_FedDAG strategy) by writing status: skipped — the old adapter's supports() filter, now owned by FedDAG itself.

There's also — sourced by the orchestrator's single generic_job.pbs on Hydra to load the fedag-env and isolate each job's Flower SuperLink port/home. Adding FedDAG to the benchmark is a single entry in the orchestrator's configs/algorithms.yaml; see Bechmarking-Suite's docs/adding_a_new_algorithm.md.

Which hyperparameters does a benchmark run actually use?

bench_entrypoint.py run doesn't invent its own config — it builds the same layered config scripts/run_experiments.py would, via the real RE.build_run_config_list(strategy, data_overrides, hyper_overrides). The final value of every knob comes from stacking these, lowest priority first:

1. GLOBAL_DEFAULTS (fedag/config.py) — the absolute default of every knob.
2. auto_scale_defaults(num_nodes, dataset, strategy) — auto-tunes rho-init/rho-scale/lambda1/l1-graph-penalty from graph size and whether dataset is simulation or a real bnlearn network.
3. STRATEGY_DEFAULTS[strategy] — the tuned per-strategy values (see Augmented Lagrangian — paper vs. tuned defaults).
4. What Bechmarking-Suite sends from the generic ExperimentSpec (dataset, num-nodes, num-samples, seeds, noniid, ...).
5. spec.extra — anything you put in the feddag: {...} block of Bechmarking-Suite's experiments.yaml grid, at the highest priority of all. Any key already declared in pyproject.toml's [tool.flwr.app.config] can be overridden this way (Flower rejects undeclared keys with Key 'X' is not present in the main dictionary), e.g.:

   algorithms:
     - feddag:
         strategy: AS_FedDAG_linear
         lambda1: 0.02        # overrides STRATEGY_DEFAULTS for this run only
         num-rounds: 200

   The one exception is num-clients — it never travels through --run-config (Flower would reject it, since it isn't declared in pyproject.toml); it's spec.num_clients instead, sent via --federation-config num_supernodes=N.

Combinations FedDAG can't run at all are skipped, not attempted: a strategy outside {AS_FedDAG_linear, GS_FedDAG, AS_FedDAG}, or heterogeneity: noniid with anything other than GS_FedDAG (the only strategy that supports Non-IID). Those write bench_result.json with "status": "skipped" and an explanatory error.

Auto-tuned hyperparameters (Optuna) — layer 4.5

If you already ran for a given (strategy, dataset), bench_entrypoint.py run automatically finds and uses the winning config — it slots in as its own priority layer, between STRATEGY_DEFAULTS and whatever you explicitly set (explicit spec.extra always wins, so you can still override a single tuned knob for one run).

• Requires pip install -e ".[tune]" (optuna + pandas — an optional extra, not a core dependency). If missing, this layer is silently skipped and you fall back to STRATEGY_DEFAULTS, same as before this feature existed — nothing breaks.
• Study lookup: by default, derives the study name from the naming convention already used in studies/*.db — <prefix>_<dataset>, where AS_FedDAG_linear→asl, GS_FedDAG→gs, AS_FedDAG→as, and simulation→sim (so AS_FedDAG_linear + sachs → asl_sachs.db). Override with study-name in the grid's extra if your study doesn't follow that convention.
• No study found: only in an interactive local terminal (not Hydra — detected via SLURM_JOB_ID, and never hangs even where isatty() lies, e.g. Git Bash on Windows), it asks whether to run tune_optuna.py right now (tune-trials in extra, default 30) before the benchmark run. On Hydra, or if you say no, it just proceeds with STRATEGY_DEFAULTS — never blocks a batch job waiting for an answer that can't come.
• Disable entirely: use-study: false in the grid's extra.
• Traceability: every run's bench_result.json records extra.study = \{study_name, used, score, trial_number, n_trials_in_study, knobs_applied\} (or a reason if it wasn't used) — so master_results.csv reporting can tell which rows used a tuned config vs. plain defaults.

algorithms:
  - feddag:
      strategy: AS_FedDAG_linear   # -> looks for studies/asl_sachs.db automatically
  - feddag:
      strategy: GS_FedDAG
      study-name: gs_child_v2      # explicit override
      tune-trials: 50              # if missing and you say "yes" to tuning now
  - feddag:
      strategy: AS_FedDAG
      use-study: false             # always use STRATEGY_DEFAULTS, ignore any study

Known limitation: one study is reused across all num_clients values

The study name (_default_study_name in bench_entrypoint.py) is keyed on (strategy, dataset, num_nodes, heterogeneity) — it deliberately does not include num_clients. In practice this means: whichever config first triggers auto-tuning for a given combo runs Optuna once, using that config's num_clients as a representative value (see _build_tune_argv), and the resulting hyperparameters get reused as-is for every other num_clients you later benchmark with (2, 3, 5, 8, 10, 12, ...) — there is no re-tune per client count. This is a deliberate cost tradeoff (re-tuning per dataset×num_clients combination would multiply tuning cost 5-6x), not a validated claim that the optimal hyperparameters are client-count-invariant. They plausibly aren't: more clients means less local data per client (when splitting a fixed total), which interacts with it_fl (local steps/round) and num_rounds, and with partition_scheme: unequal also means more between-client size skew. Future work: either confirm empirically that performance is insensitive to num_clients for the tuned strategies (e.g. by comparing SHD/F1 across the client-count sweep already in the bench, experiments_feddag_test3_{2,3,5,8,12}c.yaml), or extend the study-name convention to key on num_clients too (accepting the extra tuning cost) if it turns out not to be.

Known limitation: very large graphs (d ≳ 441) don't converge — not viable yet

Tested against pigs (d=441) and link (d=724) on 2026-07-13 (90 trials total across the 3 strategies, on Hydra, after ruling out every infra bug we could find — stale W_final.npy from concurrent runs, it-fl-cap making a round exceed flwr run's --stream refresh window): every single trial still dies mid-round-1, before the server ever aggregates once (W_final.npy is never written). About half show an explicit Ray failure, The current node timed out during startup ... the GCS has become overloaded; the rest die with no traceback at all (consistent with an OS-level OOM-kill, which doesn't leave a Python-visible error) — both point to Ray/resource exhaustion from repeatedly tearing down and restarting a federated-simulation cluster at this scale, not to anything in FedDAG's own AL/training logic (confirmed healthy up to d=37 — hundreds of real rounds per trial, early-stops working as designed). Do not add pigs, link, or any comparably large network to a benchmark grid yet — every run will burn its full time budget and end in status: error/no metrics, not a slow-but-valid result. Root cause (Ray GCS instability under repeated sequential restarts, or genuine memory exhaustion — never isolated with seff/Ray's own session logs, which hpc/hydra_env.pbs's cleanup trap deletes on job exit) is tracked as future work, not blocking the rest of the benchmark.

Known limitation: the nonlinear strategies collapse at d=40 in this port (the paper's don't)

Status (2026-07-29): open. AS_FedDAG_linear scales fine to d=40; both nonlinear strategies (GS_FedDAG, AS_FedDAG) collapse there — typically to the empty graph, deterministically (SHD = #true-edges, F1 = 0.00, std = 0.00 across 10 seeds). The original authors' code reaches SHD 30–36 / TPR 0.65–0.74 at the same size, so this is a port-specific gap, not a limit of the method.

Measured results (10 seeds each, Hydra, --hydra-exclusive)

Simulation, ER2, 10 clients × 600 obs. d=10/20 are healthy; d=40 is where it breaks:

Data	Strategy	d=10	d=20	d=40	Paper @d=40
linear iid	AS_FedDAG_linear	SHD 1.0 / F1 0.97	SHD 1.4 / F1 0.98	SHD 3.0 / F1 0.98 ✅	—
linear iid	AS_FedDAG	SHD 19.4 / F1 0.42	SHD 17.9 / F1 0.71	SHD 74.0 / F1 0.00	—
linear iid	GS_FedDAG	SHD 19.0 / F1 0.57	SHD 31.0 / F1 0.53	SHD 74.7 / F1 0.00	—
GP iid	AS_FedDAG	SHD 9.9 / TPR 0.95	—	SHD 70.3 / TPR 0.06	SHD 30.0 / TPR 0.74
GP iid	GS_FedDAG	SHD 6.7 / TPR 0.98	—	SHD 74.0 / TPR 0.00	SHD 36.5 / TPR 0.65
ANM noniid	AS_FedDAG	SHD 7.5 / TPR 0.81	—	SHD 84.2 / TPR 0.73	SHD 35.9 / TPR 0.84
ANM noniid	GS_FedDAG	SHD 6.3 / TPR 0.70	—	SHD 74.1 / TPR 0.00	SHD 24.3 / TPR 0.86

Two hypotheses were tested and both failed, which is what pointed at the real cause:

1. "The per-node MLPs are starved of local steps" — raised it-fl-cap from 6000/d (=150 at d=40) to 1000 (6.7×) and re-tuned from scratch. No improvement; AS_FedDAG d=40 linear even went from F1 0.64 → 0.00. Inspecting all 60 trials of the re-tune: no trial reached a usable F1 (range 0.00–0.18, SHD 73–572 against ~74 true edges — essentially random graphs).
2. "Just use the paper's ρ schedule" — pinned rho-init/rho-scale to the paper's exact per-size values (Table 8: 1e-11 / 120 at d=40) instead of searching them. This regressed d=20 (AS_FedDAG F1 0.71 → 0.56, GS_FedDAG 0.53 → 0.49) and did not fix d=40 (best tuning scores stayed negative for 4 of 8 studies). The paper's values were tuned as a package with everything else fixed exactly as they had it — they don't transfer one-at-a-time into this port.

Root-cause analysis: where this port diverges from the original

Comparing against the authors' TensorFlow code (FedDAG/, trainers/al_trainer.py, models/fed_dag.py) and the paper (TMLR 01/2023, App. B.6–B.7, Tables 7–8):

(a) Several of the paper's values are literally outside our search ranges — Optuna could not have found them even by chance. These are still the live ranges in tune_optuna.py (see the proposed fix):

Parameter	Paper / original code	Search range in tune_optuna.py	Reachable?
lr	3e-2	[1e-4, 5e-3]	✗ 6× above the ceiling
graph-thres	0.5 (§4.3, "median value 0.5")	[0.05, 0.25]	✗ 2× above the ceiling
h-thres (γ)	0.25 (Table 7)	[0.3, 0.85]	✗ below the floor
l1-graph-penalty	0.01 (Table 7)	[2e-4, 5e-3]	✗ 2× above the ceiling
rho-init @d=40	1e-11 (Table 8)	[1e-3, 0.2]	✗ 8 orders below the floor
rho-scale (β) @d=40	120 (Table 8)	2.0, fixed	✗ not searched at all

A low graph-thres is especially damaging at scale: with a weakly-learned W, a 0.05–0.25 cut admits hundreds of spurious edges (the SHD 572 trials above), while the paper's 0.5 keeps only confident ones.

(b) The ALM outer-loop cadence is structurally different. The original solves each ALM sub-problem with a fixed 1000 gradient steps (iter_step=1000), aggregating every it_fl=200 steps → one ALM update per 5 aggregations, and caps the whole thing at max_iter=25 outer iterations. In this port a Flower round is one aggregation, and al-every was fixed at 1 → one ALM update per aggregation (5× more often), over num-rounds=150 rounds (6× more updates than the paper's 25). Since each update does alpha += rho * h and escalates rho multiplicatively, the penalty coefficients ramp up far faster per unit of optimisation than in the original — which is exactly the mechanism that crushes W to the empty graph (h≈0 satisfies the constraint trivially) before the mechanism networks learn anything. This also explains why hypothesis (1) failed: the problem isn't how many local steps per round, it's how often the penalty escalates relative to them.

(c) it_fl is size-dependent here, fixed in the original. The paper uses it_fl=200 for every d; this port derives it-fl-cap from 6000/d, so the optimisation-per-ALM-update budget silently shrinks as graphs grow — d=40 got 150 steps where d=10 got 600.

(d) Minor: hidden-size is 8 for AS_FedDAG here (halved to keep Flower message sizes down) vs 16 in the paper for both strategies.

Why AS_FedDAG_linear is immune: it has no per-node MLPs and its own separate search space (with rho-init/rho-scale genuinely tunable), so none of (a)–(d) applies — it hits F1 0.98 at d=40.

Proposed fix — NOT yet in the code

The intended response is to widen sample_config()'s nonlinear branch so the search space contains every paper value above, instead of either the old too-narrow ranges or hard-pinned constants:

Parameter	Current range in tune_optuna.py	Proposed range
lr	[1e-4, 5e-3]	[1e-4, 5e-2]
graph-thres	[0.05, 0.25]	[0.05, 0.6]
h-thres	[0.3, 0.85]	[0.2, 0.85]
l1-graph-penalty	[2e-4, 5e-3]	[2e-4, 2e-2]
rho-init	[1e-3, 0.2]	[1e-12, 0.2]
rho-scale	2.0, fixed	[1.5, 150]
al-every	1, fixed	{1, 2, 5} (5 reproduces the paper's one-update-per-1000-steps cadence)

That would let the search reach both the region that already works at d ≤ 20 and the region the paper uses at d=40, and decide per (strategy, dataset, d, heterogeneity) rather than committing to one guess.

Status: unimplemented. still ships the "current" column above — verified against the working tree. The only knobs that were added to the nonlinear branch since this analysis are rho-runaway-factor and al-mode (see Known issues #13), which are unrelated to the d=40 collapse. Neither the widening nor the re-tune has happened; the collapse is not fixed.

Known gap: Hydra tune study names don't match the auto-lookup convention

The <prefix>_<dataset> convention above (asl_sachs, gs_child, ...) is what _default_study_name() in bench_entrypoint.py actually looks for, but the studies produced by the hpc/hydra_optuna.pbs sweeps in this repo were named <prefix>_<dataset><num_nodes>_hydra (e.g. lin_sachs11_hydra, not asl_sachs) — lin instead of the documented asl prefix, plus a node-count and _hydra suffix the lookup doesn't expect. A benchmark run against one of these datasets won't error — it silently falls back to STRATEGY_DEFAULTS and records extra.study.reason in bench_result.json explaining why, so nothing crashes, but the tuned config never gets applied unless you either (a) rename/copy the .db (and matching _trials.csv) to the canonical name FedDAG expects, or (b) set study-name explicitly per algorithm block — though a single study-name override can't vary per dataset within one feddag: {...} grid entry, so for a multi-dataset sweep (a) is the only option that actually works. Check extra.study in bench_result.json (or master_results.csv) to confirm a given row really used the tuned config before trusting its numbers over STRATEGY_DEFAULTS.

---

Expected Output

The last few lines of a successful run look like this:

[round 73] h=2.41e-07  rho=1.54e+06  alpha=3.21e+02
[round 74] h=8.30e-09  rho=3.08e+06  alpha=3.52e+02
[round 75] h=2.86e-09  rho=3.08e+06  alpha=3.59e+02  ← converged (h ≤ 1e-8)
SHD=2  TPR=0.90  FDR=0.10  F1=0.92  G-score=0.85
W_final.npy saved.

Results CSV

In the results/ folder you will find results.csv:

Each row contains the model ID (used to look up the saved W_final.npy), the experiment ID, basic configuration info, and all metrics: is_dag, SHD, TPR, FDR, F1, G-score, etc. These metrics tell you how well the recovered graph matches the ground truth — see the Metrics section for definitions.

---

Visualizing Results

The project includes visualize_W.py for inspecting a trained model visually. To generate a .png for a specific model:

python scripts/visualize_W.py --id 0001

The --id parameter is the model ID stored in results.csv.

You will get terminal output like this:

========================================
is_dag  : NO
SHD     : 24   (lower is better)
TPR     : 0.800  (higher is better)
FDR     : 0.600  (lower is better)
FPR     : 0.343  (lower is better)
True edges      : 20
Predicted edges : 40
========================================

A non-DAG result is a bad sign in general, but acceptable for short test runs (e.g., 30 rounds). The acyclicity penalty needs enough rounds to fully converge.

The figure contains seven panels:

1. Heatmap of W_final (top-left) — colored matrix showing estimated edge weights between nodes. Positive weights in red, negative in blue. Row $i$ → column $j$ means an edge from node $i$ to node $j$. Shows the magnitude and sign of the learned causal relationships.
2. Heatmap of the true graph (bottom-left) — the ground truth DAG B_true, for comparison. Shows what the algorithm should have learned.
3. Estimated graph (top-middle) — graph visualization after thresholding (|W| > graph-thres). Green edges = positive weights, red edges = negative weights.
4. True graph (bottom-middle) — the actual DAG structure that was expected.
5. Comparison graph — overlays predicted vs. true edges. Green = true positives (correctly recovered), red = false positives (predicted but not in truth), orange = false negatives (missed). Quick visual check of where the model is failing.
6. Metrics panel — is_dag, SHD, TPR, FDR, FPR, edge counts. See Metrics for definitions.
7. Histogram — distribution of |W| values (excluding zeros), with the threshold line marked.

Recovery quality depends heavily on data quality and hyperparameter choice. The defaults shipped here were the best found during tuning — if you find better ones, please share!

For named-node datasets (e.g., sachs), the visualization shows variable names instead of indices.

---

How It Works

SERVER                                  CLIENTS (k = 1..K)
_________                               _________________
Initialize W_global = 0

configure_train(round, W_global)  __→   train(msg)
  send W + (rho, alpha, tau)              receive W_global
                                          load W into local model
                                          run it-fl steps with AL loss:
                                            L = MSE + λ₁‖W‖₁ + α·h(W) + ½ρ·h(W)²
                                          return W_local + h(W_local)

aggregate_train(replies)         ←__    reply (W_local, n_examples, h)
  W_global = Σ nₖ·Wₖ / Σ nₖ              (weighted FedAvg)
  h_new = compute_h(W_global)
  _al_update(self, h_new):
    if h ≤ tol → CONVERGED
    if h > 0.25·h_prev → escalate ρ, RETRY (don't update α)
    else → α += ρ·h, h_prev = h

  → loop back to configure_train with new (rho, alpha)

The flow in plain words:

1. The server broadcasts the current global $W$ along with the dual variables $(\rho, \alpha)$ and the temperature $\tau$ (nonlinear strategies only).
2. Each client loads $W$ into its local model and runs it-fl Adam steps minimizing the augmented Lagrangian loss.
3. Clients return their updated $W$ to the server.
4. The server aggregates via FedAvg weighted by the number of local examples.
5. The server checks convergence:
   • If $h(W) \le h_\{\text\{tol\}\}$, the run converged and stops early.
   • If $h$ did not decrease enough, $\rho$ is escalated and the round is retried.
   • Otherwise, $\alpha$ is updated and the loop continues.

---

Strategies

AS_FedDAG_linear

Use case	Linear SEM data; fastest to converge; baseline for understanding the others
Model	Linear NOTEARS
Shared	Only $W$
Local persistence	Adam state via context.state
Datasets	Any — simulation or any bnlearn/bnrepository network (sachs, child, alarm, asia, insurance, …)
SEM types	linear with Gaussian / Exponential / Gumbel / Uniform / Logistic noise
Non-IID?	IID by default — all clients use the same SEM
Dirichlet?	Not used (IID only)

Recommended hyperparameters:

strategy      = "AS_FedDAG_linear"
num-rounds    = 100
it-fl         = 300      # local gradient steps per round
lambda1       = 0.01     # L1 sparsity on W
lr            = 0.003    # Adam learning rate
graph-thres   = 0.35     # binarization threshold for B_est
sem-type      = "linear"
rho-init      = 1.0
rho-scale     = 10

These are the values fedag/config.py's STRATEGY_DEFAULTS["AS_FedDAG_linear"] applies automatically — you only need to set strategy = "AS_FedDAG_linear" and they kick in unless overridden.

Both rho-init and rho-scale need to grow with the number of nodes. The paper recommends different values for synthetic vs. real data; this is currently handled by the auto-scaling logic in run_experiments.py.

Round-level behavior:

Round t:
  Server  →  broadcasts (W_global, rho, alpha) to all clients
  Clients →  run 300 Adam steps minimizing:
               loss = MSE(X, X@W) + λ·‖W‖₁ + α·h(W) + ½ρ·h(W)²
             where h(W) = trace(expm(W⊙W)) − d   (NOTEARS acyclicity)
  Server  →  W_global = Σ(Wᵢ·nᵢ) / Σnᵢ           (weighted FedAvg)
             if h > 0.25·h_prev: rho *= rho-scale
             alpha += rho · h
             stop if h ≤ h-tol

Final: B_est = (|W_final| ≥ graph-thres)

---

GS_FedDAG — Nonlinear, Gradient-Shared

Use case	Nonlinear data; the only strategy that supports Non-IID and heterogeneous clients
Model	Nonlinear MaskedNN
Shared	Only $W$ — local MLPs stay client-side
Local persistence	MLPs persist between rounds via context.state["gs-nets"]
Datasets	simulation (best fit) or any benchmark
SEM types	mlp (default), gp, mim, gp-add
Non-IID?	Yes — clients can have different SEM types (e.g. client 1 linear, client 2 quadratic)
Dirichlet?	Yes, when noniid=1. Controls heterogeneity: α=0.1 → high heterogeneity; α=10 → near-IID

Recommended hyperparameters:

strategy           = "GS_FedDAG"
num-rounds         = 150
it-fl              = 200
lr                 = 0.01     # tuned; paper uses 3e-2
hidden-size        = 16       # MLP hidden units per layer
num-hidden-layers  = 4        # MLP depth per node
temperature        = 0.2      # τ: sigmoid temperature for soft graph masking
l1-graph-penalty   = 0.0001   # tuned; paper uses 2e-3
graph-thres        = 0.5
rho-init           = 0.1
rho-scale          = 2.0
w-init             = 0.5

These are STRATEGY_DEFAULTS["GS_FedDAG"] in fedag/config.py — applied automatically when strategy = "GS_FedDAG".

Round-level behavior:

Round t:
  Server  →  broadcasts (W_logits, rho, alpha, tau) to all clients
  Clients →  W_prob = sigmoid(W / tau)               (soft adjacency)
             train MLP_j(X * W_prob[:,j]) per node j   (local NNs persist!)
             minimize: MSE + λ·‖W_prob‖₁ + α·h + ½ρ·h²
             h(W) = trace(expm(sigmoid(W/tau)²)) − d
             return only W_logits (MLP weights stay local)
  Server  →  W_global = Σ(Wᵢ·nᵢ) / Σnᵢ              (FedAvg on W only)
             AL update same as above

Note: Each client's MLP weights are not shared — they persist between rounds via context.state["gs-nets"] (Flower's per-node state, so they survive regardless of which worker process runs the client). This lets clients learn heterogeneous local causal mechanisms while still agreeing on a single global graph.

---

AS_FedDAG — Nonlinear, All-Shared

Use case	Nonlinear version of AS_FedDAG_linear. Same MaskedNN as GS_FedDAG but with a different update scheme
Model	Nonlinear MaskedNN
Shared	$W$ + all MLP weights
Local persistence	Model rebuilt from server params every round
Datasets	Any
SEM types	mlp (default), gp, mim, gp-add
Non-IID?	IID only
Dirichlet?	Not used

Recommended hyperparameters:

strategy           = "AS_FedDAG"
num-rounds         = 150
it-fl              = 200
lr                 = 0.05
hidden-size        = 8        # 16 can overload the Flower message size
num-hidden-layers  = 2
temperature        = 0.2
l1-graph-penalty   = 0.0005
graph-thres        = 0.5
rho-init           = 0.1
rho-scale          = 1.25     # best tuned value among {0.1, 1.25, 1.5, 2.0, 10.0}
lambda1            = 0.01

These are STRATEGY_DEFAULTS["AS_FedDAG"] in fedag/config.py — applied automatically when strategy = "AS_FedDAG".

---

Choosing a strategy

Scenario	Strategy
Quick baseline, synthetic linear data	AS_FedDAG_linear
Real-world federated data, heterogeneous clients	GS_FedDAG
IID data, want best possible recovery with a nonlinear model	AS_FedDAG
Large graphs (d > 20), need speed	AS_FedDAG_linear
Benchmark comparison with the paper	Match the paper's strategy per experiment

---

Paper vs. this implementation

This repo is a port of the FedDAG algorithm (Gao et al., TMLR 2023, Algorithms 1–2) to the Flower 1.2x Strategy API. The objective, the two-level model, and the optimization scheme are faithful; the differences are mostly about how the paper's nested loops map onto Flower's round model, plus a few empirically-driven choices that improved recovery on our benchmarks. References below are to the paper's equations/algorithms.

What matches the paper (similarities)

Aspect	Paper	This implementation
Objective	Maximize Σₖ Sᶜᵏ subject to h(U)=0 (Eq. 6)	Same — augmented-Lagrangian sub-problem (Eq. 8) solved per round
Score function	−1/(2nₖ)·Σ‖Dᵢⱼ − Φⱼ(g_τ(U)∘Dᵢ)‖² − λ‖g_τ(U)‖₁ (Eq. 7)	Same: 0.5/n·Σ residual² + l1-graph-penalty·Σ σ(W/τ) (compute_loss_nonlinear)
Two-level model	GSL part U (graph) + MA part Φ (d sub-networks)	GS_FedDAG / AS_FedDAG: shared W logits + per-node MLPs (MaskedNNModel)
Three variants	GS-FedDAG (share U only), AS-FedDAG (share U+Φ), linear version (share the adjacency matrix directly, Sec 5.1.1)	GS_FedDAG, AS_FedDAG, AS_FedDAG_linear — exactly these three
Acyclicity	h(U) = Tr[e^{g_τ(U)}] − d (Eq. 5); linear: Tr[e^{W∘W}] − d (NOTEARS, Eq. 3)	Same NOTEARS constraint via torch.linalg.matrix_exp (metrics.compute_h / h_func) — see the one deviation on the nonlinear square below
Soft mask	Gumbel-Sigmoid g_τ(U) (Eq. 4)	sigmoid(W/τ); the stochastic Gumbel-Sigmoid is available via gumbel=1
Optimizer	ADAM on the sub-problem (no closed form)	Adam, per client, it-fl local steps per round
Aggregation (homogeneous)	FedAvg — average the collected U (and Φ for AS)	FedAvg over W (GS) or all params (AS)
Thresholding	Median 0.5 on the mask, then iteratively cut the weakest edge until acyclic (Sec 4.3)	graph-thres + prune-cyclic=1 (same weakest-edge cycle cutting)
Coefficient update rule	α ← α + ρ·E[h]; ρ ← β·ρ iff E[h] > γ·E[h_prev] (Alg. 1, lines 9–14)	Exactly this when al-mode="paper"; β=rho-scale, γ=h-thres
Stopping	while t ≤ itmax and h ≥ htol and ρ ≤ ρmax (Alg. 1, line 5)	Early-stop on h ≤ h-tol and on the ρ ceiling (num-rounds is the itmax bound)
Guarantee	Stationary point, not the global optimum / true DAG	Same — no stronger claim made

Where it deviates (differences)

Aspect	Paper	This implementation	Why
Loop structure	Nested: itmax outer sub-problems, each an itin-step inner loop (Alg. 2) that aggregates every itfl steps	Flattened to Flower rounds: one round = it-fl local steps + one aggregation + one AL update (gated by init-iter warm-up and al-every)	Flower's round model is one aggregation per round; the paper's itin/itfl nesting collapses to "rounds"
AL coefficient update (default)	α updated every sub-problem step, ρ escalates, no rollback (Alg. 1)	Default al-mode="retry": NOTEARS-style inner loop — on insufficient h decrease, escalate ρ and re-train from the last accepted params without touching α	"retry" was more stable in tuning; the verbatim paper rule is still available as al-mode="paper"
Nonlinear acyclicity	Tr[e^{g_τ(U)}] − d (no square)	Client-side loss: Tr[e^{g_τ(U) ⊙ g_τ(U)}] − d — element-wise square. Server-side convergence check: unsquared, i.e. the paper's form — see Known issues #17	A/B tested: unsquared over-weights the training penalty and collapses the graph to empty on the benchmarks (sachs F1 ≈ 0.10 unsquared vs ≈ 0.68 squared). The two sides disagreeing is not intentional
Soft mask (default)	Stochastic Gumbel-Sigmoid every step (Eq. 4)	Deterministic sigmoid(W/τ) by default (gumbel=0)	Determinism made tuning/eval more reproducible; set gumbel=1 for the paper's exact stochastic mask
Client selection	Randomly select r of m clients each aggregation (Alg. 2, line 11)	Deterministic: first fraction_train·m clients. fraction_train is a constructor argument of FedDAGBaseStrategy, hardcoded to 1.0 — it is not a config key, so partial participation currently requires editing server_app.py	Identical to the paper when r = m, which is what every run in this repo does
FedAvg weighting	Plain average of selected U	Weighted by each client's num-examples	Identical to a plain average when client sizes are equal (our default); generalizes correctly when they aren't
MA-part persistence	Implicit per-client Φ	GS persists local MLPs across rounds via Flower's context.state["gs-nets"] (distribution-safe); AS rebuilds them from the shared params each round	Maps the paper's "local Φ" onto Flower's per-node state
Data / SEM	Gaussian ANM, ER/SF graphs, GP/MLP/MIM/GP-add mechanisms, 10 clients × 600 obs	Same simulation menu plus real bnlearn/bnrepository benchmarks (sachs, child, alarm, …)	Adds real-data evaluation on top of the paper's synthetic setup
Engineering add-ons	—	Early-stop when ρ is stuck at its ceiling (stuck-patience + rho-runaway-factor), infra-error retry, it-fl-cap, model-seed reproducibility	Practical robustness for HPC/Optuna sweeps; none change the algorithm's math

Net: run with al-mode="paper", gumbel=1, and partition-scheme="equal" (full participation is already the only mode) to get the closest match to Algorithm 1. The defaults (al-mode="retry", deterministic sigmoid, squared nonlinear h) trade verbatim fidelity for the recovery quality and stability we measured during tuning.

---

Hyperparameters Reference

Master parameter table

Defaults below are the literal GLOBAL_DEFAULTS values in fedag/config.py — the library-wide fallback used before any per-strategy or pyproject.toml override is applied. STRATEGY_DEFAULTS overrides several of these per strategy (see the two detailed tables right after this one for the actual tuned values used by each strategy).

Parameter	Default (GLOBAL_DEFAULTS)	Source	Controls
strategy	AS_FedDAG_linear	config.py	Which model and what gets federated
Data / graph
dataset	simulation	config.py	simulation, or any bnlearn/bnrepository network name
num-nodes (d)	10	config.py	DAG dimension (ignored when dataset is a benchmark — d comes from the network)
graph-degree	2.0	config.py	Expected degree of the generated Erdős–Rényi DAG (simulation only)
sem-type	linear	config.py	linear | mlp | mim | gp | gp-add | quadratic (nonlinear for GS/AS via STRATEGY_DEFAULTS)
nonlinear-type	mlp	config.py	Which nonlinear mechanism, when sem-type is nonlinear
noise-scale	1.0	config.py	σ of the additive noise term
w-min / w-max	0.5 / 2.0	config.py	Range the true edge weights are drawn from (simulation only)
graph-seed	42	config.py	Seeds the DAG topology
data-seed	2022	config.py	Seeds the SEM data draw
num-clients	2	config.py	Federation size. ⚠️ Not a Flower run-config key — it maps to --federation-config num_supernodes=N
num-samples	10000	config.py	Rows per client (see partition-scheme)
num-rounds	50 (100–150 per strategy — see below)	config.py	Outer Flower rounds
it-fl	300 (lin) / 200 (nl)	config.py	Local gradient steps per round
lr	0.003 (lin) / 0.01–0.05 (nl, tuned — see below)	config.py	Local Adam learning rate
lambda1	0.05 (0.01 tuned per strategy)	config.py	L1 sparsity on W (linear)
l1-graph-penalty	5e-4 (1e-4–5e-4 tuned per strategy)	config.py	L1 sparsity on sigmoid(W) (nonlinear)
graph-thres	0.35 (0.5 for GS/AS, tuned)	config.py	Threshold for binarizing W_final
temperature (τ)	0.2	config.py	Sigmoid temperature for soft masks. Lower → more binary
hidden-size	16 (8 for AS_FedDAG, tuned)	config.py	MLP hidden units per layer
num-hidden-layers	4 (2 for AS_FedDAG, tuned)	config.py	MLP depth
w-init	0.0 (0.5 for GS_FedDAG, tuned)	config.py	Initial value of every W logit for the nonlinear strategies. > 0 avoids the sigmoid gradient dead zone
Augmented Lagrangian
rho-init	6e-3 (1.0 lin / 0.1 nl, tuned per strategy)	config.py	Initial penalty on h(W)
rho-max	1e16 (1e14 for GS/AS)	config.py	Cap before giving up
rho-scale	10 (2.0 GS / 1.25 AS, tuned)	config.py	Multiplicative factor when escalating ρ — the most sensitive hyperparameter
alpha-init	0.0	config.py	Initial dual variable
h-tol	1e-5 (1e-6 tuned per strategy)	config.py	Convergence threshold
h-thres	0.9 (0.25 for GS/AS, tuned)	config.py	Trigger ρ escalation if h_new > h-thres·h_prev
init-iter	8 (2 GS / 10 AS, tuned)	config.py	Warm-up rounds before AL updates start
al-every	1	config.py	Run AL update every N rounds
max-retries	10	config.py	Cap on retries before advancing
stuck-patience	2	config.py	al-mode="retry" only: stop training early after this many consecutive AL give-ups with rho already at its ceiling (further retries can't escalate rho, so more rounds can't help)
rho-runaway-factor	1e4	config.py	The real ceiling for ρ is min(rho-max, rho-init × rho-runaway-factor), not rho-max alone. rho-max (1e14–1e16) is so far above a typical rho-init (1e-3…1) that a runaway trajectory can climb 5–10 orders of magnitude without ever reaching it — so stuck-patience almost never fires before the full num-rounds budget is burned. This gives it a reachable ceiling for runaway trials, while leaving healthy escalation untouched. Set 0 to disable (fall back to rho-max only)
al-mode	retry	config.py	"retry" = NOTEARS-style inner loop with rollback (the tuned defaults assume this mode) — "paper" = Algorithm 1 of Gao et al. verbatim (alpha updates every AL step, no retry/rollback, training stops once ρ would exceed its ceiling)
it-fl-cap	0	config.py	Hard cap on local gradient steps per round (0 = disabled). The client honours it for both model families — linear (local_train_step) and nonlinear (local_train_step_nonlinear); a WARNING is logged whenever it clips. Set on HPC to bound the O(d³) expm cost per step. The orchestration scripts each derive their own value from num-nodes (see it-fl-cap: who sets it and to what)
model-seed	42	config.py	Seeds the random per-node MLP weight init for GS_FedDAG / AS_FedDAG (torch.manual_seed). AS_FedDAG seeds once server-side (the shared initial arrays broadcast to every client); GS_FedDAG seeds per client (model-seed + partition-id, since its local nets are never federated). The linear model has no random init (W starts at 0), so this doesn't apply to AS_FedDAG_linear. Vary it across re-runs of the same hyperparameters to tell whether a trial's score reflects the config or just a lucky init — see
gumbel	0	config.py	1 = stochastic Gumbel-Sigmoid mask (g_τ in the paper) instead of a plain sigmoid for the nonlinear strategies; with gumbel=1, consider keeping w-init = 0.0 since the added noise breaks the dead zone
prune-cyclic	1	config.py	Paper Sec. 4.3 post-processing: if the thresholded B_est still has cycles, iteratively cut the weakest edge until it is a DAG. Set 0 to disable. n_cut > 0 in the final log is a sign the acyclicity constraint had not fully converged
Partitioning
noniid	0	config.py	Enable Non-IID (Dirichlet-based)
dirichletAlpha	None (unset)	config.py	Dirichlet concentration
partition-scheme	equal	config.py	"equal" = every client gets num-samples rows (default) — "unequal" = arithmetic split following the FedCSL paper's Eq. 11, so clients hold increasingly more data
per-client-w	0	config.py	0 = all clients share the same weighted DAG when noniid=0 (correct IID) — 1 = legacy behavior where each client gets its own W even in "IID" mode
I/O
data-dir	""	pyproject.toml	When set, partition.py loads client_<id>.npy from this folder instead of generating data — how Bechmarking-Suite feeds identical data to every algorithm. Empty = generate internally. Not in GLOBAL_DEFAULTS; read via run_config.get("data-dir")
output-dir	.	config.py	Where W_final.npy is written. run_experiments.py and bench_entrypoint.py override it with a per-run absolute path so concurrent runs don't clobber each other

pyproject.toml's [tool.flwr.app.config] is a separate, higher-priority override layer on top of everything above — see for the exact precedence order (run_config > STRATEGY_DEFAULTS[strategy] > GLOBAL_DEFAULTS).

it-fl-cap: who sets it and to what

fedag/ itself never computes this value — only reads it-fl-cap from the resolved run config and passes it down to local_train_step / local_train_step_nonlinear, which clip and log a WARNING. There is no shared helper: each orchestration script derives its own value from num-nodes, and they do not all agree.

Caller	Value it sets	Scope
pyproject.toml / GLOBAL_DEFAULTS	0 — disabled	a plain flwr run is uncapped
	max(10, 6000 // d)	nonlinear only (GS_FedDAG / AS_FedDAG); AS_FedDAG_linear is left uncapped on purpose — it is cheap per step and capping it would cut steps alarm (d=37) actually uses
	max(100, 6000 // d), via setdefault — an explicit it-fl-cap in spec.extra wins	all strategies
	max(100, 6000 // d)	all strategies
	max(100, 6000 // d) as the upper bound of the sampled it-fl, not as the cap key	all strategies

The rationale for 6000 // d is that the acyclicity term costs O(d³) per step, so bigger graphs get fewer steps per round and no single round overruns flwr run --stream's 60s refresh window. Typical values: d=10 → 600, d=20 → 300, d=37 → 162.

⚠️ The 10 vs 100 floor is a real inconsistency. tune_optuna.py lowered its floor to 10 so d ≥ 441 (pigs/link) would get ~13 steps instead of 100 and stay inside the stream window. The benchmark and verification paths kept the 100 floor. The two therefore disagree for any d > 60 (6000 // 60 = 100): a tuned config can be re-run under a different local-step budget than it was tuned with. It has no effect on the datasets currently in use (d ≤ 37), but it must be reconciled before any large-graph work resumes.

Augmented Lagrangian — paper vs. tuned defaults

These are the values applied by STRATEGY_DEFAULTS in fedag/config.py when you only set strategy. The "tuned" columns are now split because GS_FedDAG and AS_FedDAG have diverged during tuning — they no longer share identical values.

Parameter	Config key	Description	Linear (tuned)	Nonlinear (paper)	GS_FedDAG (tuned)	AS_FedDAG (tuned)
$\rho_0$	rho-init	Initial penalty weight	1.0	1e-3	0.1	0.1
$\rho_\{\max\}$	rho-max	Maximum ρ before giving up retries	1e16	1e14	1e14	1e14
$c$	rho-scale	Multiplicative factor when h doesn't decrease	10	5	2.0	1.25
$\alpha_0$	alpha-init	Initial dual variable	0.0	0.0	0.0	0.0
$h_\{\text\{tol\}\}$	h-tol	Convergence threshold	1e-6	1e-10	1e-6	1e-6
$\gamma$	h-thres	ρ-escalation trigger	0.9	0.25	0.25	0.25
warm-up	init-iter	Rounds before AL penalty activates	8	3	2	10
AL cadence	al-every	AL update every N rounds	1	1	1	1
max retries	max-retries	Max ρ escalations per AL step	10	∞	10	10
AL coefficient update	al-mode	"retry" (rollback inner loop) or "paper" (Algorithm 1 verbatim)	retry	paper	retry	retry

Local training — paper vs. tuned defaults

Parameter	Config key	Linear (tuned)	Nonlinear (paper)	GS_FedDAG (tuned)	AS_FedDAG (tuned)
Local gradient steps	it-fl	300	200	200	200
Learning rate	lr	0.003	3e-2	0.01	0.05
L1 on W	lambda1	0.01	—	—	0.01
L1 on sigmoid(W/τ)	l1-graph-penalty	—	2e-3	1e-4	5e-4
Binarization threshold	graph-thres	0.35	0.5	0.5	0.5
MLP hidden units	hidden-size	—	16	16	8 (16 can overload the Flower message size)
MLP depth	num-hidden-layers	—	4	4	2
Initial W logits	w-init	—	0.0	0.5	0.0

Bold = differs from the paper, or differs between GS_FedDAG and AS_FedDAG. These values live in fedag/config.py's STRATEGY_DEFAULTS — edit them there (or override per-run via the CSV / --run-config) rather than in pyproject.toml, since pyproject.toml is itself just one more override layer on top of these defaults.

Recommended num-rounds per dataset

Flower rounds ≈ paper's max_iter × iter_step / it_fl + retries.

Dataset	d (nodes)	Edges	Paper max_iter	Recommended num-rounds
simulation (d=5)	5	~5	25	100
simulation (d=10)	10	~20	25	100–150
sachs	11	17	25	150
child	20	25	25	150
alarm	37	46	25	200

Why more rounds than the paper? The paper's inner while loop re-trains within a single "iteration" when ρ is escalated. In Flower, each retry consumes a full round. Set num-rounds high and rely on early stopping (h ≤ h-tol).

Troubleshooting

Symptom	Cause	Fix
h(W) stays large / never converges	ρ too small or not enough rounds	Increase num-rounds (200+) or rho-init
SHD high but h(W) ≈ 0	Valid DAG but wrong edges	Tune graph-thres, l1-graph-penalty, or lr
NaN in loss	ρ exploded or lr too high	Lower lr, or set a tighter rho-max
FDR = 1.0	Threshold too low, keeping spurious edges	Increase graph-thres (try 0.3 → 0.5)
TPR = 0.0	Threshold too high, killing real edges	Decrease graph-thres (try 0.3 → 0.1)
Retries exhaust max-retries every AL step	ρ growing too fast without progress	Increase it-fl for more local training, or raise max-retries
Run stops well before num-rounds with "stuck at its ceiling" in the log	rho hit its ceiling — min(rho-max, rho-init × rho-runaway-factor) — and stayed there for stuck-patience consecutive AL steps with no h improvement — intentional early-stop (see Known issues #11 / #13)	Expected behavior; the run was provably stuck. To get more retries first, raise rho-runaway-factor (or stuck-patience, or rho-max if the ceiling is rho-max)
PowerShell --run-config fails	Scientific notation / quoting	Set values in pyproject.toml instead
Exception iterating responses: database is locked	SQLite contention in Flower's SuperLink	See note below
Key 'X' is not present in the main dictionary (Flower CLI error)	A --run-config override key isn't declared in pyproject.toml's [tool.flwr.app.config] at the time the job ran — usually a stale checkout running concurrently with a commit that adds a new key	Make sure the working copy on the HPC node has the commit that added the key before submitting; resubmitting after pulling fixes it
ActorDiedError / "the actor is dead... SYSTEM_ERROR" + SLURM error: Detected N oom_kill events in the log	Out-of-memory: the job's --mem-per-cpu × --cpus-per-task budget was exceeded, usually on heavy trials (num-rounds ~500, it-fl ~1000, larger datasets) with no client_resources reservation per simulated client	See Known issues #12 — raise mem-per-cpu in the .pbs script (now 4G in hydra_optuna.pbs), or lower concurrency via an explicit Ray client_resources config
Connection to the SuperLink is unavailable (Optuna trial pruned)	The local SuperLink for that trial's flwr run becomes unreachable. It happens at a wall-clock-consistent ~640–660s regardless of round number, concentrated on AS_FedDAG + the heaviest datasets — pointing at the flwr CLI's --stream log-reconnect loop losing the channel (see Known issues #12 / #14)	run_and_record now auto-retries the same config up to 3× on this signature before giving up, since the failure is environmental, not tied to the hyperparameters (Known issues #14). If all retries fail it's still auto-pruned (not scored -1.0, doesn't bias TPE)

---

Data

Simulation

Simulation logic is inherited from the original FedDAG repository. By default, a random DAG is generated with Erdős–Rényi sampling and a fixed graph-seed and data-seed. Each client receives its share of the simulated data.

Step 1 — Generating the causal graph

In simulation.py:

• A base undirected graph is created with NetworkX (Erdős–Rényi or Barabási–Albert).
• A random acyclic orientation is applied via _random_acyclic_orientation.
• Each edge gets a random weight sampled from [w_min, w_max], with random sign (50% positive, 50% negative).

Step 2 — Simulating data from the SEM

IIDSimulation provides simulate_linear_sem and simulate_nonlinear_sem. The selected one walks the DAG in topological order, computing each node from its parents:

X_j = f(X_parents @ W[parents, j]) + noise

The noise distribution depends on the chosen SEM type — Gaussian, Exponential, Gumbel, Uniform, or Logistic for linear; MLP, GP, MIM, GP-additive, or quadratic for nonlinear. All draws are seeded for reproducibility.

Linear SEM

A causal DAG with $d$ nodes is represented by a weighted adjacency matrix $W \in \mathbb\{R\}^\{d \times d\}$, where $W_\{ij\} \neq 0$ means node $j$ is a parent of node $i$. The data follows:

$X = X W + Z, \quad Z \sim \mathcal\{N\}(0, \sigma^2 I)$

When method = "linear", each node is computed as a weighted sum of its parents plus Gaussian noise. The sem-type parameter is ignored here — the noise is always Gaussian.

x_j = Σ w_i · x_i + ε,    ε ~ N(0, scale)

Nonlinear SEM types

All used with method = "nonlinear":

SEM type	Formula	Description
mlp	sigmoid(X @ W₁) @ W₂ + ε	Two-layer MLP with sigmoid activation. Weights ∈ [0.5, 2.0], random sign flips.
mim	tanh(X @ w₁) + cos(X @ w₂) + sin(X @ w₃) + ε	Mix of nonlinear basis functions (tanh, cos, sin) applied to parent values.
gp	GP.sample_y(X) + ε	Gaussian Process fitted jointly on all parents. Smooth nonlinear relationship.
gp-add	Σ_i GP.sample_y(X[:,i]) + ε	Additive GP — one GP per parent, summed. Simpler, more interpretable than gp.

All nonlinear types add ε ~ N(0, scale) noise.

Step 3 — IID vs. Non-IID partitioning

In partition.py:

• IID — all clients share the same SEM type.
• Non-IID — each client can have a different SEM type (linear / nonlinear), a different nonlinear function, and a different noise scale.

In short: the data is random but with a fixed underlying causal structure given by the DAG, and seeds guarantee reproducibility.

Sample-count and graph-sharing options

Two extra config keys control aspects of partitioning that are orthogonal to IID/Non-IID:

• partition-scheme (equal default | unequal) — equal gives every client num-samples rows, same as the original FedDAG behavior. unequal instead follows an arithmetic split (FedCSL paper Eq. 11 / split_dataset.m): client 0 gets the least data, client m-1 the most, while the total across all clients (num-samples × num-clients) stays the same as the equal case — so the two schemes are directly comparable. This applies to both simulated and benchmark data.
• per-client-w (0 default | 1) — with noniid=0 (IID), 0 makes every client train on the same weighted DAG W (only the noise realization differs per client — the statistically correct definition of IID). Setting it to 1 reproduces the legacy behavior where each client received its own W even in "IID" mode, which is technically not IID since the underlying causal mechanism differed across clients.

Dirichlet partitioning

The Dirichlet allocation is used by the strategies that operate on Non-IID data — currently GS_FedDAG. It acts as a "dispersion score" controlling how heterogeneous the clients are. Specifically, it governs the distribution of noise_scale across clients.

• Low alpha (e.g. 0.1) → clients concentrate on one noise level → high heterogeneity (Non-IID)
• High alpha (e.g. 10.0) → clients spread evenly → low heterogeneity (near IID)

The linearity (linear vs. nonlinear) and nonlinear_type (mlp / mim / gp / gp-add) columns in client properties are still drawn uniformly at random, independently of α. Only noise_scale is Dirichlet-distributed.

How it works in code:

In load_partition_noniid:

dirichletAlpha = run_config.get("dirichletAlpha")
if dirichletAlpha is not None:
    props = dirichlet_property_generation(num_partitions, float(dirichletAlpha), seed=graph_seed)

The dirichletAlpha value comes from the run config — set either via the CSV or the manual experiment menu.

Real-world benchmarks (bnlearn)

Three benchmarks ship with cached DAGs and are the most tested:

Dataset	Nodes	Arcs	Recommended rounds	Domain	Reference
sachs	11	17	150	Protein signaling	Sachs et al., Science 2005
child	20	25	150	Medical diagnosis	Spiegelhalter & Cowell
alarm	37	46	200	Patient monitoring	Beinlich et al., ECAI-Med 1989

You are not limited to these three. fedag.datasets.benchmarks.is_benchmark_name treats any dataset name other than "simulation" as a real network, loaded through bnlearn/pgmpy's fallbacks — so asia, insurance, hailfinder, win95pts, pigs, link, etc. work without being added to a list first (an unknown name simply raises a clear error if the network truly doesn't exist). Round/ρ/L1 defaults scale from num-nodes, not the dataset name, so they adapt to any graph size.

---

Metrics

All metrics are computed after training by thresholding |W_final| ≥ graph-thres to obtain a binary estimated graph B_est, and comparing against the ground truth B_true. The metric implementations are taken from the original FedDAG repository.

Metric	Direction	Meaning
SHD (Structural Hamming Distance)	lower = better	Edge insertions + deletions + reversals needed to match the true graph
TPR (True Positive Rate / Recall)	higher = better	Proportion of true edges correctly recovered
FDR (False Discovery Rate)	lower = better	Proportion of predicted edges that are wrong (including reversals)
FPR (False Positive Rate)	lower = better	Wrong predictions over all possible negatives
Precision	higher = better	TP / (TP + FP)
F1	higher = better	Harmonic mean of precision and recall
G-score	higher = better	max(0, TP − FP) / (TP + FN) — the paper's primary metric; penalizes false positives
h(W)	→ 0	Acyclicity constraint value; training converges when h ≤ h-tol
is_dag	true	Whether the thresholded graph is acyclic (h(W) < 1e-3)

---

Hyperparameter Tuning

Two Optuna-based scripts automate hyperparameter search. Both persist their study in a SQLite file so you can interrupt and resume at any time.

---

tune_optuna.py — Standard Bayesian search

Finds the best hyperparameters for a fixed dataset / federation setup using Optuna TPE. Each trial launches flwr run, records metrics in results/<machine>/results.csv, and reports back to Optuna.

Score: F1 - 0.01 × SHD (penalized to -1 if the result is not a DAG, -0.5 if the graph is empty).

Search space (linear, AS_FedDAG_linear): AL parameters (rho-init [1e-3, 1.0], rho-scale [1.1, 5.0], h-thres [0.1, 0.9], init-iter 2..25, al-every 1..3, max-retries 5..15, rho-runaway-factor [1e3, 1e8]), local training (lr [1e-4, 1e-2], it-fl ∈ {200, 300, 500, 1000}), thresholding (graph-thres [0.05, 0.5]), and lambda1 [5e-3, 5e-2]. Number of rounds is sampled from [50, 100, 150, 200] for simulation and from [100, 150, 200, 300, 400, 500, 600, 800] for real benchmarks, which need more rounds to converge.

Search space (nonlinear, GS_FedDAG / AS_FedDAG): narrowed to the hyperparameters that actually move F1 — lr [1e-4, 5e-3], graph-thres [0.05, 0.25], rho-init [1e-3, 0.2], l1-graph-penalty [2e-4, 5e-3], temperature [0.2, 1.0], num-hidden-layers 2..5, h-thres [0.3, 0.85] (the first four log-uniform), plus two knobs added later: rho-runaway-factor [1e3, 1e8] log-uniform and al-mode ∈ {retry, paper} — see Known issues #13 for why the flat 1e4 default had to become tunable. The rest are fixed to representative values: rho-scale = 2.0, init-iter = 10, al-every = 1, max-retries = 10, w-init = 0.15, hidden-size = 8 for AS_FedDAG / 16 for GS_FedDAG. Two notable fixed choices:

• it-fl is fixed at 1000, deliberately above any realistic it-fl-cap, so the cap — derived from num-nodes alone — is what actually decides the effective local-step budget. See it-fl-cap: who sets it and to what.
• num-rounds is fixed at 800 for real/benchmark data and 150 for simulation. 800 is a ceiling, not a target: AL early-stops long before it on graphs that converge (sachs/child); the headroom exists for the pigs/link scale, which has its own unresolved problem.

Why is num-rounds fixed and not tuned (doesn't it affect the score)? It does — but mostly as an upper bound, because of early stopping. A run ends as soon as it converges (h ≤ h-tol) or gets stuck at the ρ ceiling (stuck-patience); the remaining rounds never execute. So once num-rounds is set high enough to let training converge, raising it further changes nothing, and the Optuna analysis confirmed it (|corr| with F1 < 0.1). Fixing it to a generous value and letting early stopping decide the real round count frees the limited trial budget for the hyperparameters that genuinely move F1.

Caveat: "low correlation" isn't "zero". num-rounds interacts with it-fl (total optimization ≈ num-rounds × local steps/round) — e.g. gs_sim's best historical trial used few rounds with many local steps each. Fixing both removes the ability to rediscover that specific tradeoff via search. The fix is a deliberate budget-allocation choice, not a claim that num-rounds is irrelevant; if a particular dataset clearly wants an unusual rounds/steps balance, make num-rounds a suggest_categorical again just for that case.

Any dataset works: --dataset is no longer restricted to simulation/sachs/child/alarm. Anything that isn't "simulation" is treated as a real bnlearn/bnrepository network (asia, insurance, hailfinder, win95pts, …) via fedag.datasets.benchmarks.is_benchmark_name, and the round/rho/L1 defaults scale from num-nodes rather than a hardcoded name list.

Usage:

# Linux / macOS
python scripts/tune_optuna.py \
    --strategy AS_FedDAG_linear \
    --dataset child \
    --num-nodes 20 \
    --sem-type linear \
    --data-seed 2022 \
    --num-clients 2 \
    --n-trials 50 \
    --study-name as_linear_child20 \
    --machine VM

# Windows PowerShell
python scripts/tune_optuna.py `
    --strategy AS_FedDAG_linear `
    --dataset child `
    --num-nodes 20 `
    --sem-type linear `
    --data-seed 2022 `
    --num-clients 2 `
    --n-trials 50 `
    --study-name as_linear_child20 `
    --machine PC

Studies are saved to studies/<study-name>.db. A flat CSV of all trials is exported to studies/<study-name>_trials.csv at the end.

Resuming: re-run the exact same command — Optuna loads the existing study and continues from where it left off.

---

tune_stress.py — Ceiling / stress-test search

Finds the performance ceiling of each strategy by also tuning federation scale and data heterogeneity. Use this to answer: how well can this strategy perform under the hardest realistic conditions?

Score: (F1 - 0.01 × SHD) + 0.1 × difficulty

where difficulty = log(num_nodes / 10) + log(num_clients / 2) + log(1 / dirichletAlpha) (last term only for non-linear strategies). Optuna is rewarded for configs with many nodes, many clients, and high non-IID-ness that still produce a valid DAG with good F1.

Fixed: dataset = simulation.

Additional search dimensions (moderate ranges):

Parameter	Range
num-nodes	[10, 20, 50]
num-clients	2..10
num-rounds	[50, 100, 150, 200, 300]
dirichletAlpha	0.1..5.0 log-uniform (non-linear only)

Output paths:

Artifact	Location
Optuna study DB	studies/stress_test/<study-name>.db
Trial CSV	studies/stress_test/<study-name>_trials.csv
Results CSV	results/stress_test/<machine>/results.csv
Experiments CSV	results/stress_test/<machine>/experiments.csv

Usage:

# Linear strategy (IID, no dirichlet)
python tune_stress.py \
    --strategy AS_FedDAG_linear \
    --sem-type linear \
    --study-name stress_linear \
    --n-trials 60 \
    --machine VM

# Non-linear strategies (dirichletAlpha tuned automatically)
python tune_stress.py \
    --strategy GS_FedDAG \
    --sem-type mlp \
    --study-name stress_gs \
    --n-trials 60 \
    --machine VM

python tune_stress.py \
    --strategy AS_FedDAG \
    --sem-type mlp \
    --study-name stress_as \
    --n-trials 60 \
    --machine VM

---

inspect_study.py — Study inspection

Loads any persisted Optuna study and prints a ranked summary of trials. Useful for comparing studies across machines or picking the best config to reproduce.

python inspect_study.py --study-name as_linear_child20

---

verify_top_trials.py — Seed-robustness check

GS_FedDAG / AS_FedDAG initialize their per-node MLP weights randomly, so each Optuna trial during the search effectively got one random draw. A high-scoring trial could reflect genuinely good hyperparameters — or just a lucky init that a different seed would not reproduce. (AS_FedDAG_linear has no random init, so it doesn't need this.)

This script re-runs a finished study's top-K trials once per seed and reports the spread of the resulting score. A trial whose seeds disagree wildly (one near the original value, the rest much worse) was likely a lucky draw, not a config worth shipping. It sets model-seed per run to control the init.

python scripts/verify_top_trials.py \
    --study-name gs_sachs --strategy GS_FedDAG --dataset sachs \
    --num-nodes 11 --sem-type mlp --num-clients 5 --num-samples 10000 \
    --top-k 3 --seeds 1 2 3 4 5

It prints, per candidate, the per-seed scores plus mean / stdev / min / max, and flags high-variance candidates. Use it to pick the robust best config, not just the single highest trial value.

---

Limitations

Algorithmic

• Scalability — the acyclicity constraint uses trace(expm(W⊙W)), which is O(d³) per evaluation and is called once per local gradient step. In practice, training becomes slow or numerically unstable beyond ~20 nodes for the nonlinear strategies. The code does not warn about this on its own; the only size-related warning is the one it-fl-cap logs when it clips the local-step count.
• Convergence sensitivity — the AL parameters rho-init and rho-scale must be tuned per strategy. Values too large cause divergence; too small leads to extremely slow convergence.
• Identifiability — NOTEARS finds a DAG consistent with the data, not necessarily the ground-truth DAG. Without extra assumptions (e.g. equal noise variances), the true DAG is not identifiable from observational data alone.
• AL-update cadence mismatch — the paper updates α and ρ every ~5 outer iterations; Flower updates them every round. Setting rho-scale = 5.0 partially compensates, but may not exactly reproduce the paper's results.

Implementation

• Simulation-only data caches — fedag/datasets/benchmarks.py holds _dag_cache, _data_cache_real and _bif_model_cache at module level. This works for single-process Flower simulation but will break in a real distributed deployment where each client is a separate process. (partition.py itself no longer caches: it either regenerates the partition or, when data-dir is set, reads client_<id>.npy from disk — the latter path is distribution-safe. GS_FedDAG's per-client MLPs likewise persist correctly through Flower's context.state["gs-nets"], which is per-node and distribution-safe.)
• No client fault tolerance — if a selected client stops responding mid-round, the round fails. There is no retry or fallback mechanism.
• No privacy guarantees — the shared weight matrix W encodes information about local data distributions. There is no differential privacy, secure aggregation, or gradient noise.
• No communication compression — W is sent as a full float32 matrix every round. No quantization or sparsification is applied.
• Fixed graph size — num-nodes (d) is fixed at startup; changing it requires restarting the experiment.
• Ground-truth DAG required for evaluation — all recovery metrics (SHD, TPR, etc.) require B_true, which is only available in simulation and synthetic-benchmark modes.

database is locked — Flower SuperLink SQLite contention

Flower's SuperLink stores run state in a SQLite file (~/.flwr/local-superlink/state.db). A background cleanup task periodically deletes expired tokens (DELETE FROM token_store WHERE active_until < ?). When that cleanup runs at the exact same moment another Flower component writes to the DB, SQLite raises:

Exception iterating responses: (sqlite3.OperationalError) database is locked

This error comes from inside Flower, not from your code. The process does not crash, but the current response iterator is interrupted, so the round returns no results (SHD=? F1=? [FAIL]). It is more likely to happen on long runs with many clients because there is more concurrent write activity.

Possible fixes (not applied in this repo):

1. Enable WAL mode — SQLite's Write-Ahead Logging allows concurrent reads alongside a single writer, eliminating most lock conflicts. Add the following after the SuperLink's state.db has been created:

   import sqlite3
   conn = sqlite3.connect(r"~/.flwr/local-superlink/state.db")
   conn.execute("PRAGMA journal_mode=WAL;")
   conn.close()

2. Upgrade Flower — this race condition has been improved in newer Flower releases. Check your version with python -c "import flwr; print(flwr.__version__)" and upgrade if behind.
3. Reduce concurrent load — fewer clients or lower num-samples means less simultaneous write activity and a lower chance of collision.

---

Known issues / improvement notes

1. Aggressive rho-scale collapses W. Going from 2.0 to 10 destroys TPR. Mechanism: when ρ is huge, the 0.5·ρ·h² term dominates the loss, and gradients push every W[i,j] toward zero to minimize h regardless of how well it predicts X. Result: sparse but empty W, h ≈ 0 but TPR = 0. Keep rho-scale = 2.0 and let max_retries = 10 absorb the slower convergence.

2. Adam reset in GS / AS strategies. The linear strategy persists Adam state via context.state; the nonlinear ones recreate Adam every round. This asymmetry may explain part of the nonlinear tuning difficulty. Possible fix: persist Adam state for the MLPs (not for W, since W comes from server aggregation and resetting there makes sense).

3. Module-level data caches in benchmarks.py (_dag_cache, _data_cache_real, _bif_model_cache). Works in simulation because all clients run in the same Python process. It breaks in real distributed deployments — each client would need to load its data from disk. Partly mitigated already: setting data-dir makes partition.py read client_<id>.npy per client instead of regenerating, which is the distribution-safe path. (The old _model_cache for GS MLPs is gone — those now persist through context.state["gs-nets"].)

4. _logits_to_prob is implemented twice — once in nonlinear_model.MaskedNNModel.get_W and once in server_app._logits_to_prob. Changing τ in one requires remembering the other. Should be extracted to a single function in metrics.py.

5. set_optimizer_state can silently fail to assign (Keras variable-name matching). Stale — carried over from the TensorFlow original. LinearDAGModel.set_optimizer_state () assigns self.optimizer.state[self.W] directly from the numpy arrays; there is no name matching and no Keras involved, so the described silent-failure mode can't occur. What can happen is the reverse: get_optimizer_state returns None before Adam's first step, so round 1 legitimately restores nothing. Only AS_FedDAG_linear uses this path at all.

6. Per-strategy thresholds are hardcoded. Resolved. visualize_W.py's _STRATEGY_THRESHOLD (0.35 / 0.5 / 0.5) now explicitly mirrors STRATEGY_DEFAULTS["graph-thres"] in fedag/config.py (same values), and the code comment says so directly. No remaining inconsistency between final evaluation and visualization.

7. Evaluation always uses the linear path (_evaluate_linear). Resolved / no longer applicable. fedag/evaluation.py only contains MetricsDAG (structural metrics: SHD, TPR, FDR, F1, G-score) — there is no _evaluate_linear / _evaluate_nonlinear MSE path anymore. server_app._print_final_metrics converts W_final from logits to probabilities via sigmoid(W/τ) for GS_FedDAG/AS_FedDAG before thresholding and computing structural metrics, so the comparison against B_true is already strategy-aware. The structural metrics never depended on the local MSE loss in the first place.

8. Short warm-up for GS_FedDAG (init-iter = 2). With al-every = 1, the model barely has a reasonable W when AL updates start. The other two strategies already use longer warm-ups (AS_FedDAG_linear 8, AS_FedDAG 10); for d > 15 consider raising GS to the same range.

9. compute_loss_nonlinear is not vectorized. Stale — already fixed, verified 2026-06-21. fedag/nonlinear_model.py already batches all d per-node MLPs as single tensors (net_weights[l] shape (d, in, out)) and runs the forward pass with torch.baddbmm — there is no Python loop over nodes. The real per-step cost for GS/AS on child/alarm is h_func_nonlinear's torch.linalg.matrix_exp (O(d³)), called once per local gradient step (up to it-fl times, per client, per round) — that's the acyclicity constraint's own algorithmic cost (see Next Steps #1), not a missed vectorization. See item 11 below for what was actually making GS trials so slow.

10. compute_h (numpy/scipy) and h_func (torch) handle NaN differently. scipy_expm can return NaN matrices when W is large; compute_h does not check and silently propagates them.

11. al-mode="retry" had no early-stop once stuck at rho_max. Fixed 2026-06-21. When the inner retry loop exhausts max_retries while rho is already capped at rho_max, escalating further is impossible — but the strategy kept running the remaining num-rounds anyway, each one a full round of local training (it-fl steps × all clients) with zero chance of progress. Confirmed in logs/optuna_12573708.out: a gs_child trial sat at h≈1.8e-06 from round ~150 to round 300 while alpha kept climbing and rho stayed pinned at 1e14 — ~150 wasted rounds. Combined with num-rounds up to 500 and it-fl up to 1000 in the Optuna search space, individual trials took up to ~8h, and the 24h SLURM walltime cut the gs_child / gs_alarm studies off at 14/60 and 48/60 trials — starving the TPE search, which is most of why GS looked so much worse than AS_nonlinear on those datasets. Fix: new stuck-patience config key (default 2) — after that many consecutive AL give-ups with rho already at its ceiling, FedDAGBaseStrategy sets self.stopped = True and the round loop exits early (the same mechanism al-mode="paper" already used). num-rounds and it-fl search ranges are untouched — this only cuts off trials that are provably stuck, not what a healthy trial does. See #13 for why the original rho_max-only ceiling was almost never reached in practice.

12. HPC Optuna jobs were OOM-killed on heavy trials. Fixed 2026-06-21. hpc/hydra_optuna.pbs requested mem-per-cpu=2G × cpus-per-task=16 = 32GB for the whole node, with no client_resources/backend_config set anywhere (Flower/Ray use their defaults — no per-actor memory reservation). With num_supernodes=5 clients running concurrently each round, and the heaviest Optuna trials combining num-rounds up to 500, it-fl up to 1000, and larger datasets (sachs, alarm, 37 nodes), per-process memory crept past the 32GB cap before the run finished. Confirmed directly in the logs: optuna_12573707.out and optuna_12573710.out both show error: Detected N oom_kill events in StepId=....batch from SLURM, and optuna_12573712.out shows a raylet marked dead from missed heartbeats (the same symptom — the OS killed a process and took the node's Ray runtime down with it). The Ray/Flower traceback ("the actor is dead... SYSTEM_ERROR... connection error") is just the downstream symptom; the SLURM line is the actual cause. Fix: raised mem-per-cpu to 4G in hpc/hydra_optuna.pbs (32GB → 64GB total, matching what hydra_stress.pbs already used). num-rounds/it-fl ranges untouched.

    Softer symptom, same root cause: optuna_12573709.out (gs_alarm, GS_FedDAG on the 37-node alarm dataset — the heaviest combo in the search space). 25 of 47 trials attempted in that 24h job (53%) were pruned with Connection to the SuperLink is unavailable — no SLURM oom_kill line this time, just the local SuperLink becoming unreachable mid-run. Checked where each death happened: always within the first ~1–80 rounds of that specific trial's own flwr run process, regardless of the configured num-rounds (500, 300, 150, 100...) — i.e. death tracks wall-clock/memory growth within a single process, not a fixed round count. This is consistent with the same memory-growth pattern as above (heavy nonlinear model + many local steps accumulating per round) hitting the ceiling before SLURM's per-step OOM monitor catches it. The code already has a safety net for this exact signature — run_experiments.py flags "Connection to the SuperLink is unavailable" as infra_error, and tune_optuna.py prunes that trial (optuna.TrialPruned) instead of scoring it -1.0, so it doesn't bias the TPE search — but it still burns over half the trial budget on gs_alarm. The mem-per-cpu=4G fix above should reduce how often this triggers. A later re-investigation (2026-06-29) found a second, time-triggered flavor of this — see #14.

13. rho_max alone was an unreachable early-stop ceiling. Fixed 2026-06-29. The stuck-patience early-stop from #11 only fires once rho is pinned at rho_max, but rho_max defaults to 1e14–1e16 while a typical rho-init is 1e-3…1 — so a genuinely runaway trajectory wanders through 5–10 orders of magnitude (observed rho ≈ 7e10 in as_sim overnight logs, with h never recovering) without ever reaching rho_max, and burns the full num-rounds anyway — exactly the waste #11 was meant to prevent. Fix: new rho-runaway-factor config key (default 1e4). The strategy now caps ρ at min(rho_max, rho_init × rho-runaway-factor) (self.rho_ceiling), giving stuck-patience a ceiling it can actually reach for runaway trials while leaving healthy escalation (a few orders of magnitude above rho_init) untouched. Set 0 to fall back to rho-max only. Reuses the entire #11 stuck-patience machinery — only the ceiling moved.

14. Connection to the SuperLink is unavailable is partly time-triggered, not only OOM. Fixed 2026-06-29. Re-examining the 2026-06-28 overnight logs showed this error landing at a wall-clock-consistent ~640–660s into a trial regardless of round number (alarm dies ~round 16–22, child ~round 31–65 — different round, same elapsed time), concentrated almost entirely on AS_FedDAG + the heaviest datasets. That signature points at the flwr CLI's own --stream log-follow loop (flwr/cli/log.py::start_stream, which reconnects to the SuperLink's StreamLogs endpoint every 60s) losing the channel — not FedDAG code, and not strictly OOM. Turning --stream off isn't a safe drop-in fix (the CLI then returns before the run finishes, breaking run_experiments.py's assumption that the subprocess blocks until done). Mitigation: run_and_record now retries the identical config up to 3 attempts when result["infra_error"] is set, accumulating elapsed_s across attempts, before handing control back to the caller — so tune_optuna.py's existing TrialPruned fallback only fires after retries are exhausted. Most retries succeed because the failure isn't tied to the sampled hyperparameters. Flagged as a mitigation, not a proven root-cause fix.

15. Nonlinear MLP init was unseeded — trial scores included random-init noise. Fixed 2026-06-29. Nothing called torch.manual_seed anywhere in fedag/; only the dataset generation was seeded (data-seed). So every Optuna trial of GS_FedDAG / AS_FedDAG got a fresh, uncontrolled random MLP init, meaning part of the trial-to-trial F1 swing (e.g. as_sachs jumping between −0.5 and 0.6) was init luck, not hyperparameter signal — and a "best" trial might not reproduce. Fix: new model-seed config key (default 42). AS_FedDAG seeds once server-side before building the shared initial arrays; GS_FedDAG seeds per client as model-seed + partition-id (its local nets are never federated). AS_FedDAG_linear is unaffected (W starts at 0, no random init). New scripts/verify_top_trials.py uses this to re-run a study's top-K trials across several seeds and report score variance, so you can pick the robust best config rather than a lucky one.

16. Moderate-size nonlinear graphs (21 ≤ d ≤ 100) collapse to an empty W. Attempted 2026-07-08, then reverted — still open. The it-fl-cap = max(100, 6000 // d) heuristic gives d=40 only ~150 local steps/round (vs. 600 at d=10). Since GS_FedDAG/AS_FedDAG train one MLP per node, at d ≳ 20 those MLPs are under-trained, and the acyclicity penalty (ρ↑) pushes W to zero (h ≈ 0, empty graph) before the model learns any edges — collapsing to F1 = 0.

    The attempted fix was a tiered it_fl_cap(num_nodes) helper in fedag/config.py (full 1000-step budget for 21 ≤ d ≤ 100), plus a rho_schedule(num_nodes, strategy) companion interpolating rho-init/rho-scale in log-space between the paper's Table 8 values for d=10/20/40. ⚠️ Neither function exists in fedag/config.py today — the module only exposes auto_scale_defaults, resolve_run_config and dirichlet_alpha_or_none, and rho-init/rho-scale come from the plain SIM_NODE_RHO closest-match lookup (no interpolation). The per-caller 6000 // d values described in it-fl-cap: who sets it and to what are what actually runs.

    Independently of the missing helper, the 2026-07-29 measurements refuted the underlying hypothesis anyway: raising it-fl-cap to 1000 at d=40 and re-tuning from scratch produced no usable trial (F1 0.00–0.18 across 60 trials) — see the d=40 collapse analysis, where the AL-cadence mismatch (b) is the current leading explanation. Treat this item as superseded, not fixed.

17. The nonlinear h is squared on the client and unsquared on the server. Open — found 2026-08-06. The two sides of the AL loop do not compute the same quantity for GS_FedDAG / AS_FedDAG:

    Where	Call	Formula
    Client (training loss)	h_func_nonlinear ()	Tr[e^{g⊙g}] − d — squared
    Server (AL update + early stop)	compute_h(..., square=False) ()	Tr[e^{g}] − d — unsquared

    Since p ≥ p² for p ∈ [0,1], the server's h is systematically larger than the one the clients are minimizing — by roughly an order of magnitude, per h_func_nonlinear's own docstring. Consequences: convergence (h ≤ h-tol) is judged against a stricter number than the clients optimize, and the h > γ·h_prev ρ-escalation trigger fires on a different scale than the training signal. This is a plausible contributor to the nonlinear tuning difficulty and to the ρ-runaway behaviour of #13, though it has not been isolated experimentally.

    tests/test_nonlinear_model.py::test_h_nonlinear_matches_numpy_no_square currently fails and is the visible symptom: it asserts the client and server agree on the unsquared form. The docstring of h_func_nonlinear, by contrast, states the square was reverted to on 2026-06-15 on the strength of an A/B test (sachs F1 0.68 squared vs 0.10 unsquared). So the code, the test, and the docstring encode three different intentions and one of the three has to give. Decide which form is canonical, make both call sites use it, and update the test — do not simply relax the assertion.

---

Next Steps

1. Scale beyond 20 nodes — replace expm with an approximate differentiable acyclicity constraint such as the log-determinant barrier from DAGMA or the polynomial approximation from GOLEM, reducing per-step cost from O(d³) to O(d²).
2. Real-data support — implement a load_real_data partition function that reads CSV/Parquet files per client, enabling experiments on real federated datasets (e.g. electronic health records, IoT sensor streams).
3. Differential privacy — integrate DP-SGD (e.g. Opacus for PyTorch) on the client side before returning W, and add calibrated Gaussian noise on the server before broadcasting the aggregated W.
4. Communication efficiency — quantize W to 8-bit integers or apply top-k sparsification before upload. This reduces per-round communication from O(d²·32 bits) to O(d²·8 bits) or O(k·32 bits).
5. Client robustness — expose fraction_train as a real config key (today it is a hardcoded 1.0 constructor default on FedDAGBaseStrategy), then add server-side logic to skip rounds when fewer than the minimum number of clients respond.
6. Time-series extension — adapt the SEM simulation and client logic to support time-lagged DAGs (DYNOTEARS), enabling causal discovery on sequential or longitudinal federated data.
7. Persist Adam state in GS_FedDAG (at least for the MLPs) — likely unblocks nonlinear tuning.
8. Vectorize compute_loss_nonlinear — already done (see Known issues #9); no action needed here.
9. Centralize logits_to_prob and the threshold logic in a single module.
10. Implement _evaluate_nonlinear that uses the nonlinear model's actual loss. — moot: see Known issues #7. Final metrics are structural (SHD/TPR/FDR/F1/G-score against B_true), not MSE-based, so there's no "linear vs. nonlinear loss" path to fix.
11. Explicit success/failure logging in set_optimizer_state. — moot: see Known issues #5. The PyTorch implementation cannot fail silently.
12. Automated sweeps over (rho_init, rho_scale, lambda1, graph_thres) for nonlinear — combined with the real-time monitor, runs that clearly collapse can be early-stopped.
13. min_clients decoupled from fraction_train — currently a missing client is silently ignored; a configurable minimum would be useful.
14. Reconcile the squared vs. unsquared nonlinear h (Known issues #17) — pick one form, apply it on both the client and the server, and fix the failing test. Worth doing before the next d=40 re-tune, since it changes what "converged" means for GS_FedDAG / AS_FedDAG.
15. Unify the it-fl-cap derivation — extract the 6000 // d heuristic into one helper in fedag/config.py and have tune_optuna.py, bench_entrypoint.py and verify_top_trials.py call it, so a config is never benchmarked under a different local-step budget than it was tuned with (see it-fl-cap: who sets it and to what).

---

Glossary

Term	Meaning
FedDAG	Federated learning algorithm for learning Directed Acyclic Graphs (causal structures) from distributed data
W_final	Learned weighted adjacency matrix (model output)
B_true	True underlying DAG structure (ground truth for evaluation)
Threshold	Value used to binarize W (e.g. 0.05): edges with `
Logits → Probabilities	For GS_FedDAG / AS_FedDAG, W contains logits converted via sigmoid with temperature τ
DAG	Directed Acyclic Graph — no directed cycles (paths returning to a node)
TP / FP / FN	True Positive, False Positive, False Negative — standard classification metrics for edge recovery

---

Project Structure

FedDAG-Flower/
├__ data/
│   └__ benchmarks/                      # cached bnlearn DAGs (sachs, child, alarm)
├__ docs/images/                         # screenshots used by this README
├__ fedag/                               # — the Flower app (what `flwr run` executes) —
│   ├__ datasets/
│   │   ├__ benchmarks.py                # sachs (11), child (20), alarm (37) via bnlearn
│   │   ├__ partition.py                 # IID and Non-IID client splitting
│   │   └__ simulation.py                # random DAG + SEM data generation (original repo)
│   ├__ client_app.py                    # ClientApp — dispatches by strategy
│   ├__ server_app.py                    # FedDAGBaseStrategy + 3 variants + ServerApp
│   ├__ config.py                        # SINGLE source of defaults for everything
│   ├__ dag_model.py                     # linear NOTEARS model
│   ├__ nonlinear_model.py               # MaskedNNModel (batched/vectorized MLPs)
│   ├__ metrics.py                       # h(W), is_dag, count_accuracy (wrapper)
│   ├__ evaluation.py                    # MetricsDAG: SHD, TPR, FDR, F1, G-score
│   └__ compat.py                        # Flower version compatibility (HPC)
├__ scripts/                             # — local orchestration (run with python) —
│   ├__ run_experiments.py               # interactive launcher + CSV manager + metric parsing
│   ├__ run_batch_hpc.py                 # non-interactive batch runner (used by hpc/)
│   ├__ tune_optuna.py                   # Bayesian hyperparameter search
│   ├__ tune_stress.py                   # ceiling/stress search (nodes, clients, dirichlet)
│   ├__ verify_top_trials.py             # re-run a study's top-K trials across seeds (init robustness)
│   ├__ inspect_study.py                 # ranked summary of any persisted Optuna study
│   ├__ export_studies_excel.py          # consolidate studies/*.db into one Excel
│   └__ visualize_W.py                   # 7-panel dashboard: W heatmap, graph, metrics
├__ hpc/                                 # — jobs for the VUB Hydra cluster (SLURM) —
│   ├__ setup_hydra.pbs                  # one-time env setup
│   ├__ hydra_env.pbs                    # shared env sourced by every job
│   ├__ hydra_job.pbs                    # batch experiments from CSV
│   ├__ hydra_optuna.pbs                 # Optuna search
│   ├__ hydra_stress.pbs                 # stress test
│   └__ env_setup.sh                     # sourced by Bechmarking-Suite's generic_job.pbs
├__ tests/                               # pytest suite: test_metrics, test_al_update,
│                                        #   test_partition, test_config, test_nonlinear_model
├__ bench_entrypoint.py                  # Bechmarking-Suite contract (run / resources)
├__ run_flwr_mac_os.sh                   # macOS smoke test (gRPC fork-safety env vars)
├__ logs/     (generated)                # HPC job stdout (optuna_<jobid>.out)
├__ models/   (generated)                # saved W_final.npy / test.npy per experiment
├__ results/  (generated)                # results.csv + experiments.csv per machine tag
├__ studies/  (generated)                # Optuna .db files + <study>_trials.csv
├__ pyproject.toml                       # Flower app config + hyperparameter overrides
├__ CITATION.cff                         # how to cite this repo
├__ LICENSE                              # Apache-2.0
└__ README.md

---

Tests

The tests/ suite is plain pytest and needs no cluster, no Flower run and no downloaded benchmark:

pytest -q

It covers the metric implementations (test_metrics.py), the AL update rule in both modes (test_al_update.py), the partitioning schemes (test_partition.py), the defaults/precedence in fedag/config.py (test_config.py), and the batched MaskedNN forward pass (test_nonlinear_model.py).

⚠️ Current status: 42 passed, 1 failed. The failure is test_nonlinear_model.py::test_h_nonlinear_matches_numpy_no_square, and it is a real finding rather than a flaky test — see Known issues #17.

---

License & Citation

Released under the Apache License 2.0 — see LICENSE.

Citation metadata lives in CITATION.cff; GitHub renders it as a "Cite this repository" button. If you use this code, please cite both this repository and the original FedDAG paper (reference 1 below).

---

References

1. Gao, R. et al. FedDAG: Federated DAG Structure Learning. TMLR, 2023.
2. Zheng, X. et al. DAGs with NO TEARS: Continuous Optimization for Structure Learning. NeurIPS 2018.
3. Ng, I. et al. Masked Gradient-Based Causal Structure Learning. SDM 2022.
4. McMahan, B. et al. Communication-Efficient Learning of Deep Networks from Decentralized Data. AISTATS 2017.