arXiv AI

Depth and Scale in the Sub-150M Regime: JugnuLM-53M vs JugnuLM-110M

arXiv Machine Learning
Sep 3

Train What You Deploy: Closing the MLP Reachability Gap in Low-Rank Clone Distillation

The paper introduces a new approach to low‑rank clone distillation that ensures the student model’s training targets exactly the weights it will use at inference. By redefining the training objective to cover the full deployed matrix—without changing the model’s shape, parameter count, or FLOPs—the authors recover previously unreachable linear degrees of freedom. This results in significant performance gains across multiple teacher models, achieving comparable or superior accuracy with fewer tokens and parameters.

By Wenhui Chen, Zhifeng Li, Jie Zhou, Navan Preet Singh, Madalina Ciobanu, Chenghua Wang, Qingqing Mao, Ritankar Das
arXiv Computation and Language
3d ago

Tune the Baseline First: Most of Muon's Scaled Weight Decay Gain Is a Tuning Gap, but Its Small Real Advantage Grows With Model Size

The study re‑examines scaled weight decay (Muon‑SW) versus constant decoupled weight decay for language models, tuning learning‑rate and weight‑decay jointly across a range of model sizes. It finds that most of the reported performance gain is due to a tuning gap, with the tuned advantage shrinking to only 0.0026 BPB. However, the remaining advantage grows with model size, rising from 0.0036 to 0.0051 BPB over an eight‑fold increase in active parameters, and the optimal LR‑WD product shifts faster for Muon‑SW.

By Esther Xin
arXiv AI
Sep 25

Wiring Beats Blending: Structure-Aware Compensation for Transformer Downscaling

The paper investigates converting a large pretrained transformer (1.4 B parameters) into a smaller sibling (410 M) by studying representation alignment and parameter projection. It finds that dense weight projection destroys structure, and that a low‑budget, structure‑aware compensation—separating least‑squares function alignment from variance‑preserving rescaling—yields significant gains on token‑efficient training, outperforming subcloning and standard distillation pipelines at matched budgets.

By Ravi Satya Durga Prasad Yenugula
arXiv Machine Learning
1d ago

FastAFD: When Does Attention-FFN Disaggregation Pay on Rack-Scale NVLink?

FastAFD is an attention‑FFN disaggregation engine designed for rack‑scale NVLink systems such as GB200 NVL72. It overlaps the mixture‑of‑experts computation of one microbatch with the attention of another, achieving up to 1.17× higher throughput than the best wide‑expert‑parallelism vLLM configuration on Qwen3 and Qwen3.5, and matching it on MiniMax‑M2.5. An analytical model predicts AFD performance within 3% using only colocated vLLM profiles, revealing that speedup comes from batch and overlap gains that outweigh disaggregation costs, and that speculative decoding further boosts performance for light‑expert models.

By Yichao Fu, Yuxuan Zhang, Minshen Zhang, Ruitian Wang, Junda Chen, Hao Zhang
arXiv Computation and Language
Sep 25

Baseline Shape Decides the Verdict: A Controlled Re-Examination of Ternary Language Models at 60K Parameters

The study re‑examines a reported advantage of a routed ternary (1.58‑bit) language model over a full‑precision transformer at 60K parameters. By running controlled experiments with multiple seeds and a fixed training recipe, the authors find that the apparent benefit largely stems from the choice of baseline model shape rather than the ternary architecture itself. While the routed model does outperform other shapes at a larger 130M‑byte budget, its advantage diminishes when a plain gated diagonal‑SSM block is used, and the ternary penalty varies with architecture and quantization details.

By Gautam Veldanda
arXiv Machine Learning
Sep 7

Optimizer Memory Schedules for Outscaling the Overtraining Axis

The paper studies how different optimizers perform as training duration (overtraining) increases, focusing on matrix‑preconditioned methods (Muon, SOAP) and a momentum‑scheduled method (ADANA) compared to AdamW. Across models ranging from 51M to 253M parameters and overtraining factors up to 256×, the authors find that optimal learning‑rate schedules, weight‑decay coefficients, and memory settings shift with horizon, and that ADANA consistently outperforms AdamW, especially with log‑time weight decay and momentum cooldown. Muon and SOAP maintain roughly constant token‑efficiency advantages, with SOAP potentially improving at the highest overtraining levels.

By Katie Everett, Shikai Qiu
arXiv AI
Oct 7

Does Muon Need Fine-Grained Spectral Shaping?

The paper investigates how much spectral detail is necessary for the Muon optimizer, which traditionally uses a flat spectral profile. By analyzing singular modes, the authors find that most modes lie below a noise edge yet align positively with the gradient, leading them to propose BulkBoost—a two‑band spectral reweighting framework that adjusts bulk and spike gains while preserving matrix norms. Experiments across various pre‑training settings show that this low‑dimensional reweighting matches or surpasses fine‑grained spectral profiles, achieving notable loss reductions.

By Meher Chaitanya, Tianyi Zhou, Aristides Gionis
arXiv Machine Learning
Oct 28

Error Diffusion Replaces Latent Weights: Training Binary-Weight Language Models on a Few Bits of Dithered State

The paper introduces an error‑diffusion (ED) optimizer that replaces the full‑precision latent copy of binary‑weight language models with a small 2‑8‑bit residual buffer, reducing training state per weight from 97 bits to 3‑9 bits—up to 19× less. ED accumulates unrealized weight updates in the buffer and flips the binary weight when the accumulated value crosses a threshold aligned with the quantization grid, enabling efficient training on large transformers with minimal loss compared to latent‑weight baselines. The method scales to models with up to 19 M parameters on enwik8 and 70 M on wikitext‑103, maintaining performance within a few percent of reference while allowing on‑device adaptation within a 275 MB memory budget.

By Pierre-Olivier Bonin
arXiv Machine Learning
Jun 3

MuLoCo: Muon is a practical inner optimizer for DiLoCo

arXiv:2505. 23725v3 Announce Type: replace Abstract: DiLoCo is a powerful framework for training large language models (LLMs), enabling larger optimal batch sizes and increased accelerator utilization under networking constraints.

By Benjamin Th\'erien, Xiaolong Huang, Aaron Defazio, Irina Rish, Eugene Belilovsky
arXiv Machine Learning
Oct 28

ScaleRLT: Exact Log-Probabilities and Fast Deferred Rollouts for Reinforcement Learning on Looped Transformers

ScaleRLT introduces a hybrid learner log‑probability method that reconciles differences between dense and deferred rollout computations in looped language models like Ouro. By using a stop‑gradient replay of the rollout schedule, the approach achieves bit‑for‑bit reproducibility across FP32 and bf16 real‑weight cells while maintaining the same gradient accuracy as dense implementations. The technique boosts rollout throughput by up to 1.56× on a single GPU and significantly speeds up multi‑node GRPO updates, leading to substantial accuracy gains on GSM8K in short training runs.

By Dezhen Lu, Zixiang Ni, Runguo Li