Pith. sign in

REVIEW 2 major objections 4 minor 27 references

Geometry-Guided Layerwise FFN Width Allocation in Transformers

T0 review · 2 major / 4 minor · reviewed 2026-08-04 · deepseek-v4-flash

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

Pith's one-line read Feed-forward layer width in Transformers can be set by measuring how much each layer moves and reshapes token representations, and the paper shows such geometry-guided schedules beat uniform width and a cosine taper at matched budgets.

desk verdict Solid, honest method paper whose headline 440M result is descriptive until paired-difference variability is released; worth refereeing. read the letter →

arxiv 2608.02064 v1 pith:SNTZPZRF submitted 2026-08-03 cs.LG cs.AIcs.CL

Timur Mudarisov , Mikhail Burtsev , Radu State This is my paper
classification cs.LGcs.AIcs.CL
keywords Transformerfeed-forwardnetworkwidthallocationgeometricworkGromov-Wassersteinpersistenthomologyintrinsicdimension
verification ladder T0 review T1 audit T2 compute T3 formal

The pith

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

The reading

Transformers conventionally give every feed-forward layer the same hidden width, even though the work each layer does changes with depth. This paper proposes a data-dependent alternative: measure how each layer transforms the cloud of token representations — by mean shift, by Gromov-Wasserstein distortion of pairwise distances, and by degree-one persistent homology — and reallocate a fixed width budget across layers using an exact optimizer built on a layerwise approximation surrogate. On seven pretrained language models the scale-normalized work is concentrated in early layers, and in matched training runs at 128M, 256M, and 440M parameters, geometry-derived schedules reduce mean validation loss relative to uniform width. At 440M the best schedule lowers mean loss by about 0.019 nats versus uniform, roughly six times the mean reduction from a hand-designed cosine taper, though the paper presents this as a descriptive mean difference without paired-difference confidence intervals.

What carries the argument

The argument is carried by a layerwise approximation surrogate. A latent coefficient C_l represents how hard the layer's map is to approximate on the token distribution, lambda_l is the loss sensitivity to that layer's error, and the intrinsic dimension d_int(l) sets the error-width exponent a_l = s/d_int(l) with s=3. Geometric work measured from a forward pass is substituted for C_l, and a constrained convex program yields a unique continuous optimal width profile w*_l = max{w_min, (a_l c_l / mu*)^(1/(a_l+1))}, where mu* enforces the total budget. This formula, together with a floor width and a rounding-and-rebalancing step, turns any measured work profile into an iso-parameter architecture

What would settle it

Train a set of models where only one FFN layer's width is varied at a time while all others stay fixed, fit the validation-loss or approximation-error curve as a function of width for each layer, and compare the fitted coefficients with the geometric work profile. If the per-layer coefficients do not rank layers the same way the work does, or if the profile measured on a uniform model changes sharply once widths are reallocated, the central proxy assumption collapses.

Watch

Extended reading notes

Core claim

The paper's central discovery is that a feed-forward layer's effect on the representation cloud is a usable signal for how much width that layer should receive. Using correspondence-preserving shift, Gromov-Wasserstein distortion, and persistent-homology distance under raw, spherical, and hyperbolic geometries, the authors show that normalized work profiles are front-loaded across model families, and that substituting these measurements into an exact fixed-budget allocation rule produces schedules that lower validation loss compared with uniform width and, in the 440M five-seed comparison, compared with a hand-designed cosine taper. The paper is explicit that the headline 440M result is a di

Load-bearing premise

The load-bearing premise is that the measured geometric work of each layer, obtained from a uniform-width reference model, is an informative proxy for that layer's true approximation difficulty, and that this ordering still holds after the model is retrained with reallocated widths.

Editorial extensions

If this is right

  • If the claim holds, Transformer feed-forward width becomes a per-layer decision variable rather than a global constant, with no change to total parameter or FLOP budget.
  • A measured schedule can express non-monotone patterns that a cosine taper cannot, adapting to architecture, scale, and data distribution.
  • Scale-normalized geometric statistics (spherical and hyperbolic) are the more reliable basis for allocation; raw statistics predominantly track residual-norm growth and can mislead.
  • The allocation rule is exact for the continuous surrogate and handles depth-dependent exponents and lower bounds, so it transfers to models with different depths and width granularities.
  • The 440M comparison suggests that activation-derived schedules can improve on a tuned taper by a wide margin, once paired-difference uncertainty is supplied.

Reading between the lines

Editorial extensions of the paper, not claims the author makes directly.

  • A direct layerwise width sweep — varying one layer's width at a time and fitting the validation-loss or approximation-error curve — would calibrate geometric work to the true approximation cost; the paper itself names this as its central missing experiment.
  • The same cloud-transport statistics could plausibly guide other capacity allocations, such as attention-head width or key-value dimensions, wherever a residual-update view of a sublayer applies.
  • The transfer rule assumes the reference model's work profile remains informative after widths are reallocated and the model is retrained; testing how much the profile shifts under reallocation would clarify whether the method needs an iterative update.
  • A matched control that reverses a normalized profile while holding geometry and width histogram fixed would isolate the effect of ordering more cleanly than the paper's anti-topological/raw control, which changes both direction and metric.
Share X Bluesky LinkedIn Reddit HN

Editorial analysis

A structured set of objections, weighed in public.

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

Referee Report

2 major / 4 minor

Summary. The paper proposes a method for allocating FFN width across Transformer layers using forward-pass geometric measurements. It defines three work statistics (shift, Gromov-Wasserstein distortion, and degree-one persistent homology) under raw, spherical, and hyperbolic geometries, and derives an exact continuous fixed-budget optimizer based on a layerwise approximation surrogate. The surrogate assumes local additive loss, an intrinsic-dimension-dependent error-width exponent, and that geometric work is a valid proxy for the latent approximation coefficient. The method is evaluated on seven pretrained LLMs for diagnostics and with from-scratch paired training at 128M, 256M, and 440M scales, comparing against uniform width, a TLM-style cosine taper, and an anti-topological control. The central reported result is that at 440M, topological/hyperbolic and Gromov/spherical allocations reduce mean validation loss by 0.019 and 0.016 respectively relative to uniform, larger than the cosine taper's 0.003 mean reduction. The paper is careful in the body to describe these as descriptive mean comparisons, acknowledging that paired-difference variability is not reported.

Significance. If the empirical claims hold, the paper offers a principled, data-dependent alternative to hand-designed width tapering, with a clean convex formulation and a plausible link between activation geometry and approximation demand. The work is strengthened by its honest diagnostics: the local additivity check (R²≈0.72), the explicit separation of raw scale sensitivity from normalized reshaping, and the repeated acknowledgement that the core proxy assumption is not directly calibrated. The release of reproducible details and code is a further strength. However, the load-bearing empirical claim at 440M currently lacks inferential support, and the abstract states it more strongly than the body's own caveats. The paper is therefore a valuable contribution whose central claim needs additional statistics and more precise framing.

major comments (2)
  1. [§3.5, Table 1; Appendix A.2] The headline 440M result is a mean paired difference Δ of −0.019 for topological/hyperbolic and −0.016 for Gromov/spherical, but the table deliberately omits the standard deviation of the paired differences D_s (or the five D_s values). With five seeds and marginal SDs of 0.022–0.028, the reported mean gaps could plausibly lie within seed noise; the paper itself states that 'marginal standard deviations do not determine s_D' and that inferential claims require the paired interval. This is a load-bearing gap because the abstract states that geometry-based allocations 'improve over uniform substantially larger than the cosine taper,' which reads as a comparative superiority claim rather than a purely descriptive statement about means. Please report the paired differences or their SD/CI, and if the resulting interval includes zero, revise the abstract and conclusion to explicitly limit the
  2. [§2.3, Assumption 3; §2.5; §4] The entire allocation rule rests on Assumption 3, that the measured work Ŵ_ℓ is informative about the latent approximation coefficient C_ℓ in layerwise ordering, and on the additional assumption that this ordering transfers from a uniform-width reference model to a reallocated model trained from scratch. The paper explicitly identifies 'the absence of a direct layerwise width sweep' as the central limitation, and the diagnostics (sensitivity association, additivity) are indirect. Because the title and contributions claim 'geometry-guided' allocation, the mismatch between the strength of the claim and the indirectness of the evidence should be addressed. Either provide a calibration experiment (even on a single model) relating Ŵ_ℓ to per-layer error-width curves, or restate the conclusions so that the empirical contribution is framed as 'work profiles can predict useful allocation schedul
minor comments (4)
  1. [§3.1, §4] The anti-topological/raw control changes both the direction of the profile and the geometry, so it cannot isolate whether alignment direction or the raw geometry is responsible for the worse-than-uniform result. The paper acknowledges this; consider adding a matched reversal (same estimator, reversed order) in a future revision or at least explicitly noting that the control does not separate the two factors.
  2. [§2.2, Eq. (6)] The hyperbolic construction is acknowledged to be a fixed monotone transform of angular separation. This is fine, but the terminology 'hyperbolic geometry' may overstate the novelty; consider calling it 'angular distance with a nonlinear reweighting' or clearly stating in the main text that no radial hierarchy is learned.
  3. [§3.2, Fig. 1] The R² ≈ 0.72 for the additivity panel is described as 'positively associated'; this is honest, but the caption and text could more explicitly note that the scatter is not centered on the identity line, so the additive surrogate is an approximation, not an exact identity.
  4. [§3.5] The ratio 'approximately 6.3× the mean reduction of cosine' is a ratio of two mean differences, each without uncertainty. Even as a descriptive statement, this ratio is fragile; please report it with a caveat or derive it from paired differences only after reporting their distribution.

Circularity Check

0 steps flagged · score 0.0 of 10

No significant circularity: width schedules are derived from forward-pass measurements and a convex program, and validation loss is never used to fit the work profile or the optimizer's free parameters.

full rationale

The paper's derivation chain is self-contained and not circular. The measured work W_hat_l is computed from forward passes of a reference model under Equations (2)–(11), and Assumption 3 explicitly states that W_hat_l is a proxy for the latent coefficient C_l; this is an acknowledged modeling hypothesis, not a quantity fitted to the target. The width allocation w*_l follows from a convex fixed-budget program whose unique minimizer is proven in Lemma 1 with hand-fixed s=3 and lambda_l=1; the validation loss is not used to set the work profile, the exponent, the normalization, or any optimizer parameter. The target endpoint appears only after schedules are fixed, in paired training runs that compare the resulting architectures. There are no load-bearing self-citations: the only external TLM citation is used as a structural baseline, not as an authority for the geometry-to-width link, and no uniqueness claim is imported from the authors' prior work. The approximation ansatz of Assumption 2 is explicitly labeled an ansatz rather than smuggled via citation. The acknowledged gaps—missing paired-difference standard deviations and the unvalidated surrogate W_hat_l≈C_l—are correctness or robustness concerns, not circularity, because the measured inputs and the final validation outcomes are distinct and not linked by construction.

Assumptions & free parameters 6 free parameters · 5 assumptions · 2 invented entities

The central claim rests on four stated surrogate and transfer assumptions rather than on hidden derivations. The free parameters (s=3, lambda_l=1, epsilon_GW, w_min, rho_c, and per-layer TwoNN dimensions) are either hand-chosen or data-fitted on the reference model; none are fit to the final validation loss, which is why the circularity burden stays low. C_l is the main invented latent, with no independent evidence; lambda_l is a measurable latent that the paper chooses not to use in the actual schedules.

free parameters (6)
  • s (global expansion order) = 3
    Used in a_l = s/d_int(l) for all layers and models (Section 2.3); fixed by hand with no sensitivity analysis, so the exponent and therefore the allocation depend on this choice.
  • lambda_l (loss sensitivity weight) = 1 for all layers
    The applied schedules use the pure-work variant with lambda_l = 1 (Section 2.5); perturbation-based lambda estimates are kept as a diagnostic only, so the multiplicative sensitivity factor is a modeling choice.
  • epsilon_GW (Sinkhorn regularization) = not reported
    Fixing epsilon_GW is required in Eq. (9); the manuscript states it is fixed across layers, geometries, and models but does not give its value, so the Gromov-Wasserstein work numbers cannot be exactly reproduced.
  • w_min (minimum width floor) = not reported
    The lower bound w_min is a required input to the optimizer in Lemma 1 and affects which layers sit at the floor; its value is not stated.
  • rho_c (hyperbolic clipping radius) = 0.99
    Unit-norm states are clipped to radius 0.99 inside the Poincare ball (Eq. 4); the authors note the construction is a monotone transform of angular distance, so this constant does not change layer ordering but is still a hand choice.
  • d_int(l) (TwoNN intrinsic dimension) = per-layer estimates, roughly 1-5
    Each layer's exponent theta_l = d_int(l)/(3+d_int(l)) is fit from activation clouds via TwoNN (Section 3.1); these data-fitted values directly reshape the allocation and are used without clipping or smoothing.
assumptions (5)
  • domain assumption Assumption 1: local additive loss, Delta L approx sum_l lambda_l epsilon_l
    Section 2.3; tested indirectly by joint-vs-sum perturbation with R2 about 0.72, not proven.
  • domain assumption Assumption 2: error-width rate epsilon_l(w) approx C_l w^{-a_l} with a_l = s/d_int(l), s=3
    Section 2.3; cites nonparametric approximation theory but states explicitly it is an ansatz for Transformer FFNs, not a theorem.
  • ad hoc to paper Assumption 3: geometric proxy W_hat_l approx C_l in layerwise ordering
    Section 2.3; the load-bearing link between measured work and the allocation surrogate; the paper calls it a modeling hypothesis and admits it is model- and distribution-dependent.
  • ad hoc to paper Profile transfer: reference-model work remains informative after reallocation and from-scratch training
    Section 2.5; explicitly assumed so that schedules measured on a uniform model apply to the reallocated target.
  • domain assumption Normalization invariance of LayerNorm/RMSNorm motivates normalized metrics
    Section 2.2; Norm(cz) approx Norm(z) is used to justify spherical/hyperbolic work as readout-aligned; the paper notes spherical normalization is not identical to full LayerNorm.
invented entities (2)
  • Layerwise approximation coefficient C_l
    purpose: Latent difficulty of approximating the FFN branch at layer l; sets the scale of the error-width curve and is replaced by measured work in the allocation
    Introduced in Section 2.3 (Eq. 12); never directly measured, so it has no falsifiable handle outside the paper's surrogate framework.
  • Loss sensitivity lambda_l independent evidence
    purpose: Converts local approximation error into loss effect inside the additive surrogate
    Estimated independently via low-rank perturbation experiments (Section 3.1), so it has an external handle; however, the applied schedules set lambda_l=1 and do not use these estimates.

how reviews work

0 comments
Cite this review

Pith. "Pith review of Geometry-Guided Layerwise FFN Width Allocation in Transformers." pith.science (2026). https://pith.science/paper/SNTZPZRF

@misc{pith2026260802064,
  author       = {Pith},
  title        = {Pith review of: Geometry-Guided Layerwise FFN Width Allocation in Transformers},
  year         = {2026},
  howpublished = {\url{https://pith.science/paper/SNTZPZRF}},
  note         = {Machine review of arXiv:2608.02064}
}
read the original abstract

Feed-forward networks (FFNs) account for a large fraction of Transformer parameters, yet their hidden width is usually constant across depth. We ask whether this capacity can instead be allocated from a forward-pass measurement of layer behavior. We view each FFN as transporting a cloud of token representations and quantify the induced geometric change using correspondence-preserving shift, Gromov-Wasserstein distortion, and degree-one persistent homology under raw and scale-normalized metrics. A layerwise approximation surrogate yields an exact fixed-budget optimizer. Across seven pretrained language models, raw Euclidean work largely tracks residual-norm growth, whereas normalized work is predominantly front-loaded. Gromov-Wasserstein work is more consistently associated with perturbation-based layer sensitivity than the finite-sample topological estimate. In paired 128M and 256M training runs, several normalized-work schedules reduce mean validation loss relative to both uniform width and a hand-designed cosine taper. With the amplified paired differences at 440M, the best geometry-based allocations improve over uniform substantially larger than the cosine taper, while the anti-topological raw control is worse than uniform.

Figures

Figures reproduced from arXiv: 2608.02064 by the authors.

Figure 1
Figure 1. Diagnostics for the allocation surrogate. (a) [PITH_FULL_IMAGE:figures/full_fig_p008_1.png] view at source ↗
Figure 2
Figure 2. FFN work across relative depth. Three work statistics (rows) and three geometries (columns) are shown for the six models with complete 3 ×3 profiles; each curve is normalized to unit mean. Raw Gromov and topological work are directly scale-sensitive and rise with residual scale. Raw shift also rises empirically in these checkpoints, but is not algebraically rescaled by residual norm. Spherical and hyperbolic work ar… view at source ↗
Figure 3
Figure 3. Allocation rules compared with uniform width. [PITH_FULL_IMAGE:figures/full_fig_p013_3.png] view at source ↗

Discussion (0). Sign in to comment.

Reference graph

Works this paper leans on

27 extracted references · 3 linked inside Pith

  1. [1]

    Transformer Circuits Thread , year=

    A Mathematical Framework for Transformer Circuits , author=. Transformer Circuits Thread , year=

  2. [2]

    ICLR , year=

    Residual Connections Encourage Iterative Inference , author=. ICLR , year=

  3. [3]

    ICML , pages=

    On Layer Normalization in the Transformer Architecture , author=. ICML , pages=

  4. [4]

    Error bounds for approximations with deep

    Yarotsky, Dmitry , journal=. Error bounds for approximations with deep

  5. [5]

    Efficient Approximation of Deep

    Chen, Minshuo and Jiang, Haoming and Liao, Wenjing and Zhao, Tuo , booktitle=. Efficient Approximation of Deep

  6. [6]

    JMLR , volume=

    Adaptive Approximation and Generalization of Deep Neural Network with Intrinsic Dimensionality , author=. JMLR , volume=

  7. [7]

    NeurIPS , year=

    Intrinsic dimension of data representations in deep neural networks , author=. NeurIPS , year=

  8. [8]

    Scientific Reports , volume=

    Estimating the intrinsic dimension of datasets by a minimal neighborhood information , author=. Scientific Reports , volume=

Show all 27 references
  1. [9]

    Foundations of Computational Mathematics , volume=

    Gromov--Wasserstein Distances and the Metric Approach to Object Matching , author=. Foundations of Computational Mathematics , volume=

  2. [10]

    Discrete & Computational Geometry , volume=

    Stability of Persistence Diagrams , author=. Discrete & Computational Geometry , volume=

  3. [11]

    A Distribution-Free Theory of Nonparametric Regression , author=

  4. [12]

    Introduction to Nonparametric Estimation , author=

  5. [13]

    Computational Topology: An Introduction , author=

  6. [14]

    Foundations and Trends in Machine Learning , volume=

    Computational Optimal Transport , author=. Foundations and Trends in Machine Learning , volume=

  7. [15]

    NeurIPS , year=

    Poincar\'e Embeddings for Learning Hierarchical Representations , author=. NeurIPS , year=

  8. [16]

    Tapered Language Models , author=

  9. [17]

    2016 , eprint=

    Layer Normalization , author=. 2016 , eprint=

    2016

  10. [18]

    2019 , eprint=

    Root Mean Square Layer Normalization , author=. 2019 , eprint=

    2019

  11. [19]

    2017 , eprint=

    Attention Is All You Need , author=. 2017 , eprint=

    2017

  12. [20]

    2001 , series =

    Burago, Dmitri and Burago, Yuri and Ivanov, Sergei , title =. 2001 , series =

    2001

  13. [21]

    Geometric and Topological Inference , author=

  14. [22]

    arXiv:2407.21783 , year=

    The Llama 3 Herd of Models , author=. arXiv:2407.21783 , year=

    arXiv

  15. [23]

    arXiv:2408.00118 , year=

    Gemma 2: Improving Open Language Models at a Practical Size , author=. arXiv:2408.00118 , year=

    arXiv

  16. [24]

    arXiv:2503.19786 , year=

    Gemma 3 Technical Report , author=. arXiv:2503.19786 , year=

    arXiv

  17. [25]

    and others , journal=

    Jiang, Albert Q. and others , journal=. Mistral 7

  18. [26]

    Gokaslan, Aaron and Cohen, Vanya , year=

  19. [27]

    OpenAI blog , year=

    Language Models are Unsupervised Multitask Learners , author=. OpenAI blog , year=

Pith tools

Reviewed August 4, 2026 · model on record in the stance chip above.