GQ-FSL: Green Quantized Federated Split Learning

Overview§

GQ-FSL addresses the energy bottleneck of deploying DNNs on wireless edge devices by combining split learning (SL) with stochastic quantization during both local training and transmission. Unlike standard quantized federated learning (QFL), GQ-FSL allows asymmetric precision levels for the client-side and server-side submodels, decoupling device energy constraints from global convergence degradation. The framework jointly optimizes the DNN split point and precision levels to minimize total energy consumption under a strict accuracy guarantee.

Method & Core Math§

System Model§

In each round, a subset of clients trains a client-side submodel up to a cut layer, quantizes the cut-layer activations (and gradients during backprop), and transmits them to the edge server. The server processes the remaining layers, computes the loss, and sends quantized gradients back. Parameter updates are aggregated federated-style.

Quantization§

Stochastic quantization maps a floating-point tensor $x$ to $b$-bit values:

$$ Q(x) = \Delta \cdot \text{sign}(x) \cdot \left\lfloor \frac{|x|}{\Delta} + u \right\rfloor, \quad \Delta = \frac{2^{1-b}}{2^b-1} $$

where $u \sim U[0,1]$ provides unbiasedness. The client-side precision $b_c$ and server-side precision $b_s$ can differ, yielding the asymmetric setting.

Energy Model & Optimization§

Parameterized energy models account for computation, transmission, and idle energy at both client and server. The total expected energy $E_{\text{tot}} = E_{\text{client}} + E_{\text{server}}$ is a function of the split point $s$ and precisions $(b_c, b_s)$. The optimization problem is:

$$ \min_{s, b_c, b_s} \; E_{\text{tot}}(s, b_c, b_s) \quad \text{s.t.} \; \mathcal{A}(s, b_c, b_s) \ge \mathcal{A}_{\text{target}}, $$

where $\mathcal{A}$ is the test accuracy, estimated via a convergence bound derived under statistical heterogeneity. The authors derive a convergence bound for non-convex loss with quantization error and heterogeneous data distributions, showing the trade-off between energy and accuracy.

Convergence Bound§

The bound takes the form:

$$ \mathbb{E}[\|\nabla L(\mathbf{w})\|^2] \le \mathcal{O}\left( \frac{1}{\sqrt{T}} \right) + \mathcal{O}\left( b_c^{-2} + b_s^{-2} \right) + \mathcal{O}(\Gamma), $$

where $\Gamma$ measures data heterogeneity. This explicitly shows that client precision $b_c$ affects energy directly but convergence only through quantization variance, enabling aggressive down-quantization on device without catastrophic accuracy loss.

Joint Optimization Algorithm§

Because the search space is small (discrete split points and precision levels), GQ-FSL uses a grid search with early stopping, or a coordinate-descent style approach, to find the configuration that satisfies the accuracy constraint while minimizing energy.

Implementation Sketch§

import torch
import torch.nn as nn

class ClientSubmodel(nn.Module):
    def __init__(self, backbone, split_layer):
        super().__init__()
        # layers up to split_layer
        self.layers = backbone[:split_layer]

    def forward(self, x):
        return self.layers(x)

class ServerSubmodel(nn.Module):
    def __init__(self, backbone, split_layer):
        super().__init__()
        self.layers = backbone[split_layer:]

    def forward(self, z, targets):
        logits = self.layers(z)
        loss = nn.CrossEntropyLoss()(logits, targets)
        return logits, loss

def stochastic_quantize(tensor, bits):
    levels = 2**bits - 1
    abs_max = tensor.abs().max().item()
    delta = abs_max / levels
    noise = torch.empty_like(tensor).uniform_(0, delta)
    q = (tensor.abs() + noise).floor() * delta  # simplified
    return q * tensor.sign()

# Training loop sketch (per round, client-side)
client_model.train()
for x, y in local_batch:
    z = client_model(x)
    z_q = stochastic_quantize(z, b_c)
    # send z_q to server (simulate server forward/backward)
    grad_z_q = server.backward(z_q, y)  # returns \partial L / \partial z_q
    # dequantize approximation: use straight-through estimator
    grad_z = grad_z_q * (z.abs() < 2 * delta).float()
    z.backward(grad_z)
    optimizer.step()

The code snippet captures the core idea: client computes up to the split, quantizes activations, server processes the rest, and gradients are backpropagated with a straight-through estimator.

Results & Takeaways§

Experiments on CIFAR-10/100 and Tiny ImageNet with various DNNs show GQ-FSL reduces total energy by up to 5.5x compared to full-precision FSL and 2.3x versus quantized federated learning, while maintaining target accuracy. The asymmetric precision setup is key: lowering client precision to 2–4 bits cuts device energy dramatically with minimal accuracy degradation, since the server-side submodel remains at higher precision.

Why It Matters§

GQ-FSL offers a practical recipe for deploying large DNNs on battery-powered IoT devices: split the network, aggressively quantize only the client side, and let the server compensate. It also provides a principled framework for jointly choosing split points and bit-widths under resource constraints.