TailOPT: Efficient Distributed Optimization under Heavy-Tailed Noise

TailOPT is a framework for distributed training under heavy-tailed gradient noise, which is especially pronounced in attention-based models. It combines adaptive optimization with novel clipping techniques and comes with convergence guarantees under potentially unbounded gradient variance and local updates. Among its instantiations, Bi²Clip performs coordinate-wise clipping from above and below at both the inner and outer optimizers, capturing Adam-like benefits without the memory or communication cost of maintaining extra gradient statistics.

.
├── train_lm.py                   # Federated GLUE fine-tuning driver (includes Bi²Clip / BiClipSGD)
├── run_llm.py                    # Language-model experiment driver
├── shakespeare.py                # Shakespeare / book-corpus federated language modeling
├── synthetic.py                  # Synthetic (linear regression) heavy-tailed experiments
├── main_fed_conf_int.py          # Image-benchmark driver for confidence-interval runs
├── main_fed_eff_dataprocess.py   # Image-benchmark driver with cached data preprocessing
├── compute_privacy.py            # Differential privacy accounting entry point
├── data/                         # Dataset wrappers, GLUE loaders, LDA partitioning
├── models/                       # Local update logic, clipping optimizers, aggregation, evaluation
│   ├── ClipUpdate.py             #   BiClip / clipping optimizers and local training loops
│   ├── Update.py                 #   Base local-update module
│   ├── Fed.py                    #   Server-side aggregation rules
│   └── Nets.py, test.py, ...     #   Models and evaluation utilities
├── privacy_analysis/             # RDP accountant (vendored from TensorFlow Privacy, Apache 2.0)
├── utils/                        # Argument parsing and client sampling
├── preprocessing/                # Dataset preparation (Shakespeare, book corpus, clusters, GLDv2)
├── scripts/                      # SLURM sweep scripts and results aggregation
├── books/                        # Book corpus: public-domain texts (see provenance note below)
└── book_json/                    # Federated train/test splits derived from books/

The code is written in Python with PyTorch and Hugging Face Transformers. The main dependencies are:

torch  transformers  datasets  evaluate  numpy  wandb  nltk  scikit-learn

Experiments are tracked with Weights & Biases; run wandb login first (metrics are logged to your default entity).

The primary driver is train_lm.py, which fine-tunes RoBERTa on GLUE tasks across federated clients with configurable inner/outer optimizers, including BiClipSGD_Full (Bi²Clip):

python train_lm.py --model roberta --dataset glue --glue_task sst2 --num_users 10

Hyperparameter sweeps used in the paper are launched with the SLURM array scripts in scripts/ (e.g. noniid.sh, q_llm_bisquare.sh). Fill in the placeholder log/venv/partition values for your cluster and submit from the repository root.

shakespeare.py runs federated language modeling over the LEAF Shakespeare split or the included book corpus. To regenerate the data:

• Shakespeare: obtain the raw data via LEAF, then run preprocessing/preprocess_custom_shakespeare.py.
• Book corpus: preprocessing/process_books.py converts the texts in books/ into the per-client JSON splits in book_json/ (one client per book).

Data provenance: the texts in books/ are public-domain classic literature obtained from Project Gutenberg; book_json/ is derived from them.

• synthetic.py runs the controlled heavy-tailed linear-regression experiments.
• main_fed_conf_int.py / main_fed_eff_dataprocess.py run image benchmarks (MNIST, CIFAR) with cached preprocessing; aggregate multi-seed results with scripts/aggregate_results.py.

If you find this code or our paper helpful in your research, please consider citing:

@InProceedings{pmlr-v267-lee25ak,
  title     = {Efficient Distributed Optimization under Heavy-Tailed Noise},
  author    = {Lee, Su Hyeong and Zaheer, Manzil and Li, Tian},
  booktitle = {Proceedings of the 42nd International Conference on Machine Learning},
  pages     = {33833--33882},
  year      = {2025},
  volume    = {267},
  series    = {Proceedings of Machine Learning Research},
  publisher = {PMLR},
  url       = {https://proceedings.mlr.press/v267/lee25ak.html}
}