Paper deep dive
Training nGPT
Ilya Loshchilov, Boris Ginsburg
Intelligence
Status: not_run | Model: - | Prompt: - | Confidence: 0%
Entities (0)
Relation Signals (0)
No relation signals yet.
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
Abstract:The normalized Transformer (nGPT) realizes hyperspherical representation learning by constraining model parameter vectors and activation vectors to the unit hypersphere. In this paper, we describe a practical training recipe for nGPT and evaluate it on modern hybrid Mamba-2--Transformer Mixture-of-Experts (MoE) models. The recipe introduces Logit Gradient Preconditioning, Logarithmic Learning Rate Decay, GatedAdamW, angular update control, and optional exploration mechanisms. Compared with an unnormalized model of the same hybrid MoE architecture trained with AdamW, the 14B-total-parameter nGPT model reaches the same validation loss using approximately half as many training tokens. The recipe scales across the models considered, which contain up to 14B total parameters.
Tags
Links
- Source: https://arxiv.org/abs/2608.01284v1
- Canonical: https://arxiv.org/abs/2608.01284v1
Trouble viewing inline? Open PDF directly →
Full Text
57,134 characters extracted from source content.
Expand or collapse full text
Training nGPT Ilya Loshchilov & Boris Ginsburg NVIDIA iloshchilov,bginsburg@nvidia.com Abstract The normalized Transformer (nGPT) realizes hyperspherical representation learning by constraining model parameter vectors and activation vectors to the unit hypersphere. In this paper, we describe a practical training recipe for nGPT and evaluate it on modern hybrid Mamba-2–Transformer Mixture-of-Experts (MoE) models. The recipe introduces Logit Gradient Preconditioning, Logarithmic Learning Rate Decay, GatedAdamW, angular update control, and optional exploration mechanisms. Compared with an unnormalized model of the same hybrid MoE architecture trained with AdamW, the 14B-total-parameter nGPT model reaches the same validation loss using approximately half as many training tokens. The recipe scales across the models considered, which contain up to 14B total parameters. Figure 1: nGPT’s forward pass as a multi-step optimization on the hypersphere. 1 Preliminaries The normalized Transformer (nGPT) paper (Loshchilov et al., 2024) proposed a hyperspherical representation in which all activation vectors and parameter vectors that form matrices are normalized to unit norm. The high-level idea behind nGPT is illustrated in Figure 1, where the initial sequence “Life is” is used to predict “beautiful” via a two-step process within a single layer. The first step of the attention block considers the two blue tokens “Life” and “is”, represented as points, and suggests a prediction depicted by the green point. This prediction is then combined with the hidden state to produce a new intermediate point. The resulting intermediate point is passed to the MLP block, which in turn produces its suggestion depicted by the red point. The hidden state then takes a second learned step towards the MLP suggestion. When this process is repeated across layers, next-token prediction can be viewed as an optimization process on the hypersphere: model vectors act as anchors, dot products measure similarity to these anchors, and the attention and MLP blocks propose successive steps. nGPT was designed based on a first-principles view of the hypersphere as the representation manifold, following an earlier attempt to control the norms of Transformer parameters during training (Loshchilov, 2023). At the same time, it connects to a broad line of work that either explicitly studies hyperspherical representations (Liu et al., 2017; Wang et al., 2017; Liu et al., 2018; Xu and Durrett, 2018; Mettes et al., 2019; Wang and Isola, 2020; Liu et al., 2021; Karras et al., 2024) or arrives at related ideas through normalization, decoupled weight decay, norm control, and rotational training dynamics (Salimans and Kingma, 2016; Loshchilov and Hutter, 2019; Franke et al., 2023; Kodryan et al., 2022; Kosson et al., 2023). In this work, we extend nGPT from the dense Transformer setting originally studied in Loshchilov et al. (2024) to modern hybrid Mamba-2–Transformer MoE language models, using the Nemotron-3 architecture and an industry-grade data blend (Blakeman et al., 2025). We refer to the complete collection of normalization, architectural, and optimization modifications considered in this work as the nGPT training recipe. Its principal components include Logit Gradient Preconditioning, logarithmic learning rate decay, GatedAdamW, Pre-Moment Tangent Projection, Angular Step Cap, Second-Moment Growth Clipping, and Post-Moment Exploration Noise. These names refer to individual components of the recipe rather than to separate model architectures. The core recipe is not specific to MoE models. For a dense feed-forward block, the same normalization and optimization rules can be applied to the ordinary MLP projection matrices, while the router- and expert-specific modifications are omitted. In the experiments, we use GPT and nGPT to denote, respectively, the unnormalized and normalized variants of the same hybrid Mamba-2–Transformer MoE architecture; optimizer names are appended only when the distinction is relevant. The remainder of the paper first introduces the individual training components and then describes their application to nGPT in the hybrid MoE setting. Scope. The experiments in this paper were run with a limited compute budget. As a result, we focus on the training recipe that worked in practice and do not attempt an exhaustive ablation study. In particular, comprehensive component-wise ablations and hyperparameter scaling laws are left for future work. 2 Optimization Ingredients 2.1 Logit Gradient Preconditioning (LGP) In nGPT, the output embedding vectors and hidden states are normalized, so the unscaled logits are based on bounded dot products. As in the original nGPT formulation, we therefore use a learnable vocabulary-wise scale vector z∈ℝV s_z ^V to control the sharpness of the output distribution. We initialize this vector as z,0=. s_z,0=1. (1) Let ∈ℝV u ^V denote the logits before this scale is applied. The forward pass uses =z⊙, z= s_z u, (2) and the cross-entropy loss is computed from z. The motivation for LGP came from an empirical observation in a dense 8B nGPT model trained on 4T tokens alongside the Nemotron-H pretraining cycle (Blakeman et al., 2025). Figure 2 shows that after an initial transient, the mean value of the learned vocabulary logit scale vector z s_z grows logarithmically. While different token-specific entries of z s_z grow at different rates, the mean of z s_z appears as an explicit global multiplier on the gradient propagated through the output layer to all preceding layers, independently of the explicit learning rate schedule. This hidden, time-varying multiplier can complicate the design of optimizers and architectures, and may be one contributor to the commonly observed growth of gradient norms during long pretraining runs. Figure 2: Mean value of the logit scale vector z s_z for an 8B dense nGPT model trained on 4T tokens. The fit is computed from iteration t≥1000t≥ 1000. Standard backpropagation would multiply the gradient entering the output layer by z s_z: ∂ℒ∂=z⊙∂ℒ∂. ∂ u= s_z ∂ z. (3) This couples the learned logit scales to the effective learning rate of the output layer and all preceding layers. In particular, if the mean of z s_z grows during training, then the average gradient scale entering the network also grows, even though this change is not part of the explicit learning rate schedule. To control this coupling, we use z s_z in the forward pass but replace its explicit backward multiplier by a normalized scale: ∂ℒ∂←(zmean(z))q⊙∂ℒ∂, ∂ u← ( s_zmean( s_z) )^q ∂ z, (4) where q controls the strength of the preconditioning. We call this procedure Logit Gradient Preconditioning (LGP), since it preconditions the gradient flowing backward through the logit scale while leaving the forward logits unchanged. In our experiments, we use q=1q=1. The two endpoint settings have simple interpretations. For q=0q=0, the explicit backward scale is (zmean(z))0=, ( s_zmean( s_z) )^0=1, (5) so the direct multiplicative effect of the learned logit scale on the gradient propagated from the output layer into the rest of the network is removed. The scale vector still affects the backward pass indirectly because ∂ℒ/∂ /∂ z depends on the scaled forward logits z. For q=1q=1, the explicit backward scale is zmean(z), s_zmean( s_z), (6) which removes the global mean scale while preserving the relative vocabulary-wise variation in z s_z. Thus, q=0q=0 removes the direct vocabulary-wise multiplier from the output layer Jacobian, whereas q=1q=1 retains its relative vocabulary-wise preconditioning effect while removing the explicit time-varying global multiplier. The gradient with respect to z s_z itself is left unchanged and remains the standard gradient of the forward computation. Therefore, z s_z remains a learnable vocabulary-wise logit-scale parameter, while its explicit backward effect on the rest of the network is controlled independently. Except when the replacement scale equals the forward scale, LGP should be viewed as a preconditioned backward pass rather than the exact gradient of the forward loss with respect to u. For values of q other than zero or one, the mean of (z/mean(z))q( s_z/mean( s_z))^q is not generally equal to one. An alternative is to normalize after taking the element-wise power by defining the backward scale for coordinate i as b~q,i(z)=sz,iq1V∑ℓ=1Vsz,ℓq,i=1,…,V. b_q,i( s_z)= s_z,i^q 1V _ =1^Vs_z, ^q, i=1,…,V. (7) This alternative preserves a unit mean backward scale for every q and agrees with the formulation above at both q=0q=0 and q=1q=1. We leave the comparison of the two formulations and the study of intermediate values of q to future work. The element-wise power assumes positive entries of z s_z. We initialize all entries to one, and the reported experiments use q=1q=1, for which no fractional power is required. Experiments with noninteger values of q should enforce positivity, for example through a positive parameterization of z s_z. LGP is not specific to nGPT. In a standard Transformer, the rows of the output embedding matrix are not normalized. Decomposing each row as i=sz,i¯i w_i=s_z,i w_i, where sz,i=‖i‖2s_z,i=\| w_i\|_2 and ‖¯i‖2=1\| w_i\|_2=1, shows that the row norms act as vocabulary-wise logit scales. They therefore play the same role as the explicit z s_z parameter in nGPT. LGP can in principle be applied to a standard Transformer by computing these row norms, using the equivalent decomposition out=diag(z)¯out W_out=diag( s_z) W_out, and replacing the backward scale associated with z s_z by (z/mean(z))q( s_z/mean( s_z))^q, while leaving the forward logits unchanged. Preliminary experiments suggest that LGP with q=1q=1 removes this explicit, time-varying global factor without observable performance loss while keeping the forward pass unchanged. 2.2 Logarithmic Learning Rate Decay Figure 3: Logarithmic annealing schedules compared with cosine annealing after matching the area under the curve (AUC). All schedules use a 10% linear warmup. Each logarithmic schedule is multiplied by a constant scale factor so that its AUC matches that of the cosine schedule. Smaller ρ allocates more of the learning rate budget early, while large ρ approaches linear decay. An advantage of nGPT over the standard Transformer parameterization is that parameter-vector norms are constrained and weight decay is not used for these vectors. This removes the need to tune the interaction between the learning rate and weight decay and makes the effective step size more directly controlled by the learning rate schedule. We therefore consider a more front-loaded decay schedule, which we call Logarithmic Learning Rate Decay. After an optional warmup, let r=t−twarmuptdecay−twarmup,r∈[0,1], r= t-t_warmupt_decay-t_warmup, r∈[0,1], (8) denote the normalized decay progress. We define a dimensionless schedule multiplier that interpolates between ηmax _ and ηmin _ logarithmically in the normalized training progress: η(t)=ηmin+(ηmax−ηmin)[1−log(1+rρ)log(1+1ρ)], η(t)= _ + ( _ - _ ) [1- (1+ rρ ) (1+ 1ρ ) ], (9) where ρ>0ρ>0 controls the shape of the decay. At optimizer step t, we write ηt=η(t) _t=η(t). The corresponding scheduled adaptive step size is αt=αηt, _t=α\, _t, (10) where α is the base step size in the AdamW notation used below. When ηmax=1 _ =1, α is the peak learning rate. The schedule satisfies η(twarmup)=ηmaxη(t_warmup)= _ and η(tdecay)=ηminη(t_decay)= _ . As ρ→∞ρ→∞, it approaches linear decay. Smaller values of ρ produce a faster initial decrease and a longer tail. Figure 3 compares the proposed logarithmic decay with cosine decay (Loshchilov and Hutter, 2016). All schedules in this example use a linear warmup over the first 10% of training. For this comparison, each logarithmic schedule is multiplied by a constant so that its area under the curve (AUC) matches that of the cosine schedule. This keeps the integrated learning rate budget fixed while changing how that budget is distributed over training. Smaller values of ρ allocate more of the learning rate budget early, producing a sharp decrease after warmup and a long tail. Because of the AUC normalization, the peak schedule multiplier can exceed one. 2.3 GatedAdamW AdamW decouples weight decay from the adaptive gradient update (Loshchilov and Hutter, 2019). We follow the notation of AdamW and denote by t θ_t the parameters, by t g_t the stochastic gradient, by t m_t and t v_t the first and second moment estimates, by α the base step size, by ηt _t the dimensionless schedule multiplier, and by λ the decoupled weight decay coefficient. The scheduled adaptive step size is therefore αt=αηt _t=α _t. In AdamW, the adaptive update is controlled by ^t/(^t+ϵ) m_t/( v_t+ε). The constant ϵε is usually introduced for numerical stability, but it also suppresses updates when ^t v_t is comparable to or smaller than ϵε. We found it useful to separate these two roles. GatedAdamW uses ϵnum _num for numerical stability and a separate soft gate to control the update applied at small second-moment scales. Let t=^t+ϵnum, d_t= v_t+ _num, (11) where ϵnum≥0 _num≥ 0 is a numerical constant. We define the coordinate-wise gate t=σ(alogtϵgate), γ_t=σ (a d_t _gate ), (12) where σ(⋅)σ(·) is the sigmoid function, ϵgate>0 _gate>0 is the gate threshold, and a>0a>0 controls the sharpness of the gate. The GatedAdamW update is then t=t−1−ηt(αt⊙^t+λt−1). θ_t= θ_t-1- _t (α\, γ_t m_t d_t+λ θ_t-1 ). (13) The gate has a simple interpretation. Coordinates with t,i≪ϵgate d_t,i _gate receive a small gate value, which suppresses the numerically normalized direction m^t,i/dt,i m_t,i/d_t,i. Because ϵnum _num may be smaller than the AdamW epsilon, the resulting update can nevertheless be larger than the corresponding AdamW update, particularly when a<1a<1. Coordinates with t,i≫ϵgate d_t,i _gate receive a gate value close to one and experience little gate suppression. Thus, GatedAdamW makes the effect of Adam’s ϵε explicit and tunable: ϵnum _num controls numerical stability, while ϵgate _gate determines the second-moment scale at which the gate transitions from suppressed to active updates. Algorithm 1 GatedAdamW 1:base step size α, dimensionless schedule multiplier ηt _t, moment coefficients β1,β2∈[0,1) _1, _2∈[0,1), numerical constant ϵnum≥0 _num≥ 0, gate threshold ϵgate>0 _gate>0, gate sharpness a>0a>0, weight decay λ 2:initial parameters 0 θ_0 3:0←0 m_0← 0, 0←0 v_0← 0 4:for t=1,2,…,Tt=1,2,…,T do 5: t←∇ft(t−1) g_t←∇ f_t( θ_t-1) 6: t←β1t−1+(1−β1)t m_t← _1 m_t-1+(1- _1) g_t 7: t←β2t−1+(1−β2)t2 v_t← _2 v_t-1+(1- _2) g_t^2 8: ^t←t/(1−β1t) m_t← m_t/(1- _1^t) 9: ^t←t/(1−β2t) v_t← v_t/(1- _2^t) 10: t←^t+ϵnum d_t← v_t+ _num 11: t←σ(alog(t/ϵgate)) γ_t←σ\! (a ( d_t/ _gate ) ) 12: t←t−1−ηt(αt⊙^t+λt−1) θ_t← θ_t-1- _t (α\, Green!65$ γ_t$ Green!65$ γ_t$ Green!65$ γ_t$ Green!65$ γ_t$ m_t d_t+λ θ_t-1 ) 13:end for Figure 4: Effect of the GatedAdamW sharpness parameter a on the coordinate-wise sigmoid gate. The gate is plotted as a function of the bias-corrected second-moment scale v v. The vertical dashed line marks the gate threshold ϵgate=10−8 _gate=10^-8, chosen equal to the AdamW value of ϵε. All curves cross 0.50.5 when v^+ϵnum=ϵgate v+ _num= _gate. Smaller values of a produce a smoother transition, suppressing coordinates below the threshold less strongly while approaching one more slowly above the threshold. A useful special case recovers AdamW exactly. If a=1a=1, ϵnum=0 _num=0, and ϵgate=ϵ _gate=ε, then σ(log^tϵ)=^t^t+ϵ. σ ( v_tε )= v_t v_t+ε. (14) For positive ^t v_t, this gives t⊙^t^t=^t^t+ϵ. γ_t m_t v_t= m_t v_t+ε. (15) The equality at zero second moment is understood through the corresponding continuous extension. The update therefore reduces to AdamW. This makes GatedAdamW a conservative extension of AdamW: this special setting recovers AdamW exactly, while other settings control the transition between suppressed and fully active coordinates. An equivalent formulation of the gate without an explicit sigmoid is given in Appendix A. Figure 4 visualizes the role of the sharpness parameter a in GatedAdamW. In the figure, we use ϵgate=10−8 _gate=10^-8 and ϵnum=10−14 _num=10^-14. The gate is equal to 0.50.5 when dt,i=ϵgated_t,i= _gate, or equivalently when v^t,i=10−8−10−14 v_t,i=10^-8-10^-14, which is visually indistinguishable from 10−810^-8 on the plot. The parameter a controls the sharpness of the transition. When a=1a=1 and ϵnum=0 _num=0, the gate exactly recovers the implicit AdamW epsilon gate. With the small nonzero value of ϵnum _num used in the figure, the difference is negligible on the displayed scale. Smaller values of a make the gate softer: they assign larger gate values to coordinates below the threshold but approach one more slowly above it. Thus, a controls how abruptly coordinates transition from the epsilon-dominated regime to receiving the full adaptive update. Angular Step Cap (ASC). We optionally cap the angular displacement of normalized non-embedding parameter vectors. This provides a safeguard on the geometric step represented by the candidate after normalization. Let t−1 w_t-1 be the previous normalized vector and let ~t w_t be the post-update candidate before final normalization, including any optional post-update perturbation. We extract its component tangent to t−1 w_t-1: t=~t−t−1⟨t−1,~t⟩⟨t−1,t−1⟩. u_t= w_t- w_t-1 w_t-1, w_t w_t-1, w_t-1 . (16) The angle between the previous vector and the candidate is ϕt=atan2(‖t‖2,⟨t−1‖t−1‖2,~t⟩). _t=atan2 (\| u_t\|_2, w_t-1\| w_t-1\|_2, w_t ). (17) If ϕt _t exceeds a cap θmax(t) _ (t) and ‖t‖2>0\| u_t\|_2>0, we replace the candidate by the point at angle θmax(t) _ (t) in the same tangent direction: ~t←cos(θmax(t))t−1‖t−1‖2+sin(θmax(t))t‖t‖2. w_t← \! ( _ (t) ) w_t-1\| w_t-1\|_2+ \! ( _ (t) ) u_t\| u_t\|_2. (18) Otherwise, the candidate is left unchanged. In both cases, the usual nGPT normalization is subsequently applied as the retraction step. We also propose to optionally use the following techniques in GatedAdamW: Pre-Moment Tangent Projection (Appendix B), Post-Moment Exploration Noise (PMEN) (Appendix D), Second-Moment Growth Clipping (SMGC) (Appendix C). They are described in the Appendix because they are not central to the paper and their impact on the validation loss is modest. 3 Training nGPT In this section, we recall the principal design choices of nGPT and describe the modifications used to train it with modern MoE models such as Nemotron-3 (Blakeman et al., 2025; Waleffe et al., 2024). 3.1 Normalization and scaling factors In nGPT, we constrain all parameter vectors that form matrices to have unit norm by normalizing them along the embedding dimension (e.g., dmodel=1024d_model=1024). In distributed settings, normalization must be applied to the parameters on which the optimizer operates. A common implementation bug is to normalize the instantiated model parameters while leaving the optimizer parameters unconstrained. Activations h are also normalized to unit norm when they are recombined with the stream: h ←Norm(+A⊙(A−)), \! ( h+ _A ( h_A- h) ), (19) h ←Norm(+M⊙(M−)), \! ( h+ _M ( h_M- h) ), (20) where A∈ℝ≥0dmodel _A _≥ 0^d_model and M∈ℝ≥0dmodel _M _≥ 0^d_model are learnable parameters, called eigen learning rates, applied to the unit-normalized outputs of the attention/state-space model (SSM) and MLP/MoE blocks, A=Norm(ATTN()) h_A=Norm(ATTN( h)) and M=Norm(MLP()) h_M=Norm(MLP( h)), respectively. The presence of nonlinear elements in the network may render products of normalized vectors too constrained. nGPT introduced a trainable vector qk∈ℝdk s_qk ^d_k to rescale normalized queries q and keys k (see QKNorm (Henry et al., 2020; Nguyen and Salazar, 2019)): q ←Norm()⊙qk, ( q) s_qk, (21) k ←Norm()⊙qk. ( k) s_qk. (22) Similarly, the intermediate activations u of the MLP/MoE block are rescaled by a trainable vector u∈ℝdMLP s_u ^d_MLP: ←⊙u. u← u s_u. (23) This may be particularly important for activations such as GELU or SwiGLU, whose inputs must be at an appropriate scale to preserve nonlinearity. For squared ReLU, the role of this scale differs from that for saturating or gated nonlinearities: positive rescaling does not change the activation support, but it changes the output magnitude quadratically. For any trainable vector of scaling parameters such as a s_a, nGPT uses two scalars, sa,inits_a,init and sa,scales_a,scale. Each entry of a s_a is initialized to sa,scales_a,scale, while the forward pass uses the effective scaling vector aeff=asa,initsa,scale. s_a^eff= s_a s_a,inits_a,scale. (24) This allows us to control the effective learning rate of aeff s_a^eff by adjusting sa,scales_a,scale while keeping the global learning rate unchanged. 3.2 Changes specific to Mamba-2 and MoE We investigated various normalization options for Mamba-2 components (Dao and Gu, 2024) and found that both its input and output projections can be normalized. After removing RMSNorm from the Mamba-2 block, we introduce a trainable scalar smambas_mamba to rescale the input activations and place the inputs to SiLU at an appropriate scale. Similarly, we introduce a trainable scalar smoes_moe for the MoE block to place the inputs to the sigmoid router at an appropriate scale. In both cases, vectors, as in RMSNorm, could be used instead of scalars. However, we did not observe a significant difference between the two choices. The scalar choice preserves an isotropic scaling interpretation within the hyperspherical view. The vector choice is also compatible with this view and can be interpreted as a diagonal preconditioner. 3.3 Summary of modifications The recipe for converting the baseline hybrid Mamba-2–Transformer MoE model into its normalized version is as follows: 1. Remove normalization layers such as RMSNorm and LayerNorm. Remove weight decay. 2. Use GatedAdamW without weight decay (thus, GatedAdam). Apply Pre-Moment Tangent Projection to parameter vectors selected for normalization. Set the gating hyperparameter a to 0.5; with the Adam-equivalent epsilon setting, a=1a=1 recovers Adam. Apply an Angular Step Cap θmax(t) _ (t) to normalized non-embedding vectors, with an optional warmup over the first twarmupt_warmup iterations. After each GatedAdam update, normalize all selected rows or columns of the input and output embedding, attention, MoE, router, and Mamba projection matrices. Parameters not selected for normalization are trained with Adam (GatedAdam with a=1a=1). 3. Change the softmax scaling factor in attention from 1/dk1/ d_k to dk d_k. Normalize and rescale q and k as in Equations 21 and 22, with sqk,init=sqk,scale=1s_qk,init=s_qk,scale=1. 4. Implement the rescaling of the intermediate state of the MoE/MLP block using Equation 23, where u s_u uses su,init=1s_u,init=1 and su,scale=1s_u,scale=1. 5. Implement activation normalization for each recombination as in Equations 19 and 20 with αA,init=αM,init=0.1α_A,init=α_M,init=0.1 (which may be on the order of 1/nlayers1/n_layers) and αA,scale=αM,scale=1α_A,scale=α_M,scale=1. 6. Implement activation scaling for Mamba-2 with smamba,init=0.5dmodels_mamba,init=0.5 d_model and smamba,scale=1s_mamba,scale=1 and for MoE with smoe,init=0.5dmodels_moe,init=0.5 d_model and smoe,scale=1s_moe,scale=1. 7. Implement the rescaling of logits using Equation 2 with sz,init=1s_z,init=1 and sz,scale=1s_z,scale=1. 8. Apply Logit Gradient Preconditioning with q=1q=1. 9. Apply Post-Moment Exploration Noise using the hard-threshold profile at each iteration starting from the first one, τv=10−10 _v=10^-10, cj,t=1/djc_j,t=1/ d_j, and λnoise=10 _noise=10. 10. Apply Second-Moment Growth Clipping after the first 500 iterations, with R=100R=100, smin=10−10s_ =10^-10. 11. In contrast to the original nGPT parameterization, we change sz,scales_z,scale, αA,scale _A,scale, αM,scale _M,scale, and sqk,scales_qk,scale from 1/dmodel1/ d_model to 11. To compensate for this reparameterization, we set their effective peak learning rates to 0.50.5, 0.20.2, 0.20.2, and 0.20.2, respectively. In our experiments, we did not find it necessary to scale these learning rates with model size. 4 Experiments 4.1 Experimental Setup All experiments use models and data developed by the Nemotron Team (Blakeman et al., 2025). More specifically, we use the scaling ladder developed by Khona et al. (2026), which consists of Nemotron-3 Nano models with progressively increasing widths and depths. The model labels refer to total parameter counts. We consider the following models: 1B (0.21B active parameters per token, 88B training tokens), 2B (0.37B active, 140B training tokens), 4B (0.61B active, 215B training tokens), 7B (0.91B active, 310B training tokens), and 14B (1.74B active, 560B training tokens). Thus, the number of training tokens is approximately 320–420 times the number of active parameters, excluding the input embeddings (Khona et al., 2026). Additional details of the ladder models are omitted because they have not yet been fully described by their authors. Figure 5: Scaling results for hybrid Mamba-2–Transformer MoE models trained with AdamW (denoted as GPT AdamW) and their normalized versions trained with GatedAdam (denoted as nGPT GatedAdam). Five model sizes (by total number of parameters) are considered: 1B, 2B, 4B, 7B, and 14B (Khona et al., 2026). The 14B nGPT model trained with GatedAdam incurs about 6% overhead; one could therefore shift the corresponding nGPT results by 6% along the x-axis. We do not make this adjustment because the measured overhead is not necessarily attributable to increased FLOPs. Table 1: Scaling results for hybrid Mamba-2–Transformer MoE models trained with AdamW (denoted as GPT AdamW) and their normalized versions trained with GatedAdam (denoted as nGPT GatedAdam). Five model sizes (by total number of parameters) are considered: 1B, 2B, 4B, 7B, and 14B (Khona et al., 2026). Model size Tokens Training loss Validation loss GPT nGPT Diff. Diff. % GPT nGPT Diff. Diff. % 1B 88B 2.0561 1.9946 0.0615 2.99% 1.6435 1.5881 0.0554 3.37% 2B 140B 1.9053 1.8648 0.0405 2.13% 1.5230 1.4829 0.0401 2.63% 4B 215B 1.7994 1.7573 0.0421 2.34% 1.4330 1.3928 0.0402 2.81% 7B 310B 1.7251 1.6849 0.0402 2.33% 1.3680 1.3310 0.0370 2.70% 14B 560B 1.6207 1.5812 0.0395 2.44% 1.2827 1.2501 0.0326 2.54% For readability, we refer to a Nemotron-3-style hybrid Mamba–Transformer model as GPT and to its normalized counterpart as nGPT. The baseline model and results were provided by Khona et al. (2026), together with the following recommended hyperparameters: β1=0.9 _1=0.9, β2=0.95 _2=0.95, and weight decay of 0.10.1. After a warmup over 1B tokens, the peak learning rate was set to 2.2×10−32.2× 10^-3 for the 1B model, 2.0×10−32.0× 10^-3 for the 2B model, 1.8×10−31.8× 10^-3 for the 4B model, 1.6×10−31.6× 10^-3 for the 7B model, and 1.4×10−31.4× 10^-3 for the 14B model. The learning rate followed the Warmup-Stable-Decay (WSD) schedule (Xing et al., 2018; Hu et al., 2024) and decayed to 1% of its peak value. For nGPT, we use the proposed logarithmic decay schedule with ρ=0.05ρ=0.05 (see Figure 3). The learning rates of all parameters follow this global decay schedule, subject to the parameter-specific multipliers and warmups described below. In particular, (i) the learning rates of z s_z and all Mamba parameters except the output projections are warmed up during the first 10% of training, and (i) the Angular Step Cap for all normalized non-embedding vectors is increased from 0∘0 to 1.0∘1.0 over the same period. We use these warmups in all reported nGPT runs. Their endpoint and duration were selected in smaller-scale experiments and were not systematically retuned jointly with the global and parameter-specific learning rates for the present scaling ladder. The peak learning rate for nGPT is set to 0.24/dmodel0.24/ d_model and subsequently decays to zero. For parameter vectors in the routed and shared expert matrices, we multiply the learning rate by 1.51.5, while for router parameter vectors, we multiply it by 2.02.0. These multipliers were selected based on experiments with smaller models. They may be unnecessary; for example, some of their benefit might be recovered by adjusting the global peak learning rate. However, computational constraints did not allow us to test this possibility at scale. The baseline GPT uses ϵ=10−8ε=10^-8. For GatedAdam, we set ϵgate=10−8 _gate=10^-8 and ϵnum=10−14 _num=10^-14, and use β1=β2=0.975 _1= _2=0.975. We enable Second-Moment Growth Clipping after the first 500 optimizer steps and set R=100R=100. This deliberately permissive threshold is intended to suppress only extreme gradient spikes while leaving typical gradients unchanged. In our experiments, values as low as R=2R=2, as well as disabling clipping altogether, resulted in comparable final performance. We therefore regard this mechanism primarily as a conservative safeguard against rare optimization instabilities rather than as an essential component of the training recipe. Figure 6: Histogram of second-moment scales and the corresponding GatedAdam gate. Solid curves show the distributions of log10v _10 v across optimizer parameters for different model sizes using the left logarithmic axis. Dashed curves show the coordinate-wise GatedAdam gate on the right axis, ga(s)=σ(a[log(s+ϵnum)−log(ϵgate)])g_a(s)=σ\! (a [ (s+ _num)- ( _gate) ] ), where s=v^s= v. Vertical reference lines mark ϵgate _gate, ϵnum _num, and the medians of the distributions. The figure shows which parts of the optimizer state lie in the ϵε-dominated regime and how the choice of a changes the transition from suppressed to fully active updates. Figure 7: Training (left) and validation (right) losses as functions of the number of training tokens for the MoE model with 14B total parameters and 1.74B active parameters per token. We compare the GPT baseline trained with AdamW and its normalized nGPT counterpart trained with GatedAdam. Figure 8: Validation loss comparison between nGPT trained by GatedAdam with a=0.5a=0.5 and with a=1a=1 (gating of Adam) for models with 1B, 2B, 4B, and 7B total parameters. The a=0.5a=0.5 setting uses a peak learning rate of 0.24/dmodel0.24/ d_model, whereas the peak learning rate for a=1a=1 is tuned separately for each model size. The vertical axis reports La=1−La=0.5L_a=1-L_a=0.5, so positive values indicate an advantage for a=0.5a=0.5. Percentage labels report 100(La=1−La=0.5)/La=1100(L_a=1-L_a=0.5)/L_a=1, and filled pentagrams mark the best available a=1a=1 result for each model size. 4.2 Experimental Results Figure 5 and Table 1 show the training and validation losses obtained by GPT and nGPT for models ranging from 1B to 14B total parameters. nGPT consistently achieves losses that are approximately 2.5% lower. The 14B results for both methods deviate slightly from the overall scaling trend, potentially because this model has a different allocation of parameters across components and is trained for a shorter token horizon relative to its size. Figure 6 shows the distributions of the bias-corrected Adam second-moment scale v v used by GatedAdam at the end of training for the 4B, 7B, and 14B models. Most parameter coordinates in these distributions belong to MoE layers. The median v v decreases from 2.8×10−92.8× 10^-9 for the 4B model (dmodel=1280d_model=1280) to 1.4×10−91.4× 10^-9 for the 14B model (dmodel=2048d_model=2048). This is only a median-based estimate because the distribution is broad and the gate acts coordinate-wise. The distribution of v v is asymmetric, with a relatively flat left tail that is largely attributable to parameters in the input and output embeddings. As shown in Figure 6, coordinates with small v v can receive gate values that are numerically close to zero. Setting a to a smaller value, such as 0.50.5, increases the gate values, and therefore the effective update magnitudes, for coordinates with v^<ϵgate v< _gate. We also investigated scheduling ϵgate _gate as a function of model size or setting it based on tensor statistics. The results of this investigation are outside the scope of the current paper. Figure 7 shows the convergence curves for the GPT and nGPT versions of the MoE model with 14B total parameters and 1.74B active parameters per token. The shape of the GPT curve is strongly affected by the WSD schedule. The nGPT curve is characterized by a slower initial decrease followed by faster convergence later in training, consistent with the behavior observed in Loshchilov et al. (2024). This behavior is consistent with parameter normalization making the effective step size more directly controlled by the learning rate schedule. Training and validation losses are evaluated on different data distributions and therefore have different absolute scales; comparisons should be made between GPT and nGPT within each split. The figure shows that nGPT reaches the same training and validation losses using approximately half as many training tokens as the GPT baseline trained with AdamW. Gating with a=0.5a=0.5 versus a=1a=1. The main GPT–nGPT comparison changes both the model parameterization and the training recipe and therefore should not be interpreted as an isolated comparison between Adam and GatedAdam. To estimate the contribution of gating within the normalized model, Figure 8 compares nGPT with a=0.5a=0.5, using a peak learning rate of 0.24/dmodel0.24/ d_model, against nGPT with a=1a=1 (Adam’s gating), for which the peak learning rate is tuned separately at each model size. All other components of the nGPT training recipe are applied in both cases, so the a=1a=1 run is not a standalone Adam baseline. The advantage of a=0.5a=0.5 remains modest across model sizes and is substantially smaller than the total gap between GPT and nGPT reported in Table 1. This suggests that the overall improvement arises primarily from the normalized parameterization and the complete training recipe rather than from GatedAdam alone. At the matched validation-loss levels considered here, interpolation of the loss–token curves indicates that the a=0.5a=0.5 runs require approximately 15–20% fewer training tokens than the best available a=1a=1 runs. The ablation studies for Pre-Moment Tangent Projection and Post-Moment Exploration Noise are provided in Appendix B and Appendix D, respectively. Their effects on the validation loss are rather modest. 5 Conclusion This paper presents a practical recipe for training nGPT with modern hybrid Mixture-of-Experts models and demonstrates a substantial improvement in data efficiency. Further work may extend the evaluation to larger models and investigate hyperparameter scaling rules. Acknowledgments We thank Roger Waleffe for providing the source code and logs for the baseline GPT experiments with AdamW. We thank Mikail Khona, Kwangjun Ahn, and Roger Waleffe for providing the scaling ladder used in our experiments. Finally, we thank Mostofa Patwary, Mohammad Shoeybi, and the entire Nemotron Team (Blakeman et al., 2025) for their contributions to the development of the Nemotron open models. References A. Blakeman, A. Grattafiori, A. Basant, A. Gupta, A. Khattar, A. Renduchintala, A. Vavre, A. Shukla, A. Bercovich, A. Ficek, et al. (2025) Nemotron 3 nano: open, efficient mixture-of-experts hybrid mamba-transformer model for agentic reasoning. arXiv preprint arXiv:2512.20848. Cited by: §1, §2.1, §3, §4.1, Acknowledgments. M. Cho and J. Lee (2017) Riemannian approach to batch normalization. Advances in Neural Information Processing Systems 30. Cited by: Appendix B. T. Dao and A. Gu (2024) Transformers are SSMs: generalized models and efficient algorithms through structured state space duality. In Proceedings of the 41st International Conference on Machine Learning, Proceedings of Machine Learning Research, Vol. 235, p. 10041–10071. Cited by: §3.2. J. K. H. Franke, M. Hefenbrock, G. Koehler, and F. Hutter (2023) Constrained parameter regularization. arXiv:2311.09058. Cited by: §1. A. Henry, P. R. Dachapally, S. S. Pawar, and Y. Chen (2020) Query-key normalization for transformers. In Findings of the Association for Computational Linguistics: EMNLP 2020, p. 4246–4253. External Links: Document Cited by: §3.1. S. Hu, Y. Tu, X. Han, C. He, G. Cui, X. Long, Z. Zheng, Y. Fang, Y. Huang, W. Zhao, et al. (2024) Minicpm: unveiling the potential of small language models with scalable training strategies. arXiv preprint arXiv:2404.06395. Cited by: §4.1. T. Huang, Z. Zhu, G. Jin, L. Liu, Z. Wang, and S. Liu (2025) SPAM: spike-aware adam with momentum reset for stable LLM training. arXiv preprint arXiv:2501.06842. Cited by: Appendix C. T. Karras, M. Aittala, J. Lehtinen, J. Hellsten, T. Aila, and S. Laine (2024) Analyzing and improving the training dynamics of diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, p. 24174–24184. Cited by: §1. M. Khona, K. Ahn, R. Waleffe, M. Patwary, and M. Shoeybi (2026) Personal communication. Note: Personal communication Cited by: Figure 5, §4.1, §4.1, Table 1. M. Kodryan, E. Lobacheva, M. Nakhodnov, and D. P. Vetrov (2022) Training scale-invariant neural networks on the sphere can happen in three regimes. NeurIPS. Cited by: §1. A. Kosson, B. Messmer, and M. Jaggi (2023) Rotational equilibrium: how weight decay balances learning across neural networks. arXiv:2305.17212. Cited by: §1. W. Liu, Z. Liu, Z. Yu, B. Dai, R. Lin, Y. Wang, J. M. Rehg, and L. Song (2018) Decoupled networks. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, p. 2771–2779. Cited by: §1. W. Liu, Y. Zhang, X. Li, Z. Yu, B. Dai, T. Zhao, and L. Song (2017) Deep hyperspherical learning. Advances in neural information processing systems 30. Cited by: §1. Y. Liu, J. Bernstein, M. Meister, and Y. Yue (2021) Learning by turning: neural architecture aware optimisation. In International Conference on Machine Learning, p. 6748–6758. Cited by: §1. I. Loshchilov, C. Hsieh, S. Sun, and B. Ginsburg (2024) NGPT: normalized transformer with representation learning on the hypersphere. arXiv e-prints, p. arXiv–2410. Cited by: §1, §1, §4.2. I. Loshchilov and F. Hutter (2016) SGDR: stochastic gradient descent with warm restarts. arXiv preprint arXiv:1608.03983. Cited by: §2.2. I. Loshchilov and F. Hutter (2019) Decoupled weight decay regularization. In ICLR, Cited by: §1, §2.3. I. Loshchilov (2023) Weight norm control. arXiv preprint arXiv:2311.11446. Cited by: §1. P. Mettes, E. Van der Pol, and C. Snoek (2019) Hyperspherical prototype networks. NeurIPS. Cited by: §1. T. Q. Nguyen and J. Salazar (2019) Transformers without tears: improving the normalization of self-attention. In Proceedings of the 16th International Conference on Spoken Language Translation, Cited by: §3.1. T. Salimans and D. P. Kingma (2016) Weight normalization: a simple reparameterization to accelerate training of deep neural networks. NeurIPS. Cited by: §1. R. Waleffe, W. Byeon, D. Riach, B. Norick, V. Korthikanti, T. Dao, A. Gu, A. Hatamizadeh, S. Singh, D. Narayanan, et al. (2024) An empirical study of mamba-based language models. arXiv preprint arXiv:2406.07887. Cited by: §3. F. Wang, X. Xiang, J. Cheng, and A. L. Yuille (2017) Normface: l2 hypersphere embedding for face verification. In Proc. of the 25th ACM nternational conference on Multimedia, Cited by: §1. T. Wang and P. Isola (2020) Understanding contrastive representation learning through alignment and uniformity on the hypersphere. In ICML, Cited by: §1. C. Xing, D. Arpit, C. Tsirigotis, and Y. Bengio (2018) A walk with sgd. arXiv preprint arXiv:1802.08770. Cited by: §4.1. J. Xu and G. Durrett (2018) Spherical latent spaces for stable variational autoencoders. arXiv:1808.10805. Cited by: §1. Appendix A GatedAdamW’s gate without an explicit sigmoid The sigmoid-of-log gate has an equivalent power-law form. For each coordinate with dt,i>0d_t,i>0, γt,i _t,i =11+exp(−alogdt,iϵgate) = 11+ \! (-a d_t,i _gate ) (25) =11+(ϵgate/dt,i)a = 11+ ( _gate/d_t,i )^a (26) =dt,iadt,ia+ϵgatea. = d_t,i^\,ad_t,i^\,a+ _gate^\,a. (27) Thus, in vector notation, t=tata+ϵgatea, γ_t= d_t^\,a d_t^\,a+ _gate^\,a, (28) where powers, divisions, and additions are applied coordinate-wise. The adaptive part of the update can therefore be written without an explicit sigmoid as t⊙^t=ta−1ta+ϵgatea⊙^t. γ_t m_t d_t= d_t^\,a-1 d_t^\,a+ _gate^\,a m_t. (29) This form makes the power-law dependence on t d_t explicit. In particular, when a=1a=1, t⊙^t=^t+ϵgate. γ_t m_t d_t= m_t d_t+ _gate. (30) Hence, if ϵnum=0 _num=0 and ϵgate=ϵ _gate=ε, the adaptive update reduces to the standard AdamW update. Appendix B Pre-Moment Tangent Projection For parameter vectors constrained to the unit hypersphere, only the tangent component of the gradient contributes to the first-order change after normalization. We therefore optionally project the gradient onto the tangent space before the first- and second-moment updates of GatedAdamW. Tangent-space gradient projection is an established operation in Riemannian optimization methods for scale-invariant parameters (Cho and Lee, 2017). We do not claim the projection itself as a novel contribution; we use the descriptive term Pre-Moment Tangent Projection to emphasize that it is applied before both optimizer moments are updated and to provide a concise reference within this paper. For each normalized vector w and its gradient g, we replace ←−⟨,⟩max(⟨,⟩,ϵproj). g← g- w w, g ( w, w , _proj). (31) When ‖2=1\| w\|_2=1, this projection simply removes the radial component ⟨,⟩ w, g w, and the safeguard ϵproj=10−12 _proj=10^-12 is inactive. GatedAdamW’s moment estimates are then updated using the projected gradient. Thus, the moments do not accumulate radial gradient components that would subsequently be discarded by normalization. Unlike a fully Riemannian Adam update, this procedure does not parallel-transport the first-moment estimate after the parameter update. In the distributed implementation, the dot products and squared norms are computed over each complete normalized row or column vector, even when the vector is split across distributed optimizer shards. The required statistics are all-reduced before the local gradient shards are modified. Figure 9: Validation loss change after removing Pre-Moment Tangent Projection from nGPT for models with 1B and 7B total parameters. The default nGPT uses a peak learning rate of 0.24/dmodel0.24/ d_model, whereas for the case without projections the peak learning rate is tuned separately for each model size. Figure 9 illustrates the effect of Pre-Moment Tangent Projection. It slightly improves the final validation loss for the 1B model, while its effect is smaller for the 7B model. Appendix C Second-Moment Growth Clipping (SMGC) To limit the effect of isolated gradient spikes on the moment estimates, we optionally clip each scalar gradient coordinate using only its own second-moment history. Let viv_i denote the existing bias-uncorrected second-moment state for coordinate i, and let k be the number of completed optimizer steps. We define v¯i v_i =max(vi,(1−β2k)smin2), = (v_i,\, (1- _2^k )s_ ^2 ), (32) CR C_R =R−β21−β2, = R- _21- _2, (33) where smins_ is a floor expressed in bias-corrected v v units and R≥1R≥ 1 controls the maximum permitted one-step growth relative to v¯i v_i. The gradient is then replaced coordinate-wise by gi←sign(gi)min(|gi|,CRv¯i). g_i (g_i) (|g_i|,C_R v_i ). (34) Because vi≤v¯iv_i≤ v_i, the subsequent second-moment update satisfies vi+=β2vi+(1−β2)gi2≤Rv¯i. v_i^+= _2v_i+(1- _2)g_i^2≤ R v_i. (35) For coordinates above the floor, v¯i=vi v_i=v_i, and the rule directly bounds the growth of the second moment by vi+≤Rviv_i^+≤ Rv_i. For coordinates below the floor, the bound is instead defined relative to the floor. The clipped gradient is used to update both the first- and second-moment estimates. Clipping is applied after gradient unscaling and, when enabled, Pre-Moment Tangent Projection, but before the first- and second-moment updates of GatedAdamW. It may be activated only after a prescribed number of optimizer steps to allow the second-moment statistics to initialize. Unlike global norm clipping, SMGC uses no statistics from other coordinates or parameters. Because it operates coordinate-wise, however, it may change the direction of a multidimensional gradient. SMGC is closely related to the spike-aware gradient clipping used in SPAM (Huang et al., 2025), which identifies unusually large gradients relative to AdamW’s running second-moment estimate and rescales them before they enter the moment updates. Unlike SPAM, we do not periodically reset the first- and second-moment states. Instead, the clipping threshold is parameterized directly through the maximum permitted one-step second-moment growth factor R, together with a floor for coordinates with little second-moment history and an optional delayed activation period. Figure 10: Evolution of the number of parameter coordinates with exactly zero gradient for models with 1B, 2B, 4B, and 7B total parameters. Solid curves show the default nGPT runs with Post-Moment Exploration Noise (PMEN, λnoise=10 _noise=10), while dashed curves show otherwise matched runs without exploration noise (λnoise=0 _noise=0). The zero-gradient count is divided by dmodeld_model to facilitate comparison across model widths. The annotations show the validation loss change. Appendix D Post-Moment Exploration Noise (PMEN) We optionally inject a parameter-space perturbation after the optimizer moment update. Let ~j,t w_j,t denote the post-GatedAdamW candidate for a parameter vector j∈ℝdj w_j ^d_j, and let αj,t _j,t denote the current scheduled learning rate of parameter group j, including any parameter-specific multiplier and warmup. We associate each coordinate with an activity scale sj,i,t=vj,i,t, s_j,i,t= v_j,i,t, (36) where vj,i,tv_j,i,t is Adam’s raw, bias-uncorrected second-moment state after the current moment update. Let ψ(s;τv)≥0ψ(s; _v)≥ 0 be a non-increasing noise profile and let cj,tc_j,t be a reference coordinate scale. At selected optimizer steps, we draw independent Rademacher signs ξj,i,t∈−1,+1 _j,i,t∈\-1,+1\ and apply ~j,t←~j,t+αj,tλnoisecj,t(j,t⊙j,t), w_j,t← w_j,t+ _j,t _noisec_j,t ( ψ_j,t ξ_j,t ), (37) where [j,t]i=ψ(sj,i,t;τv). [ ψ_j,t]_i=ψ(s_j,i,t; _v). (38) Here, λnoise _noise controls the overall perturbation strength, ψ distributes it across second-moment scales, and cj,tc_j,t specifies its reference scale. A sparse hard-threshold variant is obtained with ψstep(s;τv)=[s≤τv]. _step(s; _v)=I[s≤ _v]. (39) Setting τv=0 _v=0 restricts the perturbation to coordinates whose second-moment state is exactly zero. A smooth alternative is ψb(s;τv)=11+(s/τv)b,τv>0, _b(s; _v)= 11+ (s/ _v )^b, _v>0, (40) where b>0b>0 controls the sharpness of the transition. As b→∞b→∞, this profile approaches the hard-threshold rule, apart from its value exactly at s=τvs= _v. One possible reference scale is cj,t=1dj, c_j,t= 1 d_j, (41) which normalizes the perturbation at the vector level. An Adam-relative alternative is cj,t=cAdam=1−β11+β1, c_j,t=c_Adam= 1- _11+ _1, (42) which approximates the coordinate-wise root-mean-square normalized Adam update under stationary random gradients. In this parameterization, λnoise _noise measures the perturbation relative to the coordinate-wise stochastic update scale of Adam. For the hard-threshold profile, if kjk_j coordinates are selected, then ‖Δj,tnoise‖22=αj,t2λnoise2cj,t2kj. \| w_j,t^noise \|_2^2= _j,t^2 _noise^2c_j,t^2k_j. (43) For a unit-normalized vector and a sufficiently small perturbation, the first-order angular displacement is bounded by this Euclidean perturbation norm. The perturbation is applied after the Adam moment update and the GatedAdamW parameter correction, and therefore does not enter either moment estimate. For normalized non-embedding parameters covered by ASC, the Angular Step Cap is subsequently applied to the combined optimizer and noise displacement, followed by the usual nGPT normalization. PMEN may be delayed, applied periodically, and enabled or disabled for embedding and output matrices. Figure 10 illustrates the effect of PMEN. Solid curves denote GatedAdam runs with λnoise=10 _noise=10, the value used in our baseline experiments, whereas dashed curves denote otherwise matched noise-free runs with λnoise=0 _noise=0. The runs with PMEN achieve slightly lower validation loss, although the improvement is small and does not exceed 0.1%0.1\%. Without PMEN, the number of parameter coordinates with exactly zero gradients increases with model size even after normalization by dmodeld_model. With PMEN, the normalized count remains approximately stable, with only a small increase for the largest model. At the end of training, approximately 1.4%1.4\% of the input-embedding parameter coordinates in the 4B model have v^=0 v=0, indicating that their Adam second-moment states have never received a nonzero gradient contribution. For the input embeddings, this is consistent with some token IDs never appearing in the training data. Approximately 0.7%0.7\% of the routed-expert parameter coordinates also have v^=0 v=0. PMEN is intended to reduce the likelihood that weakly updated parameter coordinates become permanently inactive and to preserve their opportunity to receive useful gradients later in training. Although PMEN reduces both the number of coordinates with zero gradients and the number with v^=0 v=0, we do not yet observe a substantial improvement in final loss.