Paper deep dive
Geometry-Guided Layerwise FFN Width Allocation in Transformers
Timur Mudarisov, Mikhail Burtsev, Radu State
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: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.
Tags
Links
- Source: https://arxiv.org/abs/2608.02064v1
- Canonical: https://arxiv.org/abs/2608.02064v1
Trouble viewing inline? Open PDF directly →
Full Text
58,769 characters extracted from source content.
Expand or collapse full text
Geometry-Guided Layerwise FFN Width Allocation in Transformers Timur Mudarisov1 Mikhail Burtsev2 Radu State1 1University of Luxembourg 2London Institute of Mathematical Sciences 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. 1 Introduction Feed-forward networks are the principal pointwise computation inside a Transformer block and typically dominate its parameter count (Vaswani et al., 2017). Nevertheless, most decoder-only language models use one global intermediate dimension: every layer receives the same dffd_f. This design is simple, but it assumes that the amount of useful nonlinear capacity is uniform even though representations and layer functions change with depth (Elhage and others, 2021; Jastrzębski et al., 2018). Uniform width is therefore an architectural convention rather than an evident optimum. Tapered Language Models (TLMs) provide direct evidence that this convention can be improved (Bayat et al., 2026). Their 440M Transformer uses the same total parameter and FLOP budget while reallocating MLP width from later to earlier layers. A smooth cosine taper reduces their in-distribution validation perplexity from 16.2816.28 to 14.4414.44, and the selected schedule is transferred to larger scales and other token-mixing architectures. This result establishes that depth-aware allocation matters. It does not, however, determine the width of an individual layer from the computation performed at that layer: the cosine shape and its endpoints remain hand-designed hyperparameters. A prescribed taper and a measured allocation answer different questions. A taper asks which low-dimensional family of depth profiles should be searched; once the family is selected, every model of a given depth receives essentially the same shape. A measured rule asks which layers of this particular model and data distribution exhibit the strongest evidence of approximation demand. It can reproduce a monotone taper when the evidence is monotone, but it can also retain local peaks or plateaus that a cosine necessarily smooths away. This distinction matters for transfer: a hand-designed schedule is cheap and does not require a reference checkpoint, whereas an activation-derived schedule can adapt to architecture, scale, and pretraining distribution. Our experiments therefore use cosine as a strong structural baseline. We ask whether activations of a reference model can provide a data-dependent schedule. At FFN layer ℓ , the residual states of a sequence form a point cloud ZℓZ_ , and the residual update maps it to a cloud HℓH_ . We call the change between the two clouds the layer’s geometric work. The definition must be aligned with the FFN readout. LayerNorm and RMSNorm largely remove positive per-token rescaling before the branch is evaluated (Ba et al., 2016; Zhang and Sennrich, 2019; Xiong and others, 2020). Raw pairwise-distance and persistence statistics are therefore directly scale-sensitive. By contrast, the raw correspondence-preserving shift equals the magnitude of the FFN branch output and is not algebraically multiplied by the residual norm, although the two quantities may covary empirically across depth. We compare raw Euclidean geometry with spherical and hyperbolic metrics built from unit-normalized states, and distinguish pointwise motion from changes in relational and topological structure. To turn a work profile into widths, we introduce a layerwise approximation surrogate. A latent coefficient CℓC_ represents how difficult the FFN map is to approximate on the representation distribution, λℓ _ represents the loss consequence of an approximation error, and dint(ℓ)d_int( ) controls the error-width exponent. A constrained optimizer then assigns the fixed total FFN budget. Geometric work W^ℓ W_ is used as a proxy for CℓC_ . This proxy is a modeling hypothesis rather than a consequence of topology or optimal transport: the present experiments diagnose its ingredients and test the resulting schedules end to end, but do not directly fit layerwise error-width curves. The paper makes four contributions. First, it defines three complementary forward-pass work statistics under three geometries, separating raw magnitude from scale-normalized reshaping. Second, it derives the unique continuous budget-optimal width allocation, including lower bounds and the depth-dependent-exponent case in which a single proportionality constant is not valid; hardware rounding is treated separately and preserves the budget but is not claimed to solve the discrete optimization problem. Third, it evaluates scale sensitivity, local loss additivity, intrinsic dimension, sensitivity association, and work profiles across seven pretrained decoders. Fourth, it compares geometry-derived schedules with uniform width, a TLM-style cosine taper, and an anti-work control at fixed parameter count. The strongest experiment is the five-seed 440M comparison, where topological/hyperbolic allocation has a substantially larger mean within-protocol loss reduction than cosine. 2 Theory 2.1 FFN Sublayers as Token Transport A decoder block processes N token states in ℝdR^d. Let Zℓ=zℓ(i)i=1NZ_ =\z_ ^(i)\_i=1^N denote the post-attention residual cloud at block ℓ . The FFN update is hℓ(i)=zℓ(i)+FFNℓ(Norm(zℓ(i))),Hℓ=hℓ(i)i=1N.h_ ^(i)=z_ ^(i)+FFN_ \! (Norm(z_ ^(i)) ), H_ =\h_ ^(i)\_i=1^N. (1) Thus the network can be viewed as a sequence of residual transports (Jastrzębski et al., 2018). For an embedding π:ℝd→(ℳ,ρ)π:R^d→(M,ρ) and a cloud comparison D, define W^ℓ=D(π(Zℓ),π(Hℓ)). W_ =D\! (π(Z_ ),π(H_ ) ). (2) For positive c, LayerNorm and RMSNorm satisfy Norm(cz)≃Norm(z)Norm(cz) (z) up to the numerical stabilizer. This invariance applies to the branch input rather than to the full residual state, but it motivates a readout-aligned criterion: a work statistic should not be dominated by variations that the FFN input normalization largely removes. Raw work is therefore retained as a control, while normalized metrics are the primary candidates for allocation. 2.2 Geometries and Work Statistics We evaluate the same input/output clouds in three metric spaces. For a nonzero state define u(z)=z/∥z∥u(z)=z/ z . Raw geometry uses the identity map and Euclidean distance. Spherical geometry uses u(z)u(z) and angular distance ρsph(z,z′)=arccos(⟨u(z),u(z′)⟩). _sph(z,z )= \! ( u(z),u(z ) ). (3) This unit normalization is a scale-invariant proxy inspired by branch-input normalization but it is not identical to full LayerNorm, which recenters the vector and may apply learned coordinatewise gains. Hyperbolic geometry uses the unit Poincaré ball d=x∈ℝd:∥x∥<1B^d=\x ^d: x <1\ with curvature −1-1 and ball radius R=1R=1. Because u(z)u(z) lies on the boundary, we clip it one percent inward: πhyp(z)=ρcu(z),ρc=1−0.01=0.99. _hyp(z)= _cu(z), _c=1-0.01=0.99. (4) The geodesic distance is d(x,y)=arcosh(1+2∥x−y∥2(1−∥x∥2)(1−∥y∥2)).d_B(x,y)=arcosh\! (1+ 2 x-y ^2(1- x ^2)(1- y ^2) ). (5) All embedded points therefore have the same radial coordinate 0.990.99. If α is their spherical angle, then d(ρcu,ρcv)=arcosh(1+8ρc2sin2(α/2)(1−ρc2)2).d_B( _cu, _cv)=arcosh\! (1+ 8 _c^2 ^2(α/2)(1- _c^2)^2 ). (6) Thus this hyperbolic construction is a fixed nonlinear monotone transform of angular separation, and it does not introduce a learned radial hierarchy. We use it as an empirical reweighting of angular distances rather than as evidence that the token cloud is intrinsically hierarchical. For each geometry, we compute three notions of work. The first keeps the known token correspondence and measures mean shift, W^ℓsh=1N∑i=1Nρ(πzℓ(i),πhℓ(i)). W^sh_ = 1N _i=1^Nρ\! (π z_ ^(i),π h_ ^(i) ). (7) This statistic is simple and directly measures update magnitude, but it does not distinguish a nearly rigid transport from a change in the internal shape of the cloud. In raw geometry, Eq. (1) gives W^ℓsh=1N∑i=1N‖FFNℓ(Norm(zℓ(i)))‖. W^sh_ = 1N _i=1^N _ \! (Norm(z_ ^(i)) ) . (8) Consequently, under Norm(cz)≃Norm(z)Norm(cz) (z), positive rescaling of the residual state does not algebraically rescale raw mean shift for a fixed FFN. Any depthwise association between raw shift and residual norm is therefore an empirical co-variation of branch-output magnitude, not a direct metric-scale confound. The second statistic measures metric distortion. With uniform token weights μi=1/N _i=1/N, let Aij=ρ(zi,zj)A_ij=ρ(z_i,z_j) and Bkm=ρ(hk,hm)B_km=ρ(h_k,h_m). We use the entropically regularized Sinkhorn approximation to the quadratic Gromov-Wasserstein objective W^ℓmet W^met_ , GW2,ε2,Sink(A,B):=minT∈Π(μ,μ)∑i,j,k,m|Aij−Bkm|2TikTjm _2, ^2,Sink(A,B)= _T∈ (μ,μ) _i,j,k,m A_ij-B_km ^2T_ikT_jm (9) +εGWKL(T∥μ⊗μ), + _GW\,KL\! (T\, \|\,μ μ ), where εGW>0 _GW>0 is fixed across layers, geometries, and models, and the coupling is optimized by Sinkhorn iterations (Mémoli, 2011; Peyré and Cuturi, 2019). We report the regularized quadratic objective itself: Eq. (9) is a squared-discrepancy convention, not its square root and not a debiased Sinkhorn divergence. The optimized coupling treats the clouds as unlabeled metric-measure spaces and asks whether pairwise relations can be preserved after allowing a coupling, whereas Eq. (7) is the correspondence-preserving baseline. The third statistic focuses on ordinary degree-one persistent homology with coefficients in 2F_2. For a finite metric cloud, the closed Vietoris–Rips filtration includes a simplex at scale t when all its pairwise distances are at most t. We compute the full finite H1H_1 filtration up to the cloud diameter, without filtration clipping or truncation. Ordinary and reduced homology coincide in degree one. Let PD1(Zℓ)PD_1(Z_ ) and PD1(Hℓ)PD_1(H_ ) denote the resulting finite diagrams, with the diagonal available with infinite multiplicity. We define W^ℓtop=2(PD1(Zℓ),PD1(Hℓ)), W^top_ =W_2\! (PD_1(Z_ ),PD_1(H_ ) ), (10) where 2W_2 is the unsquared 22-Wasserstein distance using the ℓ∞ _∞ ground metric in the birth–death plane: 2(P,Q)=[infγ∑p∈P∥p−γ(p)∥∞2]1/2.W_2(P,Q)= [ _γ _p∈ P p-γ(p) _∞^2 ]^1/2. (11) The bijection γ may match points to the diagonal, so an empty diagram is handled by matching every off-diagonal point of the other diagram to its nearest diagonal point. This statistic responds to the creation or destruction of loops, not to component merging, which belongs to H0H_0 (Edelsbrunner and Harer, 2010; Boissonnat et al., 2018). Classical Vietoris–Rips stability is naturally stated using interleaving/Gromov–Hausdorff control and bottleneck distance. Here 2W_2 is used as a finite-sample summary on finite diagrams rather than claimed to satisfy an unrestricted universal stability theorem (Cohen-Steiner et al., 2007). 2.3 Approximation Surrogate Let wℓ=dff(ℓ)w_ =d_f( ) be the width assigned to FFN layer ℓ , let fℓ(z)=FFNℓ(Norm(z))f_ (z)=FFN_ (Norm(z)), and let f~ℓ,w f_ ,w denote a width-w approximation to the trained branch map. The approximation error in the surrogate is the expected squared branch-output error on the calibration distribution PℓP_ , ϵℓ(w)=z∼Pℓ[∥fℓ(z)−f~ℓ,w(z)∥22]. _ (w)=E_z P_ \! [ f_ (z)- f_ ,w(z) _2^2 ]. (12) Thus CℓC_ has the units of squared branch-output error times waℓw^a_ , while λℓ _ converts that local error to a loss effect. We separate the latent approximation coefficient CℓC_ from its observed geometric proxy W^ℓ W_ through three assumptions. Assumption 1 (Local additive loss). Near a trained solution, the loss effect of layerwise approximation errors is locally additive: Δℒ≈∑ℓλℓϵℓ ≈ _ _ _ , where λℓ≥0 _ ≥ 0 is the sensitivity of layer ℓ . Assumption 2 (Intrinsic error-width rate). The residual cloud at layer ℓ concentrates near a set of intrinsic dimension dint(ℓ)≪d_int( ) d, and the FFN approximation error obeys ϵℓ(wℓ)≈Cℓwℓ−aℓ,aℓ=sdint(ℓ)>0, _ (w_ )≈ C_ w_ ^-a_ , a_ = sd_int( )>0, (13) where s is a global expansion-order parameter. All experiments use the fixed value s=3s=3. This form follows the dependence of nonparametric approximation rates on intrinsic rather than ambient dimension (Györfi et al., 2002; Tsybakov, 2009; Yarotsky, 2017; Chen et al., 2019; Nakada and Imaizumi, 2020). It is an ansatz for Transformer FFNs, not a theorem about their learned functions. Assumption 3 (Geometric proxy). On the distribution of residual states used for calibration, the measured work W^ℓ W_ is informative about CℓC_ , so that substituting Cℓ←W^ℓC_ ← W_ preserves enough of the layerwise ordering to construct a useful allocation. A smooth map can alter the observed cloud substantially, and a complex map can approximately preserve it, so Assumption 3 is necessarily model- and distribution-dependent. A direct test would fit error-width curves for individual layers. Without that experiment, correlations with sensitivity and end-to-end training are indirect evidence only; moreover, CℓC_ and λℓ _ are distinct quantities and need not be strongly correlated. 2.4 Budget-Optimal Width Combining Assumptions 1 and 2 gives the lower-bounded continuous surrogate minwℓ≥wmin∑ℓ=1LλℓCℓwℓ−aℓsubject to∑ℓ=1Lwℓ=B. _w_ ≥ w_ _ =1^L _ C_ w_ ^-a_ to _ =1^Lw_ =B. (14) The feasibility condition is B≥LwminB≥ Lw_ . When hidden dimension and FFN parameterization are fixed across layers, the width budget is proportional to both FFN parameters and FFN FLOPs. Lemma 1 (Lower-bounded budget-optimal FFN width). Let cℓ=λℓCℓ≥0c_ = _ C_ ≥ 0, aℓ>0a_ >0, wmin>0w_ >0, and B≥LwminB≥ Lw_ . If B>LwminB>Lw_ , assume that at least one cℓ>0c_ >0. Then Problem (14) has the unique continuous minimizer wℓ⋆=maxwmin,(aℓcℓμ⋆)1/(aℓ+1).w_ = \! \w_ , ( a_ c_ μ )^1/(a_ +1) \. (15) For B>LwminB>Lw_ , μ⋆>0μ >0 is the unique value satisfying ∑ℓwℓ⋆=B _ w_ =B. At B=LwminB=Lw_ , the allocation wℓ⋆=wminw_ =w_ is unique, although the KKT multiplier need not be. If every cℓ>0c_ >0 and the lower bound is inactive, Eq. (15) reduces to wℓ⋆=(aℓcℓμ⋆)1/(aℓ+1).w_ = ( a_ c_ μ )^1/(a_ +1). (16) If additionally aℓ=a_ =a for all layers and θ=1/(a+1)θ=1/(a+1), then wℓ⋆=B(λℓCℓ)θ∑k(λkCk)θ.w_ =B ( _ C_ )^θ _k( _kC_k)^θ. (17) Proof sketch. The feasible set is compact when B≥LwminB≥ Lw_ , and the objective is continuous and convex. If B>LwminB>Lw_ and at least one coefficient is positive, every zero-coefficient layer is optimally placed at wminw_ ; otherwise transferring width from it to a positive-coefficient layer strictly reduces the objective. The objective is strictly convex on the remaining positive-coefficient coordinates, which gives a unique allocation. The KKT condition for an active coordinate is −aℓcℓwℓ−aℓ−1+μ=0-a_ c_ w_ ^-a_ -1+μ=0; an inactive coordinate is clamped at the lower bound, yielding Eq. (15). For B>LwminB>Lw_ , the sum of the right-hand side is continuous and strictly decreasing at the budget level, so it determines a unique μ⋆μ . At B=LwminB=Lw_ the feasible set is a single point. If all cℓ=0c_ =0 and B>LwminB>Lw_ , the objective is constant and the optimizer is not unique; the practical rule then falls back to uniform allocation. ∎ With aℓ=s/dint(ℓ)a_ =s/d_int( ) and s=3s=3, define θℓ=dint(ℓ)/(3+dint(ℓ)) _ =d_int( )/(3+d_int( )). For depth-dependent intrinsic dimension, Eq. (15) must be solved with a shared multiplier; one cannot in general write wℓ∝(λℓCℓ)θℓw_ ( _ C_ ) _ using a single layer-independent proportionality constant because the multiplier is raised to different powers. One-dimensional bisection finds μ⋆μ . Equation (17) also clarifies the role of the exponent. As θ→0θ→ 0, the optimizer approaches uniform width even when work varies strongly; as θ→1θ→ 1, width becomes nearly proportional to the sensitivity-weighted complexity. Intermediate values compress noisy work estimates. Uniform allocation is thus a limiting strongly regularized case, while a hand-designed taper can be viewed as replacing (λℓCℓ)θ( _ C_ )^θ by a prescribed depth prior. The geometry-guided schedule instead estimates that factor from activations and applies the exact continuous budget normalization. The continuous solution is subsequently rounded to hardware-compatible multiples and rebalanced to recover the exact total parameter budget. 2.5 From a Reference Model to a Schedule The practical geometry-guided rule substitutes Cℓ←W^ℓC_ ← W_ . The experiments use the pure-work variant λℓ=1 _ =1 and the fixed expansion-order parameter s=3s=3; perturbation sensitivity is kept independent so that it can serve as a diagnostic rather than a tuning signal. A zero work value is admissible because the continuous program includes wmin>0w_ >0: the corresponding layer is assigned the floor by Eq. (15). If an entire profile is identically zero, the practical fallback is uniform allocation. A schedule is constructed by (1) collecting FFN input/output clouds from a uniform-width reference model, (2) estimating and aggregating the work profile, (3) interpolating that profile to the target depth when needed, (4) solving the lower-bounded continuous fixed-budget program, and (5) rounding and rebalancing the widths. The target model is then trained from scratch. This procedure assumes that the relative layer roles measured in the reference model remain informative after the architecture is reallocated. 3 Experiments 3.1 Experimental Protocols Pretrained-model diagnostics. We analyze seven decoder-only checkpoints: Llama 3.2 at 1B and 3B parameters and Llama 3.1 at 8B (Dubey and others, 2024); Gemma 2 at 2B and 9B (Gemma Team, 2024); Gemma 3 at 1B (Gemma Team, 2025); and Mistral 7B v0.3 (Jiang and others, 2023). A held-out OpenWebText sample (Gokaslan and Cohen, 2019) is passed through each model, and the post-attention FFN input and output clouds are recorded. Raw statistics use the original residual states. Spherical and hyperbolic statistics use the corresponding normalized embeddings. Persistence and Gromov-Wasserstein computations use the same fixed token subsample so that differences are not caused by different examples. Perturbation diagnostics. Layer sensitivity is estimated by low-rank FFN perturbations. For each layer we remove increasing rank fractions and record the validation-loss increase; the local slope of the small-perturbation region defines λℓ _ . To probe Assumption 1, sets of layers are perturbed jointly and their loss increase is compared with the sum of the corresponding single-layer increases. Intrinsic dimension is estimated independently at each relative depth using TwoNN (Facco et al., 2017). For every point xix_i in the layer-ℓ point cloud, let ri,1r_i,1 and ri,2r_i,2 denote its distances to the first and second nearest neighbors, and define μi=ri,2ri,1. _i= r_i,2r_i,1. Under the local-uniformity model of TwoNN, the cumulative distribution satisfies F(μ)=1−μ−d.F(μ)=1-μ^-d. After sorting the ratios, we estimate dint(ℓ)d_int( ) as the slope through the origin in −log(1−F^(μ(i)))≈dint(ℓ)logμ(i).- \! (1- F( _(i)) )≈ d_int( ) _(i). Equivalently, writing xi=logμ(i)x_i= _(i) and yi=−log(1−F^(μ(i)))y_i=- (1- F( _(i))), the fitted slope is dint^(ℓ)=∑i∈ℐxiyi∑i∈ℐxi2, d_int( )= _i x_iy_i _i x_i^2, where ℐI and the empirical-CDF plotting positions are specified below. The resulting TwoNN estimates are used directly, without clipping, winsorization, or smoothing, and the fixed value s=3s=3 gives aℓ=3/dint^(ℓ)a_ =3/ d_int( ). These measurements do not train alternative architectures and therefore provide inexpensive diagnostics of the surrogate. From-scratch allocation runs. We train approximately 128M and 256M Transformers on byte-level-BPE-tokenized OpenWebText (Radford et al., 2019). The comparison includes all nine work/geometry combinations, uniform width, the TLM-style cosine taper, and an anti-topological/raw control. Each size uses three paired seeds: within a seed, all rules share initialization and data order. Every architecture has the same total parameter count, token budget, and AdamW configuration, and final validation loss is the primary endpoint. The 440M experiment allocates the larger compute budget to five paired seeds for a focused set of rules: uniform, cosine, topological/hyperbolic, Gromov/spherical, and anti-topological/raw. Figure 1: Diagnostics for the allocation surrogate. (a) Joint perturbation losses are associated with the sum of the corresponding single-layer losses. (b) Small rank-removal perturbations are approximately linear in the removed rank fraction. (c) TwoNN intrinsic dimension varies smoothly with relative depth. These experiments diagnose the local surrogate; they do not directly calibrate geometric work to an error-width coefficient. The anti-topological/raw control reverses a topological profile while also changing to raw geometry. It tests whether a deliberately misaligned rule can fail in the expected direction, but it is not a fully matched reversal that isolates layer order while holding the estimator and width histogram fixed. Reporting and comparison logic. For each rule and seed define the paired loss difference Ds=Lrule,s−Luniform,sD_s=L_rule,s-L_uniform,s.For the 128M and 256M screens, Fig. 3 reports the mean paired difference D¯ D and descriptive normal-approximation intervals D¯±1.96sD/n D± 1.96\,s_D/ n, with n=3n=3. Because only three paired seeds are available, these intervals are used as exploratory summaries rather than as a basis for formal hypothesis testing. These intervals are descriptive rather than a basis for selecting a definitive winner among eleven alternatives. The 440M experiment narrows the comparison to five schedules and uses five paired seeds. Table 1 reports marginal seed means and standard deviations, together with D¯ D; marginal standard deviations do not determine sDs_D because they omit the within-seed covariance. Therefore, until the five paired differences (or sDs_D, its standard error, and a t4t_4 interval) are supplied from the run logs, the 440M comparisons are stated as differences in mean loss rather than inferential claims of superiority. All qualitative claims distinguish a within-protocol comparison from an absolute comparison with values reported by another paper. 3.2 Diagnostics of the Allocation Model Scale sensitivity. The first experiment asks whether work measures geometry or merely residual scale. Raw Gromov and persistence work grow strongly in late layers and track residual-norm growth, as expected because their pairwise distances and filtration scales change under Euclidean rescaling. Per-token normalization largely removes that late-depth growth. Raw mean shift also tends to rise in the observed profiles, but Eq. (8) shows that this is not a direct metric-scale effect: it is an empirical co-variation between residual norm and FFN-output magnitude. The contrast is visible across work families in Fig. 2. We therefore treat raw work as a control, while keeping the causal interpretation statistic-specific. Local additivity and linearity. The joint perturbation loss is positively associated with the sum of the corresponding single-layer losses, with R2≈0.72R^2≈ 0.72 (Fig. 1, left). The rank-removal curves are approximately linear for small and moderate perturbations and become nonlinear for the most aggressive interventions in several models. This is the regime expected for a first-order surrogate, so λℓ _ is estimated only from the local part of the curve. The correlation does not by itself prove equality to the identity line, but it supports using an additive approximation for schedule derivation. Intrinsic dimension. TwoNN estimates stay far below the ambient hidden dimension and change smoothly with depth (Fig. 1, right), consistent with prior evidence that learned representations occupy low-dimensional sets (Ansuini et al., 2019). Early layers generally have the smallest estimates, while some larger Gemma models reach values around 44–55 in the middle of the stack. The estimates shown and used by the allocator are the unmodified TwoNN outputs: no clipping or smoothing is applied. With the fixed expansion-order parameter s=3s=3, θℓ=dint(ℓ)/(3+dint(ℓ)) _ =d_int( )/(3+d_int( )) remains in an intermediate range and changes smoothly. The exponent therefore compresses large differences in measured work rather than transferring them one-for-one to width. This supports the use of an intrinsic-dimensional exponent, while not establishing the exact rate in Eq. (13). Sensitivity association. After standardization within each model, Gromov work has a stronger association with λℓ _ than topological work in five of seven models; the pooled association of the current topological estimator is weak. The result argues against claiming that topology universally predicts layer importance. It also does not invalidate topological allocation: λℓ _ measures the loss consequence of an approximation error, whereas CℓC_ is intended to measure the difficulty of realizing the layer map. The two factors enter the surrogate multiplicatively and can rank estimators differently. Taken together, the diagnostics support using a normalized, regularized profile and the local additive objective, but they do not establish Cℓ=W^ℓC_ = W_ . This boundary is important for interpreting the later training experiments: a successful schedule is evidence that the proxy carries useful ordering information, not a measurement of the true approximation seminorm or a validation of a universal nonparametric exponent. 3.3 Work Profiles Across Depth Figure 2: FFN work across relative depth. Three work statistics (rows) and three geometries (columns) are shown for the six models with complete 3×33× 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 are predominantly front-loaded and can depart from a fixed monotone taper. Figure 2 compares the depth profile after normalizing each model/statistic curve to unit mean. In raw geometry, topological and Gromov work tend to increase near the end of the network, consistent with direct scale-sensitivity of pairwise distances and persistence filtrations. Raw shift also tends to increase, but this is interpreted only as empirical co-variation of FFN-output magnitude with residual norm. Under spherical and hyperbolic geometry, the ordering mostly reverses: work is high in the earliest layers, decreases through the middle of the stack, and sometimes rises mildly near the end. The qualitative front-loading transfers across model families, but the local deviations differ by estimator and checkpoint. Thus the measurements support the broad direction identified by TLM while supplying a layer-specific, potentially non-monotone profile rather than imposing a cosine shape. Relative depth ℓ/L /L is used because the checkpoints have different numbers of blocks. Unit-mean normalization removes the arbitrary scale of each estimator and makes the experiment about allocation shape; the total width is set later by the budget constraint. The cross-family agreement is strongest for the early peak and weaker for fine structure near the middle and final blocks. That pattern motivates transferring a smoothed profile rather than copying individual noisy layer values. It also explains why geometry and tapering are complementary: both favor earlier capacity, but only the measured profile can preserve repeatable departures from monotonicity. Figure 3: Allocation rules compared with uniform width. Paired change in final validation loss for the 128M and 256M iso-parameter runs, negative is better. Points are means over three paired seeds and bars are 95% confidence intervals. 3.4 Paired 128M and 256M Training Figure 3 evaluates the full 3×33× 3 work design space at two tractable model sizes. Most schedules based on normalized geometry have negative mean deltas relative to their seed-matched uniform baselines. Topological work in hyperbolic geometry has the lowest mean validation loss among the tested configurations; Gromov/spherical and shift/hyperbolic also improve on average. The TLM-style cosine schedule remains close to uniform, while the anti-topological/raw control is centered near or above zero. These runs show that the profile contains information beyond parameter count: every rule uses the same FFN budget, and pairing removes initialization and data order from the within-seed comparison. At the same time, three seeds are too few for a definitive estimator ranking. The 95% intervals overlap zero and one another, so the small-scale result is best interpreted as a broad screen that identifies normalized schedules for the more focused larger experiment. The fact that Gromov is strongest in the sensitivity diagnostic while topological work has the best training mean reinforces that sensitivity association and allocation quality are related but non-equivalent tests. These results covering all combinations of statistics and geometries allow to disentangle two design choices, the first one is whether scale normalization improves allocation, and the second is how much correspondence and relational structure the estimator should preserve. Normalized geometries perform more consistently favorably than raw geometry, whereas no single work statistic is best across all settings. The results therefore provide stronger support for normalization as a general design principle than for any particular topological or transport-based estimator. Because this comparison is exploratory, the larger experiment evaluates a representative set of leading rules rather than choosing a separate winner for each seed or model size. 3.5 440M Comparison with TLM TLM uses a 440M Transformer to select a hand-designed taper and reports that its best cosine schedule improves validation perplexity from 16.2816.28 to 14.4414.44 in its training setup (Bayat et al., 2026). We include the same model scale and treat the cosine schedule as a direct baseline within our own protocol. Unlike an absolute cross-paper comparison, this within-protocol experiment controls the data, tokenizer, optimization, parameter count, and random seeds across allocation rules. Table 1 gives the focused descriptive comparison in this study. The TLM-style cosine taper has mean validation loss 0.0030.003 below uniform. Topological/hyperbolic allocation has mean loss 0.0190.019 below uniform, approximately 6.3×6.3× the mean reduction of cosine, while Gromov/spherical has mean loss 0.0160.016 below uniform. The two geometry-derived schedules therefore have lower mean validation loss than the prescribed cosine baseline in the same five-seed experiment. The anti-topological/raw control moves in the opposite direction and has higher mean loss and perplexity. These statements do not establish statistical separation because the paired-difference variability is not reported in the current table. Table 1: Five-seed 440M iso-parameter results. Validation loss is reported as the marginal mean ± marginal standard deviation over five paired seeds. Perplexity is exp(mean loss) (mean loss), and Δ=D¯ = D is the mean paired loss difference relative to uniform (equivalently, the arithmetic difference of means under complete pairing). The table does not contain the standard deviation or confidence interval of DsD_s, so it supports descriptive mean comparisons rather than an inferential superiority claim. Lower is better. Rule Val. loss ↓ PPL ↓ Δ loss ↓ uniform 3.449±0.0223.449± 0.022 31.4731.47 - TLM cosine 3.446±0.0223.446± 0.022 31.3731.37 −0.003-0.003 topological / hyper 3.430±0.0233.430± 0.023 30.8830.88 −0.019-0.019 Gromov / sphere 3.433±0.0283.433± 0.028 30.9730.97 −0.016-0.016 anti-topological / raw 3.515±0.0983.515± 0.098 33.6233.62 +0.066+0.066 The ratio is a within-protocol descriptive statement, not a claim that the absolute losses are directly comparable with the perplexities reported by TLM: the two studies use different training pipelines. The relevant connection is methodological. TLM shows that front-loading capacity can work; the present matched experiment indicates that an activation-derived profile can have a larger mean reduction than a fixed cosine. Five paired seeds are more informative than the three-seed 128M/256M screen, but sample count alone does not establish stability, and the paired interval is needed for inferential claims. 4 Discussion The experiments support two levels of conclusion. First, FFN capacity should not be assumed uniform. This agrees with TLM and with the observation that layer updates change their role across depth. Second, a fixed monotone taper is not the only useful alternative. Normalized work profiles recover the broad early-heavy shape but preserve model- and layer-specific structure, and the five-seed 440M comparison shows a substantially larger mean loss reduction for the geometry-derived schedules than for cosine under identical conditions. This is a descriptive statement about means until paired-difference uncertainty is reported. The results do not establish one universally superior geometry. Gromov work is more consistently associated with perturbation sensitivity, whereas topological/hyperbolic allocation obtains the best training mean. This split is plausible under the theory: sensitivity λℓ _ and approximation coefficient CℓC_ play different roles, and the training outcome also depends on profile smoothing, rounding, and interactions introduced by retraining. A conservative reading is therefore that normalized geometric change is useful, not that persistent homology is always the correct estimator. The central limitation is the absence of a direct layerwise width sweep. Such an experiment would test whether W^ℓ W_ predicts the coefficient of an error-width curve, rather than relying on sensitivity association and end-to-end outcomes. The geometric statistics are also finite-sample estimates as persistence and optimal transport can have different variance and computational cost. The fixed-radius hyperbolic embedding is only a nonlinear transform of angular distance and does not test a learned radial hierarchy. The schedule is measured on a reference model and transferred to a target model trained from scratch, so profile stability under architectural change is assumed. Finally, the anti-work control changes both direction and geometry, and therefore does not isolate alignment as cleanly as a matched reversed or permuted normalized profile would. The three-seed small-scale intervals remain wide. The five-seed 440M means are more informative, but the current table omits uncertainty of the paired differences and therefore does not justify inferential or fine-grained claims about estimator ordering. 5 Conclusion Geometric work provides a principled way to replace a global FFN-width constant with a fixed-budget, layer-dependent schedule. Raw pairwise and topological statistics are scale-sensitive, whereas raw mean shift measures FFN-output magnitude and only empirically covaries with residual scale. Normalized metric and topological changes reveal a transferable early-heavy profile. Across paired training runs, geometry-derived schedules have lower mean validation loss than uniform width in several comparisons and, at 440M, have a larger mean reduction than the TLM-style cosine baseline within the same five-seed protocol. Paired-difference intervals are required before interpreting these mean gaps as statistical separation. These findings position activation-derived capacity allocation as a practical extension of tapering and motivate direct calibration of geometric work to layerwise approximation demand. Appendix A Experimental Details and How the Figures Are Produced This appendix documents the experiments of the main paper in reproducible detail and specifies exactly how each figure and table is produced: what quantity is on each axis, how it is estimated, what normalization and aggregation are applied, and what the error bars mean. It contains no new claims; it expands the experimental protocols so that every panel can be regenerated from the released code. Section A.1 gives the shared setup (data, models, estimators, budgets, seeds, hardware, and statistics); Section A.2 describes each figure and table panel by panel. A.1 Experimental Details Data and Preprocessing All measurements and all training use OpenWebText (Gokaslan and Cohen, 2019). For the pretrained-model diagnostics we stream the corpus and pass a held-out sample through each network in inference mode (bfloat16, scaled dot-product attention, no gradient computation), recording for every FFN sublayer the post-attention input cloud ZℓZ_ and the post-FFN output cloud HℓH_ . Raw statistics use the original residual states; spherical and hyperbolic statistics use the corresponding normalized embeddings. Persistence and Gromov-Wasserstein statistics for a given layer use the same fixed token subsample for ZℓZ_ and HℓH_ , so that a measured difference reflects the action of the sublayer rather than a change of examples. For the from-scratch runs, OpenWebText is tokenized with a byte-level BPE tokenizer (Radford et al., 2019) and packed into a flat token stream; training and validation shards are disjoint at the document level. Pretrained Models Seven decoder-only checkpoints spanning three families and a range of sizes are used: Llama-3.2-1B, Llama-3.2-3B, and Llama-3.1-8B (Dubey and others, 2024); Gemma-2-2B and Gemma-2-9B (Gemma Team, 2024); Gemma-3-1B (Gemma Team, 2025); and Mistral-7B-v0.3 (Jiang and others, 2023). Depth profiles are plotted against relative depth ℓ/L /L because the checkpoints have different numbers of blocks. Six of the seven have complete 3×33× 3 work profiles and are the ones shown in the profile grid. Work Statistics and Geometries For each geometry we compute three notions of work on the matched clouds. With u(z)=z/∥z∥u(z)=z/ z , raw geometry uses the identity map and Euclidean distance; spherical geometry uses u(z)u(z) and the angular distance arccos⟨u(z),u(z′)⟩ u(z),u(z ) ; hyperbolic geometry uses a fixed monotone radial transform into the unit Poincaré ball with its geodesic distance. The three statistics are the topological work W^ℓtop=W2(PD1(Zℓ),PD1(Hℓ)) W^top_ =W_2(PD_1(Z_ ),PD_1(H_ )), where W2W_2 is the unsquared 2-Wasserstein distance between degree-one persistence diagrams under the ℓ∞ _∞ birth–death ground metric (diagonal matching handles empty diagrams); the metric work W^ℓmet W^met_ , a fixed-regularization Gromov-Wasserstein discrepancy between the intra-cloud distance matrices of ZℓZ_ and HℓH_ ; and the shift work W^ℓsh W^sh_ , the mean transported distance between matched input and output points under the chosen geometry. Persistence is computed from the Vietoris–Rips filtration of the distance matrix at degree one. Every curve in the profile figure is normalized to unit mean over depth, which removes the arbitrary scale of each estimator and makes the comparison about allocation shape; the total width is set separately by the budget constraint. Perturbation Diagnostics Layer sensitivity λℓ _ is estimated by low-rank FFN perturbations. For a layer we remove increasing rank fractions from the FFN branch and record the resulting increase in validation loss; λℓ _ is the slope of the small-perturbation (locally linear) part of that curve. To probe the local additive-loss assumption, sets of layers are perturbed jointly and the joint loss increase is compared against the sum of the corresponding single-layer increases. These diagnostics do not train alternative FFN widths and therefore do not calibrate geometric work to an error–width coefficient; they test the local surrogate only. Intrinsic Dimension Intrinsic dimension is estimated independently at each relative depth with TwoNN (Facco et al., 2017). For each point xix_i in the layer-ℓ cloud, with ri,1r_i,1 and ri,2r_i,2 its first two nearest-neighbour distances, let μi=ri,2/ri,1 _i=r_i,2/r_i,1. Under the TwoNN local-uniformity model F(μ)=1−μ−dF(μ)=1-μ^-d, dint(ℓ)d_int( ) is the origin-slope of −log(1−F^(μ(i)))- (1- F( _(i))) against logμ(i) _(i) on the sorted ratios. The TwoNN outputs are used directly: no clipping, winsorization, or smoothing. With the fixed global expansion order s=3s=3, the exponent is θℓ=dint(ℓ)/(3+dint(ℓ)) _ =d_int( )/(3+d_int( )) and aℓ=3/dint(ℓ)a_ =3/d_int( ). From-Scratch Allocation Runs We train approximately 128M and 256M Transformers on tokenized OpenWebText. The 128M/256M design space is the full set of nine work/geometry combinations, plus uniform width, the TLM-style cosine taper, and an anti-topological/raw control (eleven rules in total). Each size uses three paired seeds: within a seed, all rules share initialization and data order, so a within-seed comparison removes those nuisance factors. Every rule at a given size has the same total parameter count, the same token budget, and the same AdamW configuration, so the comparison is iso-parameter and iso-token; the final validation loss is the primary endpoint. The 440M experiment spends a larger compute budget on five paired seeds for a focused set of five rules: uniform, cosine, topological/hyperbolic, Gromov/spherical, and anti-topological/raw. The anti-topological/raw control reverses a topological profile and simultaneously switches to raw geometry; it tests whether a deliberately misaligned rule fails in the expected direction, but it is not a fully matched reversal that isolates layer order while holding the estimator and width histogram fixed. Width Schedules from a Reference Profile A measured, unit-mean work profile is smoothed across relative depth and mapped to a width schedule by the budget-optimal rule of the main paper, so that ∑ℓwℓ _ w_ equals the uniform budget. Widths are quantized to the architecture’s width granularity and floored at a minimum width; the quantization residual is redistributed so that the total width matches the uniform baseline exactly, which is what makes every rule iso-parameter by construction rather than approximately. Profiles are transferred by relative depth ℓ/L /L because the reference and target models differ in depth, and a smoothed profile is transferred rather than individual noisy layer values. Computing Infrastructure All experiments run on a single NVIDIA RTX A6000 (GA102; 48 GB GDDR6 with ECC; approximately 768 GB/s memory bandwidth; 300 W board power; PCIe 3.0 ×16× 16 in this host); no multi-GPU, model-parallel, or sharded training is used. At the training batch used, the largest configuration occupies roughly 26 GiB of the 48 GiB available, so no activation checkpointing or offloading is required. On a CUDA out-of-memory event the trainer halves the microbatch and raises gradient accumulation to keep the tokens per optimizer step unchanged, turning memory pressure into reduced throughput rather than a failed run. Between runs, GPU memory is released and downloaded model and dataset artifacts are purged from the local cache; runs are checkpointed and resumable. Randomness and Seeds Python, NumPy, and PyTorch seeds (including all CUDA devices) are set at the start of each run. Within a seed, initialization and data order are shared across all allocation rules, so each rule-versus-uniform comparison is exactly paired and isolates the width profile. The 128M and 256M screens use three paired seeds each; the 440M experiment uses five paired seeds. The pretrained-model diagnostics draw a single fixed token subsample per layer. Evaluation Metrics and Reporting The endpoint is final validation loss (cross-entropy in nats); perplexity is exp(mean loss) (mean loss). For a rule and seed we define the paired difference Ds=Lrule,s−Luniform,sD_s=L_rule,s-L_uniform,s; D¯ D is its mean over seeds. For the 128M/256M screens we report D¯ D with descriptive normal-approximation intervals D¯±1.96sD/n D± 1.96s_D/ n at n=3n=3; because only three paired seeds are available, these intervals are exploratory summaries, not a basis for formal testing or for selecting a single winner among eleven alternatives. For the 440M experiment the table reports marginal seed means and standard deviations together with D¯ D; marginal standard deviations do not determine the standard deviation of DsD_s because they omit the within-seed covariance, so the 440M comparisons are stated as differences in mean loss rather than as inferential superiority claims until the paired-difference variability is supplied from the run logs. All qualitative claims distinguish a within-protocol comparison from an absolute comparison against numbers reported by another paper. A.2 How Each Figure and Table Is Produced Diagnostics of the Allocation Surrogate The three-panel diagnostic figure (Figure 1; source identifier fig_assumptions_abc) visualizes the three assumptions behind the width rule; each panel is a scatter or profile over the pretrained models of Section A.1. Panel (a): local additivity. Each point is one joint perturbation of a set of layers. Its x-value is the sum of the single-layer validation-loss increases of the layers in that set (each measured separately as above); its y-value is the validation-loss increase when those layers are perturbed jointly. The identity line y=xy=x is exact additivity. The reported association is R2≈0.72R^2≈ 0.72; the panel shows a positive association with the identity line rather than exact equality, which is the evidence used to justify an additive objective for schedule derivation. Because a point is a set of layers, the panel contains many more points than there are layers, and the perturbation strength is held in the locally linear regime (panel (b)) so that the scatter reflects cross-layer interaction rather than within-layer nonlinearity. Panel (b): local linearity of rank removal. For a layer, the curve plots the validation-loss increase (vertical axis) against the removed rank fraction of the FFN branch (horizontal axis). Curves are shown per layer and are approximately linear for small and moderate removal, becoming nonlinear for the most aggressive interventions in several models. λℓ _ is read as the slope of the linear part only; the panel’s purpose is to show that this linear region exists and to delimit the perturbation strength used in panel (a). Panel (c): intrinsic dimension versus depth. For each model, the curve plots the TwoNN estimate dint(ℓ)d_int( ) (vertical axis) against relative depth ℓ/L /L (horizontal axis). Values are the unmodified TwoNN outputs. The panel shows that dintd_int stays far below the ambient hidden dimension, changes smoothly with depth, is smallest in early layers, and reaches roughly 4–5 in the middle of the larger Gemma stacks. This motivates the intrinsic-dimensional exponent θℓ=dint(ℓ)/(3+dint(ℓ)) _ =d_int( )/(3+d_int( )), which compresses large differences in measured work rather than transferring them one-for-one to width. Optional companion panels. Where the additivity and dimension panels are shown separately (source identifiers fig_e_additivity and fig_e_dim), the axes and estimators are identical to panels (a) and (c) above. FFN Work Across Relative Depth The profile grid (Figure 2; source identifier fig_ffn_grid_clean) is a 3×33× 3 array: rows are the three work statistics (W^top W^top, W^met W^met, and W^sh W^sh) and columns are the three geometries (raw, spherical, and hyperbolic). Within each cell, one curve per model plots that statistic against relative depth ℓ/L /L; only the six models with complete 3×33× 3 profiles are drawn. Every curve is normalized to unit mean over depth, so the panels compare shape, not magnitude. The figure is read for two contrasts. First, raw Gromov and raw topological work rise toward the end of the network, tracking residual-norm growth, because their pairwise distances and filtration scales change under Euclidean rescaling; raw shift also rises, but as an empirical co-variation of FFN-output magnitude with residual norm rather than an algebraic rescaling. Second, under spherical and hyperbolic geometry the ordering largely reverses: work is front-loaded (high in the earliest layers, decreasing through the middle, occasionally rising mildly near the end) and can depart from a fixed monotone taper. The cross-family agreement is strongest for the early peak and weakest for fine structure near the middle and final blocks, which is why a smoothed profile is transferred rather than individual noisy layer values. Allocation Rules Versus Uniform (128M and 256M) The forest plot (Figure 3; source identifier figure6_allocation_vs_uniform_forest) evaluates the full eleven-rule design space at the two smaller sizes. Each row is one allocation rule. The plotted point is the mean paired change in final validation loss relative to the seed-matched uniform baseline, D¯=means(Lrule,s−Luniform,s) D=mean_s(L_rule,s-L_uniform,s), over three paired seeds; more negative is better. The bars are the descriptive 95% normal-approximation intervals D¯±1.96sD/3 D± 1.96s_D/ 3 described above. The 128M and 256M sizes are shown as grouped rows so that the same rule can be compared across the two scales. Because the intervals overlap zero and one another, the panel is read as an exploratory screen that identifies normalized-geometry schedules as consistently favorable and the anti-topological/raw control as centered near or above zero, not as a definitive ranking of estimators. Five-Seed 440M Comparison with TLM Table 1 reports, for five rules (uniform, cosine, topological/hyperbolic, Gromov/spherical, and anti-topological/raw), the marginal validation-loss mean ± marginal standard deviation over five paired seeds, the perplexity exp(mean loss) (mean loss), and Δ=D¯ = D, the mean paired loss difference relative to uniform. Under complete pairing Δ equals the arithmetic difference of the two reported means. The table deliberately omits the standard deviation and confidence interval of the paired difference DsD_s, because the marginal standard deviation does not determine it (the within-seed covariance is not reported); accordingly the 440M results support descriptive mean comparisons and are not stated as inferential superiority claims. The cosine schedule is treated as a direct in-protocol baseline for the same 440M scale used by TLM, with data, tokenizer, optimization, parameter count, and seeds controlled across rules; the comparison is within-protocol and is not an absolute comparison against the perplexities reported by TLM, whose training pipeline differs. References A. Ansuini, A. Laio, J. H. Macke, and D. Zoccolan (2019) Intrinsic dimension of data representations in deep neural networks. In NeurIPS, Cited by: §3.2. J. L. Ba, J. R. Kiros, and G. E. Hinton (2016) Layer normalization. External Links: 1607.06450, Link Cited by: §1. R. Bayat, A. Behrouz, and A. Courville (2026) Tapered language models. Note: Preprint Cited by: §1, §3.5. J. Boissonnat, F. Chazal, and M. Yvinec (2018) Geometric and topological inference. Cambridge Texts in Applied Mathematics, Cambridge University Press. Cited by: §2.2. M. Chen, H. Jiang, W. Liao, and T. Zhao (2019) Efficient approximation of deep ReLU networks for functions on low dimensional manifolds. In NeurIPS, Cited by: §2.3. D. Cohen-Steiner, H. Edelsbrunner, and J. Harer (2007) Stability of persistence diagrams. Discrete & Computational Geometry 37 (1), p. 103–120. Cited by: §2.2. A. Dubey et al. (2024) The llama 3 herd of models. arXiv:2407.21783. Cited by: §A.1, §3.1. H. Edelsbrunner and J. Harer (2010) Computational topology: an introduction. American Mathematical Society. Cited by: §2.2. N. Elhage et al. (2021) A mathematical framework for transformer circuits. Transformer Circuits Thread. Cited by: §1. E. Facco, M. d’Errico, A. Rodriguez, and A. Laio (2017) Estimating the intrinsic dimension of datasets by a minimal neighborhood information. Scientific Reports 7 (1), p. 12140. Cited by: §A.1, §3.1. Gemma Team (2024) Gemma 2: improving open language models at a practical size. arXiv:2408.00118. Cited by: §A.1, §3.1. Gemma Team (2025) Gemma 3 technical report. arXiv:2503.19786. Cited by: §A.1, §3.1. A. Gokaslan and V. Cohen (2019) OpenWebText corpus. Note: http://Skylion007.github.io/OpenWebTextCorpus Cited by: §A.1, §3.1. L. Györfi, M. Kohler, A. Krzyżak, and H. Walk (2002) A distribution-free theory of nonparametric regression. Springer. Cited by: §2.3. S. Jastrzębski, D. Arpit, N. Ballas, V. Verma, T. Che, and Y. Bengio (2018) Residual connections encourage iterative inference. In ICLR, Cited by: §1, §2.1. A. Q. Jiang et al. (2023) Mistral 7B. arXiv:2310.06825. Cited by: §A.1, §3.1. F. Mémoli (2011) Gromov–wasserstein distances and the metric approach to object matching. Foundations of Computational Mathematics 11 (4), p. 417–487. Cited by: §2.2. R. Nakada and M. Imaizumi (2020) Adaptive approximation and generalization of deep neural network with intrinsic dimensionality. JMLR 21 (174), p. 1–38. Cited by: §2.3. G. Peyré and M. Cuturi (2019) Computational optimal transport. Foundations and Trends in Machine Learning 11 (5–6), p. 355–607. Cited by: §2.2. A. Radford, J. Wu, R. Child, D. Luan, D. Amodei, and I. Sutskever (2019) Language models are unsupervised multitask learners. OpenAI blog. Cited by: §A.1, §3.1. A. B. Tsybakov (2009) Introduction to nonparametric estimation. Springer. Cited by: §2.3. A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. Kaiser, and I. Polosukhin (2017) Attention is all you need. External Links: 1706.03762, Link Cited by: §1. R. Xiong et al. (2020) On layer normalization in the transformer architecture. In ICML, p. 10524–10533. Cited by: §1. D. Yarotsky (2017) Error bounds for approximations with deep ReLU networks. Neural Networks 94, p. 103–114. Cited by: §2.3. B. Zhang and R. Sennrich (2019) Root mean square layer normalization. External Links: 1910.07467, Link Cited by: §1.