← Search

Martin Jaggi

91 accepted papers

2026

Beyond URLs: Metadata Diversity and Position for Efficient LLM Pretraining

ICLR 2026poster

Incorporating metadata in Large Language Models (LLMs) pretraining has recently emerged as a promising approach to accelerate training. However prior work highlighted only one useful signal—URLs, leaving open the question of whether other forms of metadata could yield greater benefits. In this study…

Cited by 0SourceScholar
2026

Gradient-Normalized Smoothness for Optimization with Approximate Hessians

ICLR 2026poster

In this work, we develop new optimization algorithms that use approximate second-order information combined with the gradient regularization technique to achieve fast global convergence rates for both convex and non-convex objectives. The key innovation of our analysis is a novel notion called Gradi…

Cited by 0SourcecodeScholar
2026

Weight Decay may matter more than µP for Learning Rate Transfer in Practice

ICLR 2026poster

Transferring the optimal learning rate from small to large neural networks can enable efficient training at scales where hyperparameter tuning is otherwise prohibitively expensive. To this end, the Maximal Update Parameterization (µP) proposes a learning rate scaling designed to keep the update dyna…

Cited by 0SourcecodeScholar
2025

Attention with Markov: A Curious Case of Single-layer Transformers

ICLR 2025spotlight

Attention-based transformers have achieved tremendous success across a variety of disciplines including natural languages. To deepen our understanding of their sequential modeling capabilities, there is a growing interest in using Markov input processes to study them. A key finding is that when trai…

Cited by 3SourcePDFScholar
2025

CoTFormer: A Chain of Thought Driven Architecture with Budget-Adaptive Computation Cost at Inference

ICLR 2025poster

Scaling language models to larger and deeper sizes has led to significant boosts in performance. Even though the size of these models limits their application in compute-constrained environments, the race to continually develop ever larger and deeper foundational models is underway. At the same time…

Cited by 1SourcePDFScholar
2025

Effective Interplay between Sparsity and Quantization: From Theory to Practice

ICLR 2025spotlight

The increasing size of deep neural networks (DNNs) necessitates effective model compression to reduce their computational and memory footprints. Sparsity and quantization are two prominent compression methods that have been shown to reduce DNNs' computational and memory footprints significantly whil…

Cited by 5SourcePDFScholar
2025

Enhancing Multilingual LLM Pretraining with Model-Based Data Selection

NeurIPS 2025poster

Dataset curation has become a basis for strong large language model (LLM) performance. While various rule-based filtering heuristics exist for English and multilingual datasets, model-based filtering techniques have primarily focused on English. To address the disparity stemming from limited researc…

Cited by 0SourceScholar
2025

GRAPE: Optimize Data Mixture for Group Robust Multi-target Adaptive Pretraining

NeurIPS 2025poster

The performance of large language models (LLMs) across diverse downstream applications is fundamentally governed by the quality and composition of their pretraining corpora. Existing domain reweighting algorithms primarily optimize data mixtures for a single target task, thereby resulting in models…

Cited by 4SourceScholar
2025

Intrinsic User-Centric Interpretability through Global Mixture of Experts

ICLR 2025poster

In human-centric settings like education or healthcare, model accuracy and model explainability are key factors for user adoption. Towards these two goals, intrinsically interpretable deep learning models have gained popularity, focusing on accurate predictions alongside faithful explanations. Howev…

2025

On-Device Collaborative Language Modeling via a Mixture of Generalists and Specialists

ICML 2025poster

On-device LLMs have gained increasing attention for their ability to enhance privacy and provide a personalized user experience. To facilitate private learning with scarce data, Federated Learning has become a standard approach. However, it faces challenges such as computational resource heterogenei…

2024

Analyzing & Reducing the Need for Learning Rate Warmup in GPT Training

NeurIPS 2024poster

Learning Rate Warmup is a popular heuristic for training neural networks, especially at larger batch sizes, despite limited understanding of its benefits. Warmup decreases the update size $\Delta \mathbf{w}_t = \eta_t \mathbf{u}_t$ early in training by using lower values for the learning rate $\eta_…

Cited by 3SourcePDFScholar
2024

DenseFormer: Enhancing Information Flow in Transformers via Depth Weighted Averaging

NeurIPS 2024poster

The transformer architecture by Vaswani et al. (2017) is now ubiquitous across application domains, from natural language processing to speech processing and image understanding. We propose DenseFormer, a simple modification to the standard architecture that improves the perplexity of the model with…

Cited by 6SourcePDFScholar
2024

LASER: Linear Compression in Wireless Distributed Optimization

ICML 2024poster

Data-parallel SGD is the de facto algorithm for distributed optimization, especially for large scale machine learning. Despite its merits, communication bottleneck is one of its persistent issues. Most compression schemes to alleviate this either assume noiseless communication links, or fail to achi…

Cited by 7SourcePDFScholar
2024

Layer-wise linear mode connectivity

ICLR 2024poster

Averaging neural network parameters is an intuitive method for fusing the knowledge of two independent models. It is most prominently used in federated learning. If models are averaged at the end of training, this can only lead to a good performing model if the loss surface of interest is very parti…

2024

On Convergence of Incremental Gradient for Non-convex Smooth Functions

ICML 2024poster

In machine learning and neural network optimization, algorithms like incremental gradient, single shuffle SGD, and random reshuffle SGD are popular due to their cache-mismatch efficiency and good practical convergence behavior. However, their optimization properties in theory, especially for non-con…

Cited by 1SourcePDFScholar
2024

QuaRot: Outlier-Free 4-Bit Inference in Rotated LLMs

NeurIPS 2024poster

We introduce QuaRot, a new Quantization scheme based on Rotations, which is able to quantize LLMs end-to-end, including all weights, activations, and KV cache in 4 bits. QuaRot rotates LLMs in a way that removes outliers from the hidden state without changing the output, making quantization easier.…

2024

Rotational Equilibrium: How Weight Decay Balances Learning Across Neural Networks

ICML 2024poster

This study investigates how weight decay affects the update behavior of individual neurons in deep neural networks through a combination of applied analysis and experimentation. Weight decay can cause the expected magnitude and angular updates of a neuron's weight vector to converge to a steady stat…

Cited by 10SourcePDFScholar
2024

Scaling Laws and Compute-Optimal Training Beyond Fixed Training Durations

NeurIPS 2024spotlight

Scale has become a main ingredient in obtaining strong machine learning models. As a result, understanding a model's scaling properties is key to effectively designing both the right training setup as well as future generations of architectures. In this work, we argue that scale and training researc…

2024

Spectral Preconditioning for Gradient Methods on Graded Non-convex Functions

ICML 2024poster

The performance of optimization methods is often tied to the spectrum of the objective Hessian. Yet, conventional assumptions, such as smoothness, do often not enable us to make finely-grained convergence statements—particularly not for non-convex problems. Striving for a more intricate characteriza…

Cited by 10SourcePDFScholar
2024

The Privacy Power of Correlated Noise in Decentralized Learning

ICML 2024poster

Decentralized learning is appealing as it enables the scalable usage of large amounts of distributed data and resources without resorting to any central entity, while promoting privacy since every user minimizes the direct exposure of their data. Yet, without additional precautions, curious users ca…

2023

Agree to Disagree: Diversity through Disagreement for Better Transferability

ICLR 2023top-5%

Gradient-based learning algorithms have an implicit \emph{simplicity bias} which in effect can limit the diversity of predictors being sampled by the learning procedure. This behavior can hinder the transferability of trained models by (i) favoring the learning of simpler but spurious features --- p…

2023

Fast Attention Over Long Sequences With Dynamic Sparse Flash Attention

NeurIPS 2023poster

Transformer-based language models have found many diverse applications requiring them to process sequences of increasing length. For these applications, the causal self-attention---which is the only component scaling quadratically w.r.t. the sequence length---becomes a central concern. While many wo…

Cited by 12SourcePDFScholar
2023

MultiMoDN—Multimodal, Multi-Task, Interpretable Modular Networks

NeurIPS 2023poster

Predicting multiple real-world tasks in a single model often requires a particularly diverse feature space. Multimodal (MM) models aim to extract the synergistic predictive potential of multiple data types to create a shared feature space with aligned semantic meaning across inputs of drastically va…

2023

SIMSUM: Document-level Text Simplification via Simultaneous Summarization

ACL 2023long

Document-level text simplification is a specific type of simplification which involves simplifying documents consisting of several sentences by rewriting them into fewer or more sentences. In this paper, we propose a new two-stage framework SIMSUM for automated document-level text simplification. Ou…

2023

Special Properties of Gradient Descent with Large Learning Rates

ICML 2023poster

When training neural networks, it has been widely observed that a large step size is essential in stochastic gradient descent (SGD) for obtaining superior models. However, the effect of large step sizes on the success of SGD is not well understood theoretically. Several previous works have attribute…

Cited by 10SourcePDFScholar
2022

Beyond spectral gap: the role of the topology in decentralized learning

NeurIPS 2022accept

In data-parallel optimization of machine learning models, workers collaborate to improve their estimates of the model: more accurate gradients allow them to use larger learning rates and optimize faster. We consider the setting in which all workers sample from the same dataset, and communicate over…

2022

Byzantine-Robust Learning on Heterogeneous Datasets via Bucketing

ICLR 2022spotlight

In Byzantine robust distributed or federated learning, a central server wants to train a machine learning model over data distributed across multiple workers. However, a fraction of these workers may deviate from the prescribed algorithm and send arbitrary messages. While this problem has received s…

2022

FLamby: Datasets and Benchmarks for Cross-Silo Federated Learning in Realistic Healthcare Settings

NeurIPS 2022accept

Federated Learning (FL) is a novel approach enabling several clients holding sensitive data to collaboratively train machine learning models, without centralizing data. The cross-silo FL setting corresponds to the case of few ($2$--$50$) reliable clients, each holding medium to large datasets, and i…

2022

Masked Training of Neural Networks with Partial Gradients

AISTATS 2022poster

State-of-the-art training algorithms for deep learning models are based on stochastic gradient descent (SGD). Recently, many variations have been explored: perturbing parameters for better accuracy (such as in Extragradient), limiting SGD updates to a subset of parameters for increased efficiency (s…

Cited by 32SourcePDFScholar
2022

SKILL: Structured Knowledge Infusion for Large Language Models

NAACL 2022long

Large language models (LLMs) have demonstrated human-level performance on a vast spectrum of natural language tasks. However, it is largely unexplored whether they can better internalize knowledge from a structured data, such as a knowledge graph, or from text. In this work, we propose a method to i…

Cited by 100SourcePDFScholar
2022

Sharper Convergence Guarantees for Asynchronous SGD for Distributed and Federated Learning

NeurIPS 2022accept

We study the asynchronous stochastic gradient descent algorithm, for distributed training over $n$ workers that might be heterogeneous. In this algorithm, workers compute stochastic gradients in parallel at their own pace and return them to the server without any synchronization. Existing convergen…

Cited by 106SourcePDFScholar
2021

A Linearly Convergent Algorithm for Decentralized Optimization: Sending Less Bits for Free!

AISTATS 2021poster

Decentralized optimization methods enable on-device training of machine learning models without a central coordinator. In many scenarios communication between devices is energy demanding and time consuming and forms the bottleneck of the entire system. We propose a new randomized first-order method…

Cited by 93SourcePDFScholar
2021

Breaking the centralized barrier for cross-device federated learning

NeurIPS 2021poster

Federated learning (FL) is a challenging setting for optimization due to the heterogeneity of the data across different clients which gives rise to the client drift phenomenon. In fact, obtaining an algorithm for FL which is uniformly better than simple centralized training has been a major open pro…

Cited by 102SourcePDFScholar
2021

Consensus Control for Decentralized Deep Learning

ICML 2021spotlight

Decentralized training of deep learning models enables on-device learning over networks, as well as efficient scaling to large compute clusters. Experiments in earlier works reveal that, even in a data-center setup, decentralized training often suffers from the degradation in the quality of the mode…

Cited by 100SourcePDFScholar
2021

Critical Parameters for Scalable Distributed Learning with Large Batches and Asynchronous Updates

AISTATS 2021poster

It has been experimentally observed that the efficiency of distributed training with stochastic gradient (SGD) depends decisively on the batch size and—in asynchronous implementations—on the gradient staleness. Especially, it has been observed that the speedup saturates beyond a certain batch size a…

Cited by 23SourcePDFScholar
2021

Exact Optimization of Conformal Predictors via Incremental and Decremental Learning

ICML 2021spotlight

Conformal Predictors (CP) are wrappers around ML models, providing error guarantees under weak assumptions on the data distribution. They are suitable for a wide range of problems, from classification and regression to anomaly detection. Unfortunately, their very high computational complexity limits…

2021

LENA: Communication-Efficient Distributed Learning with Self-Triggered Gradient Uploads

AISTATS 2021poster

In distributed optimization, parameter updates from the gradient computing node devices have to be aggregated in every iteration on the orchestrating server. When these updates are sent over an arbitrary commodity network, bandwidth and latency can be limiting factors. We propose a communication fra…

Cited by 27SourcePDFScholar
2021

Learning from History for Byzantine Robust Optimization

ICML 2021spotlight

Byzantine robustness has received significant attention recently given its importance for distributed and federated learning. In spite of this, we identify severe flaws in existing algorithms even when the data across the participants is identically distributed. First, we show realistic examples whe…

2021

Lightweight Cross-Lingual Sentence Representation Learning

ACL 2021long

Large-scale models for learning fixed-dimensional cross-lingual sentence representations like LASER (Artetxe and Schwenk, 2019b) lead to significant improvement in performance on downstream tasks. However, further increases and modifications based on such large-scale models are usually impractical d…

2021

Quasi-global Momentum: Accelerating Decentralized Deep Learning on Heterogeneous Data

ICML 2021spotlight

Decentralized training of deep learning models is a key element for enabling data privacy and on-device learning over networks. In realistic learning scenarios, the presence of heterogeneity across different clients’ local datasets poses an optimization challenge and may severely deteriorate the gen…

2021

RelaySum for Decentralized Deep Learning on Heterogeneous Data

NeurIPS 2021poster

In decentralized machine learning, workers compute model updates on their local data. Because the workers only communicate with few neighbors without central coordination, these updates propagate progressively over the network. This paradigm enables distributed training on networks without all-to-al…

2021

Semantic Perturbations With Normalizing Flows for Improved Generalization

ICCV 2021poster

Data augmentation is a widely adopted technique for avoiding overfitting when training deep neural networks. However, this approach requires domain-specific knowledge and is often limited to a fixed set of hard-coded transformations. Recently, several works proposed to use generative models for gene…

Cited by 14PDFcodeScholar
2021

Taming GANs with Lookahead-Minmax

ICLR 2021poster

Generative Adversarial Networks are notoriously challenging to train. The underlying minmax optimization is highly susceptible to the variance of the stochastic gradient and the rotational component of the associated game vector field. To tackle these challenges, we propose the Lookahead algorithm f…

2021

Understanding the effects of data parallelism and sparsity on neural network training

ICLR 2021poster

We study two factors in neural network training: data parallelism and sparsity; here, data parallelism means processing training data in parallel using distributed systems (or equivalently increasing batch size), so that training can be accelerated; for sparsity, we refer to pruning parameters in a…

Cited by 9SourcePDFScholar
2020

A Unified Theory of Decentralized SGD with Changing Topology and Local Updates

ICML 2020poster

Decentralized stochastic optimization methods have gained a lot of attention recently, mainly because of their cheap per iteration cost, data locality, and their communication-efficiency. In this paper we introduce a unified convergence analysis that covers a large variety of decentralized SGD metho…

Cited by 610SourcePDFScholar
2020

Context Mover’s Distance & Barycenters: Optimal Transport of Contexts for Building Representations

AISTATS 2020poster

We present a framework for building unsupervised representations of entities and their compositions, where each entity is viewed as a probability distribution rather than a vector embedding. In particular, this distribution is supported over the contexts which co-occur with the entity and are embedd…

2020

Decentralized Deep Learning with Arbitrary Communication Compression

ICLR 2020poster

Decentralized training of deep learning models is a key element for enabling data privacy and on-device learning over networks, as well as for efficient scaling to large compute clusters. As current approaches are limited by network bandwidth, we propose the use of communication compression in the d…

Cited by 280SourcecodeScholar
2020

Ensemble Distillation for Robust Model Fusion in Federated Learning

NeurIPS 2020poster

Federated Learning (FL) is a machine learning setting where many devices collaboratively train a machine learning model while keeping the training data decentralized. In most of the current training schemes the central model is refined by averaging the parameters of the server model and the updated…

2020

Evaluating The Search Phase of Neural Architecture Search

ICLR 2020poster

Neural Architecture Search (NAS) aims to facilitate the design of deep networks for new tasks. Existing techniques rely on two stages: searching over the architecture space and validating the best architecture. NAS algorithms are currently compared solely based on their results on the downstream tas…

Cited by 437SourcecodeScholar
2020

Linearly Convergent Frank-Wolfe with Backtracking Line-Search

AISTATS 2020poster

Structured constraints in Machine Learning have recently brought the Frank-Wolfe (FW) family of algorithms back in the spotlight. While the classical FW algorithm has poor local convergence properties, the Away-steps and Pairwise FW variants have emerged as improved variants with faster convergence.…

2020

On the Relationship between Self-Attention and Convolutional Layers

ICLR 2020poster

Recent trends of incorporating attention mechanisms in vision have led researchers to reconsider the supremacy of convolutional layers as a primary building block. Beyond helping CNNs to handle long-range dependencies, Ramachandran et al. (2019) showed that attention can completely replace convoluti…

Cited by 767SourcecodeScholar
2020

Optimizer Benchmarking Needs to Account for Hyperparameter Tuning

ICML 2020poster

The performance of optimizers, particularly in deep learning, depends considerably on their chosen hyperparameter configuration. The efficacy of optimizers is often studied under near-optimal problem-specific hyperparameters, and finding these settings may be prohibitively costly for practitioners.…

Cited by 62SourcePDFScholar
2020

Practical Low-Rank Communication Compression in Decentralized Deep Learning

NeurIPS 2020poster

Lossy gradient compression has become a practical tool to overcome the communication bottleneck in centrally coordinated distributed training of machine learning models. However, algorithms for decentralized training with compressed communication over arbitrary connected networks have been more comp…

2019

Decentralized Stochastic Optimization and Gossip Algorithms with Compressed Communication

ICML 2019oral

We consider decentralized stochastic optimization with the objective function (e.g. data samples for machine learning tasks) being distributed over n machines that can only communicate to their neighbors on a fixed communication graph. To address the communication bottleneck, the nodes compress (e.g…

Cited by 609SourcePDFScholar
2019

Efficient Greedy Coordinate Descent for Composite Problems

AISTATS 2019poster

Coordinate descent with random coordinate selection is the current state of the art for many large scale optimization problems. However, greedy selection of the steepest coordinate on smooth problems can yield convergence rates independent of the dimension $n$, requiring $n$ times fewer iterations.…

2019

Error Feedback Fixes SignSGD and other Gradient Compression Schemes

ICML 2019oral

Sign-based algorithms (e.g. signSGD) have been proposed as a biased gradient compression technique to alleviate the communication bottleneck in training large neural networks across multiple workers. We show simple convex counter-examples where signSGD does not converge to the optimum. Further, even…

2019

Overcoming Multi-model Forgetting

ICML 2019oral

We identify a phenomenon, which we refer to as multi-model forgetting, that occurs when sequentially training multiple deep networks with partially-shared parameters; the performance of previously-trained models degrades as one optimizes a subsequent one, due to the overwriting of shared parameters.…

2019

PowerSGD: Practical Low-Rank Gradient Compression for Distributed Optimization

NeurIPS 2019poster

We study gradient compression methods to alleviate the communication bottleneck in data-parallel distributed optimization. Despite the significant attention received, current compression schemes either do not scale well, or fail to achieve the target test accuracy. We propose a low-rank gradient com…

2019

Unsupervised Scalable Representation Learning for Multivariate Time Series

NeurIPS 2019poster

Time series constitute a challenging data type for machine learning algorithms, due to their highly variable lengths and sparse labeling in practice. In this paper, we tackle this challenge by proposing an unsupervised method to learn universal embeddings of time series. Unlike previous works, it is…

2018

A Distributed Second-Order Algorithm You Can Trust

ICML 2018oral

Due to the rapid growth of data and computational resources, distributed optimization has become an active research area in recent years. While first-order methods seem to dominate the field, second-order methods are nevertheless attractive as they potentially require fewer communication rounds to c…

2018

Adaptive balancing of gradient and update computation times using global geometry and approximate subproblems

AISTATS 2018poster

First-order optimization methods comprise two important primitives: i) the computation of gradient information and ii) the computation of the update that leads to the next iterate. In practice there is often a wide mismatch between the time required for the two steps, leading to underutilization of…

Cited by 0SourcePDFScholar
2018

On Matching Pursuit and Coordinate Descent

ICML 2018oral

Two popular examples of first-order optimization methods over linear spaces are coordinate descent and matching pursuit algorithms, with their randomized variants. While the former targets the optimization by moving along coordinates, the latter considers a generalized notion of directions. Exploiti…

Cited by 29SourcePDFScholar
2017

A Unified Optimization View on Generalized Matching Pursuit and Frank-Wolfe

AISTATS 2017poster

Two of the most fundamental prototypes of greedy optimization are the matching pursuit and Frank-Wolfe algorithms. In this paper we take a unified view on both classes of methods, leading to the first explicit convergence rates of matching pursuit methods in an optimization sense, for general sets o…

Cited by 71SourcePDFScholar
2017

Efficient Use of Limited-Memory Accelerators for Linear Learning on Heterogeneous Systems

NeurIPS 2017poster

We propose a generic algorithmic building block to accelerate training of machine learning models on heterogeneous compute systems. Our scheme allows to efficiently employ compute accelerators such as GPUs and FPGAs for the training of large-scale machine learning models, when the training data exc…

Cited by 18SourcePDFScholar
2017

Greedy Algorithms for Cone Constrained Optimization with Convergence Guarantees

NeurIPS 2017poster

Greedy optimization methods such as Matching Pursuit (MP) and Frank-Wolfe (FW) algorithms regained popularity in recent years due to their simplicity, effectiveness and theoretical guarantees. MP and FW address optimization over the linear span and the convex hull of a set of atoms, respectively. In…

Cited by 34SourcePDFScholar
2015

Adding vs. Averaging in Distributed Primal-Dual Optimization

ICML 2015poster

Distributed optimization methods for large-scale machine learning suffer from a communication bottleneck. It is difficult to reduce this bottleneck while still efficiently and accurately aggregating partial work from different machines. In this paper, we present a novel generalization of the recent…