Token Geometry

Kathan Shah

Embedding and LM-head gradients need only O(V+D) optimizer state: Ember matches Adam while shrinking memory by orders of magnitude and revealing that token trajectories are mostly a straight ray.

Reviewed by Pith at T0; open to challenge. T0 means a machine referee read the full paper against a public rubric. the ladder, T0–T4 →

Challenge this review Re-run · record.json Download PDF Read on arXiv ↗

T0 review

2026-07-12 08:47 UTC pith:SUOGEZM2

load-bearing objection Practical memory win for embeddings/LM-heads with a coherent Fisher story; Adafactor-adjacent but token-specific and well-supported at the scales they actually ran. the 3 major comments →

arxiv 2607.01455 v3 pith:SUOGEZM2 submitted 2026-07-01 cs.LG cs.AI

Token Geometry

Kathan Shah This is my paper

classification cs.LG cs.AI

keywords Ember optimizertoken embeddingsLM-headFisher metricrank-1 second momentoptimizer memoryTransformer trainingsquare-root Fisher

verification ladder T0 review T1 audit T2 compute T3 formal T4 reserved

The pith

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

Language models interface discrete tokens with continuous weights through the embedding table and LM-head. This paper argues those two matrices have a different gradient geometry from dense linear layers: their per-parameter second moments are essentially rank-1, so a cheap outer product of row and column energy buffers recovers the same square-root Fisher preconditioning that Adam uses. The resulting optimizer, Ember, drops Adam's first-moment buffer entirely, keeps only O(V+D) state instead of O(2VD), and matches or slightly beats AdamW across supervised fine-tuning, GRPO-style RL, and pretraining while remaining stable at low batch sizes. As a side observation, the path each token embedding takes during training collapses almost entirely onto a single linear ray (roughly 90 percent of trajectory variance), suggesting these parameters do not wander through a highly nonconvex landscape. The practical payoff is that token-table optimizer state no longer needs to be sharded and can fit in kilobytes even for multi-billion-parameter models.

Core claim

The second-moment geometry of embedding and LM-head gradients is well approximated by a rank-1 outer product of token-wise (row) and feature-wise (column) energies; conditioning the instantaneous gradient by the square root of that factored matrix recovers Adam-quality updates with O(V+D) state and no first-moment buffer, while the resulting token trajectories are largely one-dimensional rays.

What carries the argument

Ember's factored square-root Fisher preconditioner: bias-corrected EMA row and column second-moment vectors whose geometric-mean-normalized outer product supplies a per-entry denominator, so each update is a unitless z-score of the raw gradient.

Load-bearing premise

The dense per-parameter squared-gradient matrix of an embedding or LM-head is accurately recovered by a simple outer product of one vector per token and one vector per feature; residual cross-correlations are small enough not to matter.

What would settle it

On a held-out large-vocabulary model, measure the pointwise correlation between Ember's factored second-moment estimate and the true dense E[g_ij^{2}]; if the correlation collapses or validation loss systematically exceeds AdamW by more than seed noise, the rank-1 claim fails.

Watch this falsifier. Get emailed when new claim-graph text bears on it.

Share X Bluesky LinkedIn Reddit HN

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, simulated authors' rebuttal, and a circularity audit.

Desk Editor's Note

Practical memory win for embeddings/LM-heads with a coherent Fisher story; Adafactor-adjacent but token-specific and well-supported at the scales they actually ran.

read the letter

The useful takeaway is simple: Ember drops Adam’s first moment and replaces the dense second-moment buffer on embedding and LM-head tables with a rank-1 outer product of row and column factors (O(V+D) state, no sharding), and across the runs they report it matches or slightly beats AdamW on SFT, FineWeb pretraining, GRPO, and LlamaGen image-AR while cutting token-interface optimizer state by three to four orders of magnitude.

What is actually new is not the factored second moment itself—Adafactor already lives there, and the paper says so—but the token-specific geometry: one-hot participation makes row energy a clean proxy for pi, the square-root Fisher then yields isotropic updates, the LM-head is treated as a smoothed embedding, and they give a clean unit-correct geometric-mean normalization. The distributed ZeRO/FSDP-friendly code is real and immediately usable. The SVD trajectory observation (PC1 ~90% variance, monomial structure) is a nice empirical curiosity even if it is currently Ember-specific and observational.

Soft spots are real but proportionate. Most runs are short (250–300 steps) and modest in scale (up to ~3B); multi-seed error bars are sparse. The load-bearing rank-1 claim rests mainly on one 500-step gpt2-small collection (corr 0.974). That is thinner than ideal, especially for the low-batch and RL regimes where they advertise an edge, but the derivation is transparent and the residual is not hidden. Novelty relative to Adafactor is moderate; the framing and the token-table focus are the differentiators. Citations and math look solid; free parameters are the usual β2/α/ε.

This is for people who train or serve models ≤7B, multimodal setups with multiple tables, or anyone tired of sharding embedding optimizer state. It does not rewrite learning theory. I would send it to peer review: the engineering result plus open code and a coherent geometric story deserve referee time, with the natural ask for longer multi-seed runs and a denser residual check. Worth engaging.

Referee Report

3 major / 7 minor

Summary. The paper argues that embedding and LM-head matrices have gradient geometry distinct from dense linear layers, and exploits this with Ember: a square-root-Fisher-style optimizer that drops Adam’s first-moment buffer and replaces the dense second-moment matrix with a bias-corrected rank-1 outer product of O(V) row and O(D) column second-moment estimates (Alg. 1), using O(V+D) state instead of O(2VD). Empirically, Ember matches or slightly improves AdamW on token tables across GPT-2/Pythia SFT (Table 1), FineWeb pretraining (Figs. 1–2), GRPO RL on Qwen2.5-3B (Fig. 6), batch-size sweeps (Fig. 8), and LlamaGen image-AR (Fig. 9), while cutting token-interface optimizer state by orders of magnitude. The authors further report that token trajectories under Ember admit a near-1D ray structure (PC1 ≈ 90% variance; Figs. 5, 7) and frame competitive Transformer optimizers as square-root Fisher preconditioners in different bases.

Significance. The practical contribution is clear and timely: token-table Adam state is a genuine VRAM and sharding burden (e.g., multi-GB at multi-billion scale; Fig. 4), and a drop-in O(V+D) alternative that preserves quality would simplify single-GPU iteration and ZeRO/FSDP setups. The open-sourced distributed implementation is a concrete strength. Empirically, parity-or-better results across SFT, RL, pretraining, and vision-AR, plus stability at very low batch sizes, would improve the memory–quality Pareto frontier for a parameter class that has seen little optimizer progress. The Fisher-based unification and the trajectory SVD analysis are interesting if they survive stronger controls; even as supporting narrative they help explain why factored second moments and no first-moment state can suffice for embeddings.

major comments (3)

App. 6.7–6.8 and Fig. 12 are load-bearing for the claim that Ember recovers Adam-quality square-root Fisher geometry via E[g_ij²]≈R_i C_j. The only quantitative residual analysis is a single 500-step gpt2-small FineWeb collection (pointwise corr. 0.974, OLS slope 0.93). The derivation itself notes the factorization is exact only to first order under uncorrelated energies. Please report the same dense-vs-factored comparison at least at (i) larger model scale, (ii) longer training, and (iii) the low-batch / sparse-gradient regime where Ember’s advantage is advertised (Fig. 1B, Fig. 2), so that residual Cov(δ_i²,x_j²) growth can be ruled out as a confounder of the “matches AdamW” result.
§3.2 and Figs. 5, 7 claim that token optimization trajectories are “well described by a simple 1D ray… counter to the popular belief that neural net parameters navigate a heavily nonconvex landscape.” The SVD analysis is performed under the Ember update rule only. Without the same PC energy spectrum for AdamW (and ideally SGD) on identical runs, it is unclear whether the 1D structure is a property of token-table geometry or an artifact of Ember’s isotropic 1/p-style preconditioning. A side-by-side AdamW trajectory decomposition is needed for the landscape claim as stated in the abstract and §3.2.
The main quality claims rest on short horizons: SFT for 250 steps (Table 1, §6.3), batch-size sweeps for 300 steps with Val@200 (Fig. 8, §6.3), and limited detail on FineWeb pretraining length. Table 1 also shows a clear Ember regression on pythia-160m (+0.035). To support “matches or improves AdamW” as a general replacement rather than a short-run phenomenon, please include longer SFT/pretraining curves (or intermediate checkpoints with error bars) at least for one GPT-2 and one Pythia scale, and discuss the 160M outlier.

minor comments (7)

Alg. 1 line 13 says “Ascend the reward R_t”; for supervised CE this is descent of loss. Clarify the sign convention or split SFT vs RL notation.
Fig. 2 and the row-only ablation are important: at bs=8, Ember row-only is noticeably worse than full Ember/Adam (7.665 vs 7.219). The paper focuses analysis on row-only; state more clearly when the column factor is required in practice.
§1.1 / App. 6.6 comparison to Adafactor is useful but should list the exact Adafactor config used in Fig. 2 (relative_step, scale_parameter, clip_threshold, decay schedule) so the 0.25-nat gap is reproducible.
Abstract and §5 claim “improve the Pareto frontier” and “universally matches or improves”; Table 1 and Fig. 2 show mostly parity with occasional small wins and one clear loss. Soften absolute language to match the tables.
Fig. 1 caption and body refer to “Fig 1B/C/D” without fully specifying axes/metrics in the caption; add loss/reward definitions and seeds.
Typos/clarity: “forgoes the need to shard both token table optimizer states”; “we open-source… available atgithub.com” (missing space); DeepSeek V4 / 2026 citations look like placeholders—verify bibliography.
§2.2 “Scale with 1/p” vs “Scale isotropically”: the Fisher scaling argument (App. 6.9) is clear, but a one-line pointer from Alg. 1’s row factor to F^{-1/2}∝1/p would help readers connect the algorithm to the metric discussion.

Circularity Check

0 steps flagged

No circular derivation: Ember follows from Fisher geometry plus the exact rank-1 outer-product structure of embedding/LM-head gradients, then is tested empirically; the 1-D ray is a post-hoc SVD observation, not a forced prediction.

full rationale

The paper’s load-bearing chain is self-contained and non-circular. It begins from the classical Fisher information identity for cross-entropy (second derivative equals squared score near the optimum; Chentsov uniqueness of the Fisher metric for KL), notes that Adam already approximates the square-root empirical Fisher, and observes that embedding and LM-head gradients are exactly rank-1 outer products (g_ij = δ_i x_j with x one-hot or soft-max error). Consequently E[g_ij²] factors as R_i C_j to first order, licensing an O(V+D) outer-product second-moment buffer (Alg. 1, App. 6.7–6.8). The geometric-mean normalization is derived from unit matching and variance optimality under that factorization, not fitted to the target losses. First-moment removal is justified by a separate empirical autocorrelation measurement (Fig. 11), not by definition. All performance claims (SFT, GRPO, FineWeb, image-AR) are ordinary held-out comparisons against AdamW/Adafactor; none are forced by construction. The “1-D ray” claim is an observational SVD of recorded trajectories under Ember (PC1 ≈ 90 % variance, monomial structure), not a prediction derived from a fitted parameter. Citations (Adam, Muon, Adafactor, Amari, Chentsov) are external; there is no self-citation chain, no uniqueness theorem imported from the same author, and no ansatz smuggled via prior work. Mild interpretive reframing of existing optimizers as “square-root Fisher” does not make the empirical results circular. Residual cross-correlation is acknowledged and measured once (corr 0.974 on a short gpt2-small run); that is a possible correctness limitation, not circularity.

Axiom & Free-Parameter Ledger

3 free parameters · 6 axioms · 1 invented entities

The central claim rests on standard information-geometry facts (Fisher = second-order KL term, Chentsov uniqueness), the empirical success of Adam-style square-root second-moment preconditioning, and the modeling assumption that embedding/LM-head squared gradients factor as rank-1 outer products of participation and feature energies. Free parameters are essentially the usual Adam β2 and learning rate; no new physical constants or fitted scales are introduced. Ember itself is an algorithmic construction, not a new ontological entity.

free parameters (3)

β2 (second-moment EMA decay) = 0.999
Lifted from Adam default 0.999; controls how much history enters the row/column curvature estimates. Ablated but still a free hyperparameter.
learning rate α / trust-region scale = task-dependent (e.g. 1e-3, 2e-5)
Step size for the z-scored update; treated as a free hyperparameter shared with Adam baselines (e.g. 1e-3 or 2e-5 depending on experiment).
ε numerical stabilizer = 1e-8
Standard 1e-8 floor under the square root; minor but present.

axioms (6)

standard math For cross-entropy, the Fisher information equals the expected squared score and coincides with the Hessian near an optimum (regularity assumptions).
Invoked in §1 and App. 6.1 to justify second-moment preconditioning.
standard math Chentsov’s theorem: the Fisher metric is the unique canonical local geometry for the KL divergence (up to scaling).
Used to argue that competitive CE optimizers must essentially use Fisher geometry.
domain assumption Adam’s per-parameter second-moment EMA is a serviceable empirical square-root Fisher for all parameter classes, including embeddings.
Stated in §2.1 as the starting point for simplifying the second-moment buffer.
domain assumption Embedding and LM-head per-example gradients are rank-1 outer products, and the expected squared gradient factors as E[δ_i²] E[x_j²] with negligible residual correlation.
Core modeling step in App. 6.7–6.8 and Fig. 12 that licenses O(V+D) state.
domain assumption Token participation frequency p_i makes the row-wise Fisher scale as p_i², so square-root Fisher yields p-independent (isotropic) updates.
Derived under a coherent-signal binomial model in App. 6.9; used to interpret Ember’s row factor.
ad hoc to paper First-moment (momentum) EMA carries little useful signal on token tables because temporal gradient autocorrelation is near zero or negative.
Empirical observation in Fig. 11 used to justify dropping Adam’s first-moment state entirely.

invented entities (1)

Ember optimizer (row/column factored second-moment preconditioner without first-moment state) independent evidence
purpose: Replace Adam on embedding and LM-head matrices with O(V+D) state while preserving or improving loss.
Algorithmic construction; independent evidence is the empirical match to Adam across tasks, not an external physical prediction.

reviewed 2026-07-12 · how reviews work

0 comments

Cite this review

Pith. "Pith review of Token Geometry." pith.science (2026). https://pith.science/paper/SUOGEZM2

@misc{pith2026260701455,
  author       = {Pith},
  title        = {Pith review of: Token Geometry},
  year         = {2026},
  howpublished = {\url{https://pith.science/paper/SUOGEZM2}},
  note         = {Machine review of arXiv:2607.01455}
}

Share X Bluesky LinkedIn Reddit HN

read the original abstract

Language models learn continuous programs over discrete symbols, with the embedding table and LM-head acting as the read/write interface between them. We show that this interface has gradient geometry distinct from dense hidden weights which can be exploited to improve the Pareto frontier across supervised finetuning, RL, and pretraining, while only utilizing kilobytes of optimizer state. We introduce Ember, a lightweight optimizer for embedding and LM-head matrices that utilizes O(V + D) VRAM, instead of Adam's O(2VD), and forgoes the need to shard both token table optimizer states. We provide empirical evidence that Ember scales effectively across batch size and parameter count. We show that the optimization trajectory of tokens can be well described by a simple 1D ray, counter to the popular belief that neural net parameters navigate a heavily nonconvex landscape. We provide a principled view on the surprisingly narrow space of optimizers that suffice for Transformer training. Finally, we open-source our distributed Ember implementation that merges cleanly with existing ZeRO/FSDP setups to support further research at https://github.com/katop1234/ember