Sharan Sahu

Selected publications

1. ICML

   On the Provable Suboptimality of Momentum SGD in Nonstationary Stochastic Optimization

   Sharan Sahu, Cameron J. Hogan, and Martin T. Wells

   In Proceedings of the 43rd International Conference on Machine Learning (ICML) 2026

   Abs arXiv Bib PDF Poster

   In this paper, we provide a comprehensive theoretical analysis of Stochastic Gradient Descent (SGD) and its momentum variants (Polyak Heavy-Ball and Nesterov) for tracking time-varying optima under strong convexity and smoothness. Our finite-time bounds reveal a sharp decomposition of tracking error into transient, noise-induced, and drift-induced components. This decomposition exposes a fundamental trade-off: while momentum is often used as a gradient-smoothing heuristic, under distribution shift it incurs an explicit drift-amplification penalty that diverges as the momentum parameter $\beta$ approaches 1, yielding systematic tracking lag. We complement these upper bounds with minimax lower bounds under gradient-variation constraints, proving this momentum-induced tracking penalty is not an analytical artifact but an information-theoretic barrier: in drift-dominated regimes, momentum is unavoidably worse because stale-gradient averaging forces systematic lag. Our results provide theoretical grounding for the empirical instability of momentum in dynamic settings and precisely delineate regime boundaries where vanilla SGD provably outperforms its accelerated counterparts.

   @inproceedings{sahu2026provable,
     title     = {On the Provable Suboptimality of Momentum {SGD} in Nonstationary Stochastic Optimization},
     author    = {Sahu, Sharan and Hogan, Cameron J. and Wells, Martin T.},
     booktitle = {Proceedings of the 43rd International Conference on Machine Learning},
     year      = {2026}
   }

2. NeurIPS

   Provably Reliable Classifier Guidance through Cross-entropy Error Control

   Sharan Sahu, Arisina Banerjee, and Yuchen Wu

   In The 40th Annual Conference on Neural Information Processing Systems (NeurIPS) 2026

   Abs arXiv Bib PDF

   Classifier-guided diffusion models generate conditional samples by augmenting the reverse-time score with the gradient of the log-probability predicted by a probabilistic classifier. In practice, this classifier is usually obtained by minimizing an empirical loss function. While existing statistical theory guarantees good generalization performance when the sample size is sufficiently large, it remains unclear whether such training yields an effective guidance mechanism. We study this question in the context of cross-entropy loss, which is widely used for classifier training. Under mild smoothness assumptions on the classifier, we show that controlling the cross-entropy at each diffusion model step is sufficient to control the corresponding guidance error. In particular, probabilistic classifiers achieving conditional KL divergence $\varepsilon^2$ induce guidance vectors with mean squared error $\widetilde O(d \varepsilon)$, up to constant and logarithmic factors. Our result yields an upper bound on the sampling error of classifier-guided diffusion models and bears resemblance to a reverse log-Sobolev–type inequality. To the best of our knowledge, this is the first result that quantitatively links classifier training to guidance alignment in diffusion models, providing both a theoretical explanation for the empirical success of classifier guidance, and principled guidelines for selecting classifiers that induce effective guidance.

   @inproceedings{sahu2026provably,
     title     = {Provably Reliable Classifier Guidance through Cross-entropy Error Control},
     author    = {Sahu, Sharan and Banerjee, Arisina and Wu, Yuchen},
     booktitle = {Advances in Neural Information Processing Systems},
     year      = {2026}
   }

3. arXiv

   Adapt or Forget: Provable Tradeoffs Between Adam and SGD in Nonstationary Optimization

   Sharan Sahu, Abir Sarkar, Cameron J. Hogan, and Martin T. Wells

   arXiv preprint arXiv:2605.04269 2026

   Abs arXiv Bib PDF

   We provide a theoretical analysis of Adam under non-stationary stochastic objectives, separating two regimes: Euclidean tracking under adaptive strong monotonicity of the Adam-preconditioned mean-gradient operator, and high-probability projected stationarity guarantees under general $L$-smooth objectives. In the tracking regime, we derive finite-time expected and high-probability bounds that decompose sharply into four components: initialization, objective drift, a first-moment tracking error governed by $\beta_1$, and a preconditioner perturbation governed by $\beta_2$. We characterize the burn-in time required for the transient terms to decay to the asymptotic tracking bound under constant and step-decay schedules. We also prove a high-probability bound on the average projected stationarity gap for Adam under distribution shift. Across both analyses, our bounds reveal a noise–drift tradeoff: in noise-dominated regimes, first-moment averaging and adaptive preconditioning can yield favorable upper guarantees, whereas in drift-dominated regimes, stale first-moment information and preconditioner perturbations can enlarge Adam's tracking guarantee, potentially allowing vanilla SGD to attain a smaller tracking error. Our explicit $(\beta_1,\beta_2,\epsilon)$-dependent bounds identify mechanisms through which adaptive step-sizing can help or hurt under nonstationarity and provide theoretical explanations consistent with Adam's empirical instability and stabilization under distribution shift.

   @article{sahu2026adapt,
     title   = {Adapt or Forget: Provable Tradeoffs Between {Adam} and {SGD} in Nonstationary Optimization},
     author  = {Sahu, Sharan and Sarkar, Abir and Hogan, Cameron J. and Wells, Martin T.},
     journal = {arXiv preprint arXiv:2605.04269},
     year    = {2026}
   }

4. arXiv

   HOPPER: Learnable Hop Extraction for Linearized Graph Sequence Models

   Isuru Herath, Arin Gopakumar, and Sharan Sahu

   arXiv preprint arXiv:2608.09031 2026

   Abs arXiv Bib PDF

   Graph neural networks typically propagate information through repeated message-passing layers, coupling propagation distance with the number of nonlinear transformations applied. This coupling can make deep architectures difficult to optimize and lead to over-smoothing, over-squashing, and loss of long-range information. Linearized Graph Sequence Models (LGSMs) address this issue by separating propagation depth from processing depth and representing successive propagation states of each node as a sequence. However, existing LGSMs construct these sequences using fixed graph operators, limiting their ability to adapt propagation to the input graph, node features, and downstream task. We introduce HOPPER, an end-to-end learnable extension of LGSM that learns how hop sequences should be extracted before processing by a modern state-space model. HOPPER supports feature-conditioned, structure-aware, graph and hop-adaptive propagation while preserving permutation equivariance, with standard adjacency-based and non-backtracking LGSM sequences arising as special cases of the extractor family. HOPPER is state-of-the-art or competitive across ECHO-Synth and performs strongly on City-Networks. On the LRIM physics-based long-range dependency benchmark, varying the maximum neighborhood size used for message-backtracking cancellation, corresponding to the structural memory window, substantially affects performance. Ablations further isolate the contributions of the learnable extraction mechanism and its structural and feature-adaptive components, showing that adaptive hop-sequence construction provides gains beyond the downstream sequence model alone. Together, these results demonstrate that learnable sequence extraction is a flexible and effective framework for long-range graph representation learning across synthetic, physics-based and real-world graph benchmarks.

   @article{herath2026hopper,
     title   = {{HOPPER}: Learnable Hop Extraction for Linearized Graph Sequence Models},
     author  = {Herath, Isuru and Gopakumar, Arin and Sahu, Sharan},
     journal = {arXiv preprint arXiv:2608.09031},
     year    = {2026}
   }

5. arXiv

   Online Distributionally Robust LLM Alignment via Regression to Relative Reward

   Sharan Sahu and Martin T. Wells

   arXiv preprint arXiv:2509.19104 2025

   Abs arXiv Bib PDF

   Reinforcement Learning with Human Feedback (RLHF) has become crucial for aligning Large Language Models (LLMs) with human intent. However, existing offline RLHF approaches suffer from overoptimization, where language models degrade by overfitting inaccuracies and drifting from preferred behaviors observed during training. Distributionally robust optimization (DRO) is a natural solution, but existing DRO-DPO methods are sample-inefficient, ignore heterogeneous preferences, and lean on brittle heuristics. We introduce DRO-REBEL, a family of robust online REBEL updates built on type-$p$ Wasserstein, Kullback-Leibler (KL), and $\chi^2$ ambiguity sets. Strong duality reduces each update to a relative-reward regression, retaining REBEL's scalability without PPO-style clipping or value networks. Under linear rewards, log-linear policies, and a standard coverage condition, we prove $\widetilde{O}(\sqrt{d/n})$ bounds on squared parameter error, with sharper constants than prior DRO-DPO analyses, and give the first parametric $\widetilde{O}(d/n)$ rate for DRO-based alignment under preference shift, matching non-robust RLHF in benign regimes. Each divergence yields a tractable SGD-based algorithm: gradient regularization for Wasserstein, importance weighting for KL, and a 1-D dual solve for $\chi^2$. On Emotion Alignment, the ArmoRM multi-objective benchmark, and HH-Alignment, DRO-REBEL outperforms prior robust and non-robust baselines across unseen preference mixtures, model sizes, and dataset scales.

   @article{sahu2025online,
     title   = {Online Distributionally Robust {LLM} Alignment via Regression to Relative Reward},
     author  = {Sahu, Sharan and Wells, Martin T.},
     journal = {arXiv preprint arXiv:2509.19104},
     year    = {2025}
   }