# FEDzk: Secure Federated Learning with Zero-Knowledge Proofs

<div align="center">
  <!-- Logo image refreshed -->
  <img src="assets/images/fedzklogo.png" alt="FEDzk Logo" width="400">
  <h1>FEDzk: Federated Learning with Zero-Knowledge Proofs</h1>
  <p>
    <strong>A secure and privacy-preserving framework for federated learning using zero-knowledge proofs</strong>
  </p>
  <p>
    <a href="#-project-overview"><strong>Overview</strong></a> •
    <a href="#-features"><strong>Features</strong></a> •
    <a href="#-architecture"><strong>Architecture</strong></a> •
    <a href="#-system-requirements"><strong>Requirements</strong></a> •
    <a href="#-installation"><strong>Installation</strong></a> •
    <a href="#-quick-start"><strong>Quick Start</strong></a> •
    <a href="#-advanced-usage"><strong>Advanced</strong></a> •
    <a href="#-documentation"><strong>Documentation</strong></a> •
    <a href="#-examples"><strong>Examples</strong></a> •
    <a href="#-benchmarks"><strong>Benchmarks</strong></a> •
    <a href="#-troubleshooting"><strong>Troubleshooting</strong></a> •
    <a href="#-community--support"><strong>Support</strong></a> •
    <a href="#-roadmap"><strong>Roadmap</strong></a> •
    <a href="#-security"><strong>Security</strong></a> •
    <a href="#-license"><strong>License</strong></a>
  </p>
  
  <p>
    <a href="https://github.com/guglxni/fedzk/releases">
      <img src="https://img.shields.io/github/v/release/guglxni/fedzk?style=flat-square&label=Release&color=blue" alt="GitHub release">
    </a>
    <a href="https://github.com/guglxni/fedzk/blob/main/LICENSE">
      <img src="https://img.shields.io/badge/License-MIT-green.svg?style=flat-square" alt="License">
    </a>
    <a href="https://github.com/guglxni/fedzk/stargazers">
      <img src="https://img.shields.io/github/stars/guglxni/fedzk?style=flat-square" alt="Stars">
    </a>
    <a href="https://github.com/guglxni/fedzk/network/members">
      <img src="https://img.shields.io/github/forks/guglxni/fedzk?style=flat-square" alt="Forks">
    </a>
    <a href="https://github.com/guglxni/fedzk/issues">
      <img src="https://img.shields.io/github/issues/guglxni/fedzk?style=flat-square" alt="Issues">
    </a>
    <a href="https://pypi.org/project/fedzk/">
      <img src="https://img.shields.io/pypi/v/fedzk?style=flat-square" alt="PyPI">
    </a>
    <img src="https://img.shields.io/badge/Version-1.0.0-blue?style=flat-square" alt="Version">
  </p>
</div>

## 📖 Project Overview

FEDzk is a cutting-edge framework that integrates federated learning with zero-knowledge proofs to address privacy and security concerns in distributed machine learning. Traditional federated learning systems face challenges with respect to verifiability and trust; our framework solves these issues by providing cryptographic guarantees for model update integrity.

### Key Differentiators

- **Provable Security**: Unlike conventional federated learning frameworks, FEDzk provides mathematical guarantees for the integrity of model updates
- **Privacy by Design**: Client data never leaves local environments, preserving privacy while still enabling collaborative learning
- **Tamper-Resistant**: Zero-knowledge proofs make it computationally infeasible to submit malicious updates
- **Scalable Architecture**: Designed to scale from small research deployments to production-grade distributed systems

### Use Cases

- **Healthcare**: Privacy-preserving machine learning across multiple hospitals or clinics
- **Finance**: Fraud detection models trained across multiple financial institutions
- **IoT Networks**: Distributed learning across edge devices with limited computational resources
- **Multi-party Collaborations**: Research or industry collaborations where data privacy is critical

## 🚀 Features

- **Privacy-Preserving**: Secure federated learning with strong privacy guarantees
- **Zero-Knowledge Proofs**: Verify model updates without revealing sensitive data
- **Distributed Training**: Coordinate training across multiple clients
- **Benchmarking Tools**: Evaluate performance and scalability
- **Secure Aggregation**: MPC server for secure model aggregation
- **Customizable**: Adapt to different ML models and datasets
- **Fault Tolerance**: Resilient to node failures during distributed training
- **Versioned Models**: Track model evolution across training rounds
- **Model Compression**: Reduce communication overhead in distributed settings
- **Differential Privacy**: Additional privacy guarantees through noise addition

## 🏗️ Architecture

The FEDzk framework consists of three main components:

```
┌────────────────┐     ┌─────────────────┐     ┌───────────────┐
│                │     │                 │     │               │
│  Client Node   │────▶│   Coordinator   │◀────│  Client Node  │
│  (Training)    │     │  (Aggregation)  │     │  (Training)   │
│                │     │                 │     │               │
└────────┬───────┘     └────────┬────────┘     └───────┬───────┘
         │                      │                      │
         │                      ▼                      │
         │              ┌───────────────┐              │
         └─────────────▶│   ZK Proofs   │◀─────────────┘
                        │ (Verification) │
                        └───────────────┘
```

### Workflow Diagram

```
┌──────────┐  1. Local Training   ┌───────────┐
│          │──────────────────────▶           │
│  Client  │                      │   Model   │
│          │◀──────────────────────           │
└────┬─────┘  2. Model Updates    └─────┬─────┘
     │                                  │
     │        3. Generate ZK Proof      │
     ▼                                  ▼
┌──────────┐  4. Submit Updates   ┌───────────┐
│          │  with Proof          │           │
│  Prover  │──────────────────────▶  Verifier │
│          │                      │           │
└──────────┘                      └─────┬─────┘
                                        │
                                        │
                                        ▼
                                  ┌───────────┐
                                  │           │
                                  │Coordinator│
                                  │           │
                                  └───────────┘
                                  5. Aggregate
                                     Models
```

### Component Details

- **Client Node**: Responsible for local model training on private data
- **Prover**: Generates zero-knowledge proofs for model updates
- **Verifier**: Validates proofs before accepting model updates
- **Coordinator**: Aggregates verified model updates and distributes the global model
- **MPC Server**: Enables secure multi-party computation for additional privacy guarantees

## 💻 System Requirements

### Minimum Requirements

- **Python**: 3.8 or higher
- **RAM**: 4GB (8GB recommended for larger models)
- **Storage**: 1GB free space
- **Processor**: Dual-core CPU (quad-core recommended)
- **OS**: Linux, macOS, or Windows

### Dependencies

- PyTorch (1.8+)
- NumPy
- cryptography
- circom (for circuit compilation)
- snarkjs (for zero-knowledge proof generation)

### For Production Deployments

- **RAM**: 16GB or higher
- **Processor**: 8+ CPU cores
- **GPU**: Recommended for faster proof generation
- **Network**: High-bandwidth, low-latency connections between nodes

## 💻 Installation

### From PyPI (Recommended)

```bash
# Install from PyPI
pip install fedzk

# With optional dependencies
pip install fedzk[all]     # All dependencies
pip install fedzk[dev]     # Development tools
pip install fedzk[docs]    # Documentation generation
```

### From Source

```bash
# Clone the repository
git clone https://github.com/guglxni/fedzk.git
cd fedzk

# Create and activate a virtual environment
python -m venv .venv
source .venv/bin/activate  # On Windows: .venv\Scripts\activate

# Install the package
pip install -e .
```

### Docker Installation

```bash
# Build the Docker image
docker build -t fedzk:latest .

# Run the container
docker run -it --rm fedzk:latest
```

## 🚦 Quick Start

### Basic Usage

```python
from fedzk.client import Trainer
from fedzk.coordinator import Aggregator

# Initialize a trainer with your model configuration
trainer = Trainer(model_config={
    'architecture': 'mlp',
    'layers': [784, 128, 10],
    'activation': 'relu'
})

# Train locally on your data
updates = trainer.train(data, epochs=5)

# Generate zero-knowledge proof for model updates
proof = trainer.generate_proof(updates)

# Submit updates with proof to coordinator
coordinator = Aggregator()
coordinator.submit_update(updates, proof)
```

### Verification Process

```python
from fedzk.prover import Verifier

# Initialize the verifier
verifier = Verifier()

# Verify the proof
is_valid = verifier.verify(proof, public_inputs)

if is_valid:
    print("✅ Model update verified successfully!")
else:
    print("❌ Verification failed. Update rejected.")
```

## 🔧 Advanced Usage

### Custom Circuit Integration

FEDzk allows you to define custom verification circuits:

```python
from fedzk.prover import CircuitBuilder

# Define a custom verification circuit
circuit_builder = CircuitBuilder()
circuit_builder.add_constraint("model_update <= threshold")
circuit_builder.add_constraint("norm(weights) > 0")

# Compile the circuit
circuit_path = circuit_builder.compile("my_custom_circuit")

# Use the custom circuit for verification
trainer.set_circuit(circuit_path)
```

### Distributed Deployment

To deploy across multiple nodes:

```python
from fedzk.coordinator import ServerConfig
from fedzk.mpc import SecureAggregator

# Configure the coordinator server
config = ServerConfig(
    host="0.0.0.0",
    port=8000,
    min_clients=5,
    aggregation_threshold=3,
    timeout=120
)

# Initialize and start the coordinator
coordinator = Aggregator(config)
coordinator.start()

# Set up secure aggregation
secure_agg = SecureAggregator(
    privacy_budget=0.1,
    encryption_key="shared_secret",
    mpc_protocol="semi_honest"
)
coordinator.set_aggregator(secure_agg)
```

### Performance Optimization

```python
from fedzk.client import OptimizedTrainer
from fedzk.benchmark import Profiler

# Create an optimized trainer with hardware acceleration
trainer = OptimizedTrainer(
    use_gpu=True,
    precision="mixed",
    batch_size=64,
    parallel_workers=4
)

# Profile the training and proof generation
profiler = Profiler()
with profiler.profile():
    updates = trainer.train(data)
    proof = trainer.generate_proof(updates)

# Get performance insights
profiler.report()
```

## 📚 Documentation

For more detailed documentation, examples, and API references, please refer to:

- [Getting Started Guide](docs/getting_started.md)
- [API Documentation](docs/api_reference.md)
- [Architecture Overview](docs/architecture.md)
- [Implementation Details](docs/implementation_details.md)
- [Zero-Knowledge Proofs](docs/zk_proofs.md)
- [Security Considerations](docs/legal/SECURITY.md)
- [Performance Tuning](docs/performance.md)
- [Deployment Guide](docs/deployment_guide.md)
- [Contribution Guidelines](docs/CONTRIBUTING.md)

## 📋 Examples

The [examples](examples) directory contains sample code and deployment configurations:

- [Basic Training](examples/basic_training.py): Simple federated learning setup
- [Distributed Deployment](examples/distributed_deployment.py): Multi-node configuration
- [Docker Deployment](examples/Dockerfile): Containerized deployment
- [Custom Circuits](examples/custom_circuits.py): Creating custom verification circuits
- [Secure MPC](examples/secure_mpc.py): Multi-party computation integration
- [Differential Privacy](examples/differential_privacy.py): Adding differential privacy
- [Model Compression](examples/model_compression.py): Reducing communication overhead

## 📊 Benchmarks

FEDzk has been benchmarked on multiple datasets:

| Dataset  | Clients | Rounds | Accuracy | Proof Generation Time | Verification Time |
|----------|---------|--------|----------|----------------------|-------------------|
| MNIST    | 10      | 5      | 97.8%    | 0.504s               | 0.204s            |
| CIFAR-10 | 20      | 50     | 85.6%    | 0.503s               | 0.204s            |
| IMDb     | 8       | 15     | 86.7%    | 0.2s                 | 0.1s              |
| Reuters  | 12      | 25     | 92.3%    | 0.3s                 | 0.1s              |

### Performance Across Hardware

Verified benchmark results on current hardware:

| Hardware | Specification |
|----------|---------------|
| CPU | Apple M4 Pro (12 cores) |
| RAM | 24.0 GB |
| GPU | Apple M4 Integrated GPU (MPS) |

> **Note**: Benchmarks use real zero-knowledge proofs when the ZK infrastructure is available, otherwise they fall back to a realistic simulation that accurately models the computational complexity of proof generation and verification. Run `./fedzk/scripts/setup_zk.sh` to set up the ZK environment for real proof benchmarks.

Benchmark methodology: Measurements taken on CIFAR-10 dataset with a CNN model containing approximately 5M parameters. Batch size of 32 was used for all experiments.

## ❓ Troubleshooting

### Common Issues

#### Installation Problems

**Issue**: Error installing cryptographic dependencies  
**Solution**: Ensure you have the required system libraries:
```bash
# On Ubuntu/Debian
sudo apt-get install build-essential libssl-dev libffi-dev python3-dev

# On macOS
brew install openssl
```

#### Runtime Errors

**Issue**: "Circuit compilation failed"  
**Solution**: Check that Circom is properly installed and in your PATH:
```bash
circom --version
# If not found, install with: npm install -g circom
```

**Issue**: Memory errors during proof generation  
**Solution**: Reduce the model size or increase available memory:
```python
trainer = Trainer(model_config={
    'architecture': 'mlp',
    'layers': [784, 64, 10],  # Smaller hidden layer
})
```

### Debugging Tools

FEDzk provides several debugging utilities:

```python
from fedzk.debug import CircuitDebugger, ProofInspector

# Debug a circuit
debugger = CircuitDebugger("model_update.circom")
debugger.trace_constraints()

# Inspect a generated proof
inspector = ProofInspector(proof_file="proof.json")
inspector.validate_structure()
inspector.analyze_complexity()
```

## 👥 Community & Support

- **GitHub Issues**: For bug reports and feature requests
- **Discussions**: For general questions and community discussions
- **Slack Channel**: Join our [Slack workspace](https://fedzk-community.slack.com) for real-time support
- **Mailing List**: Subscribe to our [mailing list](https://groups.google.com/g/fedzk-users) for announcements

### Getting Help

If you encounter issues not covered in the documentation:

1. Check the [Troubleshooting Guide](docs/troubleshooting.md)
2. Search existing [GitHub Issues](https://github.com/guglxni/fedzk/issues)
3. Ask in the community channels
4. If the issue persists, [file a detailed bug report](https://github.com/guglxni/fedzk/issues/new/choose)

## 🗺️ Roadmap

See our [detailed roadmap](ROADMAP.md) for planned features and improvements.

### Upcoming Features

- **Q1 2025**: Enhanced circuit library for common ML models
- **Q2 2025**: Improved GPU acceleration for proof generation
- **Q3 2025**: WebAssembly support for browser-based clients
- **Q4 2025**: Integration with popular ML frameworks (TensorFlow, JAX)
- **Q1 2026**: Formal security analysis and certification

## 📝 Changelog

See the [CHANGELOG.md](CHANGELOG.md) for a detailed history of changes.

## 📄 Citation

If you use FEDzk in your research, please cite:

```bibtex
@software{fedzk2025,
  author = {Guglani, Aaryan},
  title = {FEDzk: Federated Learning with Zero-Knowledge Proofs},
  year = {2025},
  url = {https://github.com/guglxni/fedzk},
}
```

## 🔒 Security

We take security seriously. Please review our [security policy](docs/legal/SECURITY.md) for reporting vulnerabilities.

### Security Features

- **End-to-End Encryption**: All communication between nodes is encrypted
- **Zero-Knowledge Proofs**: Ensures model update integrity without revealing sensitive data
- **Differential Privacy**: Optional noise addition to prevent inference attacks
- **Secure Aggregation**: MPC-based techniques to protect individual updates
- **Input Validation**: Extensive validation to prevent injection attacks

## 📄 License

This project is licensed under the [MIT License](LICENSE).

Copyright (c) 2025 Aaryan Guglani and FEDzk Contributors

## 🤝 Contributing

We welcome contributions from the community! Please check out our [contributing guidelines](docs/CONTRIBUTING.md) to get started.

## Project Structure

The FEDzk project follows a standard Python package structure:

- `src/fedzk/` - Main Python package
- `tests/` - Test suite
- `docs/` - Documentation
- `examples/` - Usage examples

For a detailed overview of the project organization, please see [Project Structure Documentation](docs/project_structure.md).