# FedNewton Experiments Repository

This repository contains the implementation and experimental framework for **Federated Learning with Second-Order Optimization**, focusing specifically on Newton-style server updates, namely FedNewton.

It provides a robust testbed for comparing these advanced optimizers against standard first-order baselines (FedAvg, FedProx) and second-order methods (FedSophia, FedNew) across various computer vision tasks.

## 📂 Repository Structure

```text
.
├── main.py                 # Core logic: Client training & Server aggregation algorithms
├── analysis.py             # Post-processing: Aggregates .npz logs into plots & CSVs
├── plot_grid_search.py     # Utility: Visualizes results specifically from grid search runs
│
├── run_example.sh          # Demo: Simple single-run script to verify environment
├── run_grid_search.sh      # Tuning: Hyperparameter search (e.g., Newton LR/Damping)
│
├── exp_comparison.sh       # Exp 1: Compare Methods
├── exp_solvers.sh          # Exp 2: Compare Solvers (Exact vs CG vs Approx vs Diag)
├── exp_backbones.sh        # Exp 3: Compare Models (ResNet vs ViT)
│
└── README.md               # Documentation
```

## 🛠️ Requirements

*   **Python:** 3.8+ (Tested on 3.12)
*   **Core:** PyTorch, torchvision, numpy, pandas, matplotlib
*   **Optional:** `timm` (for ViT/EfficientNet backbones), `clip`

**Quick Install:**
```bash
pip install torch torchvision timm numpy matplotlib pandas
```

## ⚡ Quick Start

### 1. Basic Run (Sanity Check)
Use the example script to run a minimal experiment (e.g., small number of rounds on CIFAR-10) to ensure your environment is set up correctly.
```bash
bash run_example.sh
```

### 2. Hyperparameter Tuning
To find the optimal Learning Rate or Damping factor for FedNewton:
```bash
bash run_grid_search.sh
```

### 3. Running Full Experiments
The repository includes three pre-configured experiment suites. These scripts automatically handle multi-GPU parallelization:

*   **Compare Algorithms (FedAvg vs. FedNewton):**
    ```bash
    bash exp_comparison.sh
    ```
*   **Compare Linear Solvers (Performance/Time Trade-off):**
    ```bash
    bash exp_solvers.sh
    ```
*   **Compare Backbones (ResNet vs. ViT):**
    ```bash
    bash exp_backbones.sh
    ```

## 🔬 Experimental Protocol

### Supported Methods
| Category | Algorithms |
| :--- | :--- |
| **Baselines** | FedAvg, FedProx |
| **Second-Order** | FedSophia, FedNew, FedNewton|

### Newton Solver Variants
Configurable via the `exp_solvers.sh` script or `--newton_solver` argument:
*   **L-BFGS:** Quasi-Newton approximation.
*   **LowRank:** Top-k Eigenvalue approximation.
*   **CG:** Conjugate Gradient (Hessian-Free).
*   **Exact:** Full matrix inversion (High memory).

### Standard Settings
*   **Heterogeneity:** Controlled by Dirichlet alpha (`--alpha`). Lower $\alpha$ means higher heterogeneity.
*   **Clients:** Default 10 clients, typically `fraction=0.5` or `1.0` participation.
*   **Data:** CIFAR-10, CIFAR-100, MNIST (auto-downloaded).

## 📊 Metrics & Analysis

### Logging
Metrics are saved per-round in `.npz` format under `runs/<experiment_id>/logs/`.
*   **Metrics:** `test_acc`, `test_loss`, `train_loss`, `wall_time`.

### Visualizing Results
Use `analysis.py` to generate comparative plots (Accuracy vs. Rounds, Accuracy vs. Time) for a specific experiment folder:

```bash
# Example: Analyze results from the solver comparison experiment
python analysis.py --log_dir runs/compare_solvers_20231130/logs
```

For grid search visualization:
```bash
python plot_grid_search.py --root_dir runs/grid_search_task
```

## 📝 Practical Notes
1.  **Memory Management:** For large models (ViT, DINOv2), use `--training_mode head`.
2.  **Concurrency:** The `exp_*.sh` scripts use FIFO pipes to manage GPU tokens. Ensure you set the `GPUS` array inside these scripts to match your hardware (e.g., `GPUS=(0 1 2 3)`).
3.  **Reproducibility:** Do not modify `main.py` while an experiment is running. The script backs up the code *at the start* of the run, but dynamic changes could affect running threads.
