Ashley Zhang

Next-gen LLM Architecture and Optimizer. Research Scientist @tilderesearch

[Ashley Zhang @ashley11829]

The first post on this account! It's my favorite work this year

Introducing our most recent work, 𝗢𝗻𝗹𝗶𝗻𝗲 𝗞𝗟 𝗦𝗵𝗮𝗺𝗽𝗼𝗼 (𝗢𝗞𝗟𝗦)! 🚀 We achieve 𝟭.𝟰𝟱× 𝘁𝗵𝗲 𝗽𝗮𝗿𝗮𝗺𝗲𝘁𝗲𝗿 𝗲𝗳𝗳𝗶𝗰𝗶𝗲𝗻𝗰𝘆 𝗼𝗳 𝗠𝘂𝗼𝗻 while retaining 98% of its training throughput.

Most Deep Learning optimizers fall into two categories: 𝗔𝗱𝗮𝗚𝗿𝗮𝗱 and 𝗦𝘁𝗲𝗲𝗽𝗲𝘀𝘁 𝗗𝗲𝘀𝗰𝗲𝗻𝘁. While the current SOTA, 𝗠𝘂𝗼𝗻, can be viewed as a spectral norm steepest descent, we developed OKLS as a new member of the AdaGrad family.

Inspired by 𝘜𝘯𝘥𝘦𝘳𝘴𝘵𝘢𝘯𝘥𝘪𝘯𝘨 𝘢𝘯𝘥 𝘐𝘮𝘱𝘳𝘰𝘷𝘪𝘯𝘨 𝘚𝘩𝘢𝘮𝘱𝘰𝘰 𝘢𝘯𝘥 𝘚𝘖𝘈𝘗 𝘷𝘪𝘢 𝘒𝘶𝘭𝘭𝘣𝘢𝘤𝘬–𝘓𝘦𝘪𝘣𝘭𝘦𝘳 𝘔𝘪𝘯𝘪𝘮𝘪𝘻𝘢𝘵𝘪𝘰𝘯 (𝘢𝘳𝘟𝘪𝘷:2509.03378), OKLS is a zero-stale, optimal KL Kronecker product approximated AdaGrad.

While designs like this are typically considered too expensive and impractical—often forcing a trade-off with staleness—we conquered it with a hardware-aligned numerical algorithm: 𝗦𝗰𝗮𝗹𝗲𝗱 𝗖𝗔𝗡𝗦 𝗖𝗼𝘂𝗽𝗹𝗲𝗱 𝗡𝗲𝘄𝘁𝗼𝗻-𝗦𝗰𝗵𝘂𝗹𝘇, maintaining 98% of Muon's training throughput. ⚡

We also derived the 𝗺𝘂𝗣 𝘀𝗰𝗮𝗹𝗶𝗻𝗴 𝗿𝘂𝗹𝗲 for OKLS from first principles, which has been validated to successfully transfer the optimal learning rate across different model sizes.

This demonstrates that by combining theoretically sound designs with meticulous numerical algorithms, there is significant potential for discovery. 🧠

Blog link: https://t.co/MEV8tJxrCo

[Tilde @tilderesearch]

Introducing Online KL Shampoo (OKLS), an optimizer that brings a KL-optimal approximation of full-matrix AdaGrad to language-model training.

Diagonal optimizers ignore correlations between gradient coordinates. Full-matrix AdaGrad captures this geometry but requires quadratic state. Muon considers correlations but not their history. OKLS closes this gap using KL-optimal Kronecker factors, whitening matrix gradients across both row and column directions while remaining naturally scale-invariant.

The main challenge is computing fresh inverse-square-root preconditioners at every step. Even one-step staleness can destabilize training. We make zero-staleness preconditioning practical with Scaled CANS Coupled Newton–Schulz: 10 iterations, 27 FP16 GEMMs, and FP32 accumulation.

OKLS achieves 1.45× the parameter efficiency of Muon while retaining 98% of its training throughput. Across 200M–1B models, an OKLS model matches a Muon model roughly 1.5× larger.

[Ashley Zhang @ashley11829]

And we found the EMA for the Online KL Shampoo preconditioners is relatively fast (0.9482). I'm curious if the first-order congruence update used here can handle tracking it

[Ashley Zhang @ashley11829]

Looks great! I also tried to maintain S^{-1/2} and update it with the Riemannian SPD cone metric, in a similar but not exactly the same style

In my early tests, it risked diverging during training due to a scale mismatch between the Riemannian SPD preconditioner (lives in log scale) and the gradient (lives in linear scale), which is why we ultimately turned to the (ultra fast) exact-root approach for Online KL Shampoo, which outperforms Muon by 1.59x parameter efficiency https://t.co/MEV8tJxrCo

I'm looking forward to seeing if KL-Root-Kron's formulation and update rules can stabilize the "maintaining -1/2 root" approach in practice and comparisons against Online KL Shampoo

[You Jiacheng @YouJiacheng]

Combining PSGD-Kron's Σ^{-1/2} (Σ=E[gg.T]) multiplicative update and KL-Shampoo's KL objective, we can get an optimizer WITHOUT eig/inverse/solve.
Wow. (work done by GPT Pro)
The method: multiplicative update P with the objective KL(N(0, Σ), N(0, P^{-2})).

[Ashley Zhang @ashley11829]

For Untra-FineWeb, the compression ratio improves marginally after 16k vocab size, so I do believe 16k- 64k is the sweet spot

[Artur Chakhvadze @norpadon]

Apparently latest claudes have 16K vocabulary, which is very unconventionally low. Qwen3.5 is around 250K for example

My theory is that they are trying to remove the softmax bottleneck

[Ashley Zhang @ashley11829]

The absence of excellent kernels for many linear attentions is one of the major reasons that they are not widely used today; we hope the kernel collection project Popcorn can change the situation

[Tilde @tilderesearch]

We ran 1,373,020 benchmarks across 222 open ML kernel implementations. The result is Popcorn: the world’s largest collection of verified ML kernels.

Today, we’re releasing it to everyone.

[Kaiyue Wen @wen_kaiyue]

A hidden detail in the recently released Muse Glimmer model: its per-matrix weight RMS norm is pinned almost exactly at ~6e-3!

If you’re curious why fixing weight RMS like this can work, feel free to checkout Hyperball: https://t.co/9AFvaJMmpx

[Jensen Huang @JensenHuang]

For my first post, I’m sharing a letter @NVIDIA signed on why open models matter.

AI will transform every industry, power every company, and be built by every country.

Open models strengthen safety and cybersecurity, accelerate innovation and diffusion, and enable sovereignty.

The world needs both frontier closed models and frontier open models.

https://t.co/AUKzoQ5Ikb

[Ashley Zhang @ashley11829]

We fixed a potential numerical issue in the Online KL Shampoo (OKLS) code release

While the Scaled CANS Coupled Newton-Schulz is incredibly stable even for low-rank matrices, the eigenvalue estimation, containing only O(n^2) operations, can overflow or underflow in FP32 for extreme inputs

Optimally utilizing floating numbers is something so high-yield

[Ashley Zhang @ashley11829]

The first post on this account! It's my favorite work this year

Introducing our most recent work, 𝗢𝗻𝗹𝗶𝗻𝗲 𝗞𝗟 𝗦𝗵𝗮𝗺𝗽𝗼𝗼 (𝗢𝗞𝗟𝗦)! 🚀 We achieve 𝟭.𝟰𝟱× 𝘁𝗵𝗲 𝗽𝗮𝗿𝗮𝗺𝗲𝘁𝗲𝗿 𝗲𝗳𝗳𝗶𝗰𝗶𝗲𝗻𝗰𝘆 𝗼𝗳 𝗠𝘂𝗼𝗻 while retaining 98% of its training throughput.

Most Deep Learning optimizers fall into two categories: 𝗔𝗱𝗮𝗚𝗿𝗮𝗱 and 𝗦𝘁𝗲𝗲𝗽𝗲𝘀𝘁 𝗗𝗲𝘀𝗰𝗲𝗻𝘁. While the current SOTA, 𝗠𝘂𝗼𝗻, can be viewed as a spectral norm steepest descent, we developed OKLS as a new member of the AdaGrad family.

Inspired by 𝘜𝘯𝘥𝘦𝘳𝘴𝘵𝘢𝘯𝘥𝘪𝘯𝘨 𝘢𝘯𝘥 𝘐𝘮𝘱𝘳𝘰𝘷𝘪𝘯𝘨 𝘚𝘩𝘢𝘮𝘱𝘰𝘰 𝘢𝘯𝘥 𝘚𝘖𝘈𝘗 𝘷𝘪𝘢 𝘒𝘶𝘭𝘭𝘣𝘢𝘤𝘬–𝘓𝘦𝘪𝘣𝘭𝘦𝘳 𝘔𝘪𝘯𝘪𝘮𝘪𝘻𝘢𝘵𝘪𝘰𝘯 (𝘢𝘳𝘟𝘪𝘷:2509.03378), OKLS is a zero-stale, optimal KL Kronecker product approximated AdaGrad.

While designs like this are typically considered too expensive and impractical—often forcing a trade-off with staleness—we conquered it with a hardware-aligned numerical algorithm: 𝗦𝗰𝗮𝗹𝗲𝗱 𝗖𝗔𝗡𝗦 𝗖𝗼𝘂𝗽𝗹𝗲𝗱 𝗡𝗲𝘄𝘁𝗼𝗻-𝗦𝗰𝗵𝘂𝗹𝘇, maintaining 98% of Muon's training throughput. ⚡

We also derived the 𝗺𝘂𝗣 𝘀𝗰𝗮𝗹𝗶𝗻𝗴 𝗿𝘂𝗹𝗲 for OKLS from first principles, which has been validated to successfully transfer the optimal learning rate across different model sizes.

This demonstrates that by combining theoretically sound designs with meticulous numerical algorithms, there is significant potential for discovery. 🧠

Blog link: https://t.co/MEV8tJxrCo

[Ashley Zhang @ashley11829]

If we find some way to test-time scale looped models - not only scale toward the "more computation and more powerful", but also "very little computation for easy problems", looped models will be the future

[Tilde @tilderesearch]

The beta period for OneLayerDeeper is over! We are very excited to collaborate on this competition with the great folks @coreauto

Rules: https://t.co/Xa7BYqOJTY
Blog post: https://t.co/hIgrrzvtsv

Competition lasts until the end of August!

[Bruno Mlodozeniec @brunorganised]

Tilde seems like such an underrated neolab (purely judging by seeming quality of work vs valuation order of magnitude bracket)

[Ashley Zhang @ashley11829]

@17Ahmetyucel @tilderesearch It's L=c+AN^{alpha}