OLALa: Online Learned Adaptive Lattice Codes for Heterogeneous Federated Learning

TL;DR

OLALa introduces an adaptive lattice quantization scheme for federated learning, allowing clients to dynamically optimize their model update compression, which enhances convergence and reduces communication costs in heterogeneous environments.

Contribution

This work develops a novel online learning framework for adaptive lattice quantization in federated learning, with proven convergence guarantees and improved performance over fixed schemes.

Findings

OLALa outperforms fixed-codebook schemes in various settings.

Adaptive lattices improve convergence bounds.

Clients can effectively tune quantizers online with minimal communication.

Abstract

Federated learning (FL) enables collaborative training across distributed clients without sharing raw data, often at the cost of substantial communication overhead induced by transmitting high-dimensional model updates. This overhead can be alleviated by having the clients quantize their model updates, with dithered lattice quantizers identified as an attractive scheme due to its structural simplicity and convergence-preserving properties. However, existing lattice-based FL schemes typically rely on a fixed quantization rule, which is suboptimal in heterogeneous and dynamic environments where the model updates distribution varies across users and training rounds. In this work, we propose Online Learned Adaptive Lattices (OLALa), a heterogeneous FL framework where each client can adjust its quantizer online using lightweight local computations. We first derive convergence guarantees for FL with non-fixed lattice quantizers and show that proper lattice adaptation can tighten the convergence bound. Then, we design an online learning algorithm that enables clients to tune their quantizers throughout the FL process while exchanging only a compact set of quantization parameters. Numerical experiments demonstrate that OLALa consistently improves learning performance under various quantization rates, outperforming conventional fixed-codebook and non-adaptive schemes.