← Search

Yushun Zhang

7 accepted papers

2026

MoFO: Momentum-Filtered Optimizer for Mitigating Forgetting in LLM Fine-Tuning

ICML 2026poster

Large language models (LLMs) have demonstrated remarkable capabilities across a wide range of tasks. Typically, LLMs are first pre-trained on large corpora and subsequently fine-tuned on task-specific datasets. However, during fine-tuning, LLMs may forget some knowledge acquired in the pre-training …

Cited by 0SourceScholar
2025

Adam-mini: Use Fewer Learning Rates To Gain More

ICLR 2025poster

We propose Adam-mini, an optimizer that achieves on-par or better performance than AdamW with $50$% less memory footprint. Adam-mini reduces memory by cutting down the learning rate resources in Adam (i.e., $1/\sqrt{v}$). By delving into the Hessian structure of neural nets, we find Adam’s $v$ might…

2024

ReMax: A Simple, Effective, and Efficient Reinforcement Learning Method for Aligning Large Language Models

ICML 2024poster

Reinforcement Learning from Human Feedback (RLHF) is key to aligning Large Language Models (LLMs), typically paired with the Proximal Policy Optimization (PPO) algorithm. While PPO is a powerful method designed for general reinforcement learning tasks, it is overly sophisticated for LLMs, leading to…

2024

Why Transformers Need Adam: A Hessian Perspective

NeurIPS 2024poster

SGD performs worse than Adam by a significant margin on Transformers, but the reason remains unclear. In this work, we provide an explanation through the lens of Hessian: (i) Transformers are "heterogeneous'': the Hessian spectrum across parameter blocks vary dramatically, a phenomenon we call "bloc…

2022

Adam Can Converge Without Any Modification On Update Rules

NeurIPS 2022accept

Ever since \citet{reddi2019convergence} pointed out the divergence issue of Adam, many new variants have been designed to obtain convergence. However, vanilla Adam remains exceptionally popular and it works well in practice. Why is there a gap between theory and practice? We point out there is a mis…

Cited by 98SourcePDFScholar
2022

HyperDQN: A Randomized Exploration Method for Deep Reinforcement Learning

ICLR 2022poster

Randomized least-square value iteration (RLSVI) is a provably efficient exploration method. However, it is limited to the case where (1) a good feature is known in advance and (2) this feature is fixed during the training. If otherwise, RLSVI suffers an unbearable computational burden to obtain the…

2021

When Expressivity Meets Trainability: Fewer than $n$ Neurons Can Work

NeurIPS 2021poster

Modern neural networks are often quite wide, causing large memory and computation costs. It is thus of great interest to train a narrower network. However, training narrow neural nets remains a challenging task. We ask two theoretical questions: Can narrow networks have as strong expressivity as wid…

Cited by 14SourcePDFScholar