Federated Learning & Privacy-Preserving AI: The Future of Distributed Training

package com.devmetrix.fl.server;

import org.springframework.stereotype.Service;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

@Service
public class FederatedAggregatorService {

    // Store updates for the current round
    private final Map<String, double[]> clientUpdates = new ConcurrentHashMap<>();

    // The current global model weights
    private double[] globalWeights;

    // Hyperparameters
    private final int EXPECTED_CLIENTS_PER_ROUND = 100;
    private final double LEARNING_RATE = 0.01;

    public FederatedAggregatorService() {
        // Initialize global weights (e.g., Xavier initialization)
        this.globalWeights = initializeWeights(1024); // 1024 parameters
    }

    /**
     * Called by REST Controller when a client pushes an update
     */
    public synchronized void submitUpdate(String clientId, double[] localWeights) {
        System.out.println("Received update from client: " + clientId);
        clientUpdates.put(clientId, localWeights);

        // Check if we have enough updates to close the round
        if (clientUpdates.size() >= EXPECTED_CLIENTS_PER_ROUND) {
            aggregateUpdates();
        }
    }

    /**
     * Performs FedAvg to compute new global weights
     */
    private void aggregateUpdates() {
        System.out.println("Aggregating " + clientUpdates.size() + " updates...");

        double[] newGlobalWeights = new double[globalWeights.length];
        int totalClients = clientUpdates.size();

        // FedAvg: Simple averaging (assuming equal weighting for simplicity)
        // In prod, you'd weight by sample size (n_k)
        for (double[] updates : clientUpdates.values()) {
            for (int i = 0; i < updates.length; i++) {
                newGlobalWeights[i] += updates[i];
            }
        }

        // Normalize
        for (int i = 0; i < newGlobalWeights.length; i++) {
            newGlobalWeights[i] /= totalClients;
        }

        // Update global state
        this.globalWeights = newGlobalWeights;

        // Reset for next round
        this.clientUpdates.clear();

        System.out.println("Round complete. New global model version available.");
        notifyClients(this.globalWeights);
    }

    private double[] initializeWeights(int size) {
        return new double[size]; // In reality, random values
    }

    public double[] getGlobalWeights() {
        return this.globalWeights;
    }
}

import * as tf from '@tensorflow/tfjs';

/**
 * Represents the FL Client running on the edge device
 */
class FederatedClient {
    model: tf.Sequential;
    serverId: string;

    constructor() {
        this.model = this.createModel();
        this.serverId = 'https://api.devmetrix.cloud/fl';
    }

    /**
     * Define the model architecture.
     * Must match the server's expectation.
     */
    createModel() {
        const model = tf.sequential();
        model.add(tf.layers.dense({units: 10, inputShape: [5], activation: 'relu'}));
        model.add(tf.layers.dense({units: 1, activation: 'sigmoid'})); // Binary classification
        model.compile({optimizer: 'sgd', loss: 'binaryCrossentropy'});
        return model;
    }

    /**
     * Orchestrates a single training round
     */
    async trainRound(globalWeights: tf.Tensor[], localData: tf.Tensor, localLabels: tf.Tensor) {
        console.log('Starting local training round...');

        // 1. Load global weights received from server
        // This syncs the local model with the global consensus
        this.model.setWeights(globalWeights);

        // 2. Train locally (e.g., 5 epochs)
        // We use a small batch size to fit in memory
        await this.model.fit(localData, localLabels, {
            epochs: 5,
            batchSize: 32,
            verbose: 0,
            callbacks: {
                onEpochEnd: (epoch, logs) => {
                    console.log(`Epoch ${epoch}: loss=${logs?.loss}`);
                }
            }
        });

        // 3. Extract new weights after training
        const localWeights = this.model.getWeights();

        // 4. Send back to server (Serialize tensors to arrays)
        const serializedWeights = await Promise.all(
            localWeights.map(w => w.data())
        );

        // Convert Float32Array to standard array for JSON
        const weightsPayload = serializedWeights.map(arr => Array.from(arr));

        await this.sendUpdateToServer(weightsPayload);

        // Cleanup memory
        tf.dispose([localData, localLabels]);
    }

    async sendUpdateToServer(weights: number[][]) {
        try {
            const response = await fetch(`${this.serverId}/update`, {
                method: 'POST',
                body: JSON.stringify({ weights }),
                headers: { 'Content-Type': 'application/json' }
            });

            if (response.ok) {
                console.log('Update successfully sent to server.');
            }
        } catch (err) {
            console.error('Failed to send update:', err);
        }
    }
}