Scholarlink

<p>Training language models (LMs) under tight GPU memory budgets rules out standard back-propagation and motivates <i>zeroth-order</i> (ZO) optimization. While ZO methods have proven effective for fine-tuning, their potential during the more memory-intensive pretraining stage has received little attention. We revisit the singular-value spectra of layer gradients during pretraining and show that the gradient information is spread across many directions; low-rank ZO methods therefore potentially discard some informative components. Building on this insight, we introduce <b>KronZO</b>, a Kronecker-structured ZO optimizer that (i)&#xa0;explores a full-rank search subspace with state-of-the-art storage compression and (ii)&#xa0;employs a criterion-driven <i>directional</i> update that selectively keeps only informative steps. When pretraining GPT-2 Small on OpenWebText, KronZO achieves a markedly lower training loss than previous ZO baselines while consuming less GPU memory. KronZO narrows the gap with first-order (FO) baselines at a fraction of their memory footprint, extending ZO optimization to larger models and paving the way for memory-efficient pretraining on commodity hardware. Beyond pretraining, KronZO-pretrained models retain strong downstream transferability, matching FO baselines on the GLUE benchmark.</p>

错误：搜索内容不能为空，请输入英文关键词
错误：关键词超出字数限制，请精简
高级检索

Zeroth-Order Kronecker Optimization for Pretraining Language Models

  • 期刊论文

  • 发表日期：

    2026年2月2日

查看全文
  • Nathan Allaire,
  • Sébastien Le Digabel,
  • Dominique Orban,
  • Vahid Partovi Nia

摘要

Training language models (LMs) under tight GPU memory budgets rules out standard back-propagation and motivates zeroth-order (ZO) optimization. While ZO methods have proven effective for fine-tuning, their potential during the more memory-intensive pretraining stage has received little attention. We revisit the singular-value spectra of layer gradients during pretraining and show that the gradient information is spread across many directions; low-rank ZO methods therefore potentially discard some informative components. Building on this insight, we introduce KronZO, a Kronecker-structured ZO optimizer that (i) explores a full-rank search subspace with state-of-the-art storage compression and (ii) employs a criterion-driven directional update that selectively keeps only informative steps. When pretraining GPT-2 Small on OpenWebText, KronZO achieves a markedly lower training loss than previous ZO baselines while consuming less GPU memory. KronZO narrows the gap with first-order (FO) baselines at a fraction of their memory footprint, extending ZO optimization to larger models and paving the way for memory-efficient pretraining on commodity hardware. Beyond pretraining, KronZO-pretrained models retain strong downstream transferability, matching FO baselines on the GLUE benchmark.