Latent-Space Mean-Field Theory for Deep BitNet-like Training: Constrained Gradient Flows with Smooth Quantization and STE Limits

TL;DR

This paper introduces a mean-field theoretical framework for understanding the training dynamics of deep quantized neural networks, revealing how smooth quantization influences convergence and stability.

Contribution

It develops a novel mean-field analysis for deep quantized networks, showing convergence of latent weight measures and establishing a foundation for gradient-based training with quantization.

Findings

Empirical measures converge to solutions of constrained continuity equations.

Exponential decay cancels singularities, ensuring uniform bounds.

Provides a rigorous mathematical basis for training quantized neural networks.

Abstract

This work develops a mean-field analysis for the asymptotic behavior of deep BitNet-like architectures as smooth quantization parameters approach zero. We establish that empirical measures of latent weights converge weakly to solutions of constrained continuity equations under vanishing quantization smoothing. Our main theoretical contribution demonstrates that the natural exponential decay in smooth quantization cancels out apparent singularities, yielding uniform bounds on mean-field dynamics independent of smoothing parameters. Under standard regularity assumptions, we prove convergence to a well-defined limit that provides the mathematical foundation for gradient-based training of quantized neural networks through distributional analysis.

Equations

$${\alpha }^{(\ell )}(W^{(\ell )})\equiv {\Psi }^{(\ell )}(W^{(\ell )}):=\frac{1}{n_{\ell }{m}_{\ell }}\sum _{i=1}^{n_{\ell }}\sum _{j=1}^{m_{\ell }}{W}_{ij}^{(\ell )},$$

$$P^{(\ell )}(W^{(\ell )}):=W^{(\ell )}-{\alpha }^{(\ell )}(W^{(\ell )})1_{n_{\ell }\times m_{\ell }},$$

$${\alpha }^{(\ell )}(W)1+P^{(\ell )}(W)$$

$$⟨1,P^{(\ell )}(W)⟩$$

$$=⟨1,W⟩-{\alpha }^{(\ell )}(W)⟨1,1⟩$$

$$=n_{\ell }{m}_{\ell }{\alpha }^{(\ell )}(W)-{\alpha }^{(\ell )}(W)(n_{\ell }{m}_{\ell })=0,$$

$$⟨A,P^{(\ell )}(B)⟩$$

$$=⟨A,B⟩-{\alpha }^{(\ell )}(B)⟨A,1⟩$$

$$=⟨A,B⟩-{\alpha }^{(\ell )}(B)\cdot n_{\ell }{m}_{\ell }{\alpha }^{(\ell )}(A)$$

$$=⟨A,B⟩-{\alpha }^{(\ell )}(A)\cdot n_{\ell }{m}_{\ell }{\alpha }^{(\ell )}(B)$$

$$=⟨A-{\alpha }^{(\ell )}(A)1,B⟩=⟨P^{(\ell )}(A),B⟩.$$

$${\alpha }^{(\ell )}(P^{(\ell )}(W))$$

$$={\alpha }^{(\ell )}(W)-{\alpha }^{(\ell )}(W)\cdot {\alpha }^{(\ell )}(1)$$

$$={\alpha }^{(\ell )}(W)-{\alpha }^{(\ell )}(W)\cdot 1=0.$$

$$\Vert W{\Vert }_F^2$$

$$=\Vert {\alpha }^{(\ell )}(W)1{\Vert }_F^2+\Vert P^{(\ell )}(W){\Vert }_F^2+2⟨{\alpha }^{(\ell )}(W)1,P^{(\ell )}(W)⟩$$

$$=|{\alpha }^{(\ell )}(W){|}^2\Vert 1{\Vert }_F^2+\Vert P^{(\ell )}(W){\Vert }_F^2+0$$

$$=n_{\ell }{m}_{\ell }|{\alpha }^{(\ell )}(W){|}^2+\Vert P^{(\ell )}(W){\Vert }_F^2.$$

$$\Vert T_c(A)-T_c(B){\Vert }_F$$

$${\mathrm{sgn}}_{\epsilon }(z)$$

$$|\cdot {|}_{\epsilon }(z)$$

$${\mathrm{clip}}_{\epsilon }(x;a,b)$$

$$\widetilde{W}^{(\ell)}_{\varepsilon}:=\mathrm{sgn}_{\varepsilon}\big{(}P^{(\ell)}(W^{(\ell)})\big{)}\in[-1,1]^{n_{\ell}\times m_{\ell}}.$$

$${\beta }_{\epsilon }^{(\ell )}(W^{(\ell )}):=\frac{1}{n_{\ell }{m}_{\ell }}\sum _{i,j}{\left|P^{(\ell )}(W^{(\ell )}{)}_{ij}\right|}_{\epsilon },$$

$${\mathrm{Quant}}_{\epsilon }^{(b)}(x):={\mathrm{clip}}_{\epsilon }\left(\frac{x}{{\gamma }_{\epsilon }(x)}\cdot Q_b,-Q_b+\delta ,Q_b-\delta \right),{\gamma }_{\epsilon }(x):=\max \{\epsilon ,\Vert x{\Vert }_{\infty }\},$$

$$h^{(\ell)}(x)=\sigma^{(\ell)}\!\Big{(}\beta^{(\ell)}_{\varepsilon}(W^{(\ell)})\,\widetilde{W}^{(\ell)}_{\varepsilon}\,\mathrm{Quant}^{(b)}_{\varepsilon}(x)\Big{)}.$$

$$h^{(0)}(x)$$

$$h^{(\ell )}(x)$$

$$f_W(x)$$

$$\mathcal{R}_{\varepsilon}(W^{(1)},\dots,W^{(L)}):=\mathbb{E}_{(X,Y)\sim\pi}\big{[}\ell\big{(}f_{W}(X),Y\big{)}\big{]}.$$

$$W^{(\ell)}(k+1)=W^{(\ell)}(k)-\eta\,\nabla_{W^{(\ell)}}\mathcal{R}_{\varepsilon}\big{(}W^{(1)}(k),\dots,W^{(L)}(k)\big{)},$$

$$\Vert f_W(x)-f_{\hat{W}}(x){\Vert }_2\le L_{\mathrm{fwd}}\sum _{\ell =1}^L\Vert W^{(\ell )}-{\hat{W}}^{(\ell )}{\Vert }_F.$$

$$|{\mathrm{sgn}}_{\epsilon }^{\prime }(z)|={\epsilon }^{-1}{\mathrm{sech}}^2(z/\epsilon )\le {\epsilon }^{-1}.$$

$$|\nabla |\cdot {|}_{\epsilon }(z)|=\left|\frac{z}{\sqrt{z^2+{\epsilon }^2}}\right|\le 1.$$

$$|{\gamma }_{\epsilon }(x)-{\gamma }_{\epsilon }(y)|\le \Vert x-y{\Vert }_{\infty }.$$

$$|{\beta }_{\epsilon }^{(\ell )}(W)-{\beta }_{\epsilon }^{(\ell )}(\hat{W})|$$

$$\le \frac{1}{n_{\ell }{m}_{\ell }}\sum _{i,j}|P^{(\ell )}(W-\hat{W}{)}_{ij}|$$

$$=\frac{1}{n_{\ell }{m}_{\ell }}\Vert P^{(\ell )}(W-\hat{W}){\Vert }_1$$

$$\le \frac{1}{n_{\ell }{m}_{\ell }}\sqrt{n_{\ell }{m}_{\ell }}\Vert P^{(\ell )}(W-\hat{W}){\Vert }_F$$

Full text

1 Introduction

The training dynamics of quantized neural networks pose fundamental theoretical challenges due to the non-differentiable nature of quantization operators. BitNet-like architectures [12] employ discrete sign and clipping functions in forward propagation while maintaining continuous latent weights for gradient-based optimization. This creates a mathematical tension between the discrete forward pass and smooth optimization requirements.

Mean-field theory [8, 3] offers a powerful framework for analyzing neural network training by treating parameters as interacting particles and studying their empirical measure evolution. However, extending this theory to quantized networks is non-trivial because quantization operators violate smoothness assumptions required for standard mean-field analysis [2].

This paper addresses this challenge by analyzing the limiting behavior of smooth quantization approximations as smoothing parameters vanish. Our key insight is that the exponential decay inherent in smooth approximations exactly compensates for apparent singularities, enabling rigorous mean-field analysis. We prove that these dynamics converge to a well-defined limit governed by constrained transport equations [1], providing the first mathematical justification for quantized network training through distributional gradient analysis.

The mathematical framework developed in this work exhibits structural parallels with key concepts in high energy physics theory. The mean‐field limit of neural network training dynamics resembles the holographic principle in AdS/CFT correspondence, where bulk gravitational dynamics relate to boundary conformal field theory.

1.1 Related work

Mean-field theory for neural networks. The mean-field analysis of neural network training was developed by [8, 3], with extensions to deep networks [11, 7]. The connection to optimal transport was established through Wasserstein gradient flows [1, 10].

Quantized neural networks. Binary neural networks were introduced by [4, 9, 5], with the straight-through estimator formalized by [2, 14]. Recent advances include BitNet-like architectures [12] and comprehensive surveys. Theoretical analysis includes approximation theory [6, 13] and generalization bounds [15].

2 Modeling BitNet-like architectures with Smooth Quantization

2.1 Notation and constraint structure

Fix depth $L\in\mathbb{N}$. For layer $\ell\in\{1,\dots,L\}$ with width $n_{\ell}$ and fan-in $m_{\ell}$, the latent weight matrix is $W^{(\ell)}\in\mathbb{R}^{n_{\ell}\times m_{\ell}}$.

For a matrix $A\in\mathbb{R}^{m\times n}$, we denote:

• •

$\|A\|_{F}:=\sqrt{\sum_{i=1}^{m}\sum_{j=1}^{n}|A_{ij}|^{2}}$ the Frobenius norm,

• •

$\|A\|_{\infty}:=\max_{i,j}|A_{ij}|$ the max norm,

• •

$\langle A,B\rangle:=\sum_{i,j}A_{ij}B_{ij}$ the Frobenius inner product.

Define the layer mean

[TABLE]

the projection onto the zero-mean subspace

[TABLE]

and the constraint sets $\mathcal{H}^{(\ell)}_{c}:=\{W:\Psi^{(\ell)}(W)=c\}$, $\mathcal{H}^{(\ell)}_{0}=\ker(\Psi^{(\ell)})$. The following properties are standard and used repeatedly.

Lemma 2.1** (Orthogonal decomposition and isometries).**

For each layer $\ell$ and $W\in\mathbb{R}^{n_{\ell}\times m_{\ell}}$: (i) $W=\alpha^{(\ell)}(W)\,\mathbf{1}+P^{(\ell)}(W)$ with $\langle\mathbf{1},P^{(\ell)}(W)\rangle=0$; (ii) $P^{(\ell)}$ is linear, self-adjoint, idempotent; (iii) $\|W\|_{F}^{2}=n_{\ell}m_{\ell}\,\left\lvert\alpha^{(\ell)}(W)\right\rvert^{2}+\|P^{(\ell)}(W)\|_{F}^{2}$; (iv) $T_{c}(W)=W+c\,\mathbf{1}$ is an isometry $\mathcal{H}^{(\ell)}_{0}\to\mathcal{H}^{(\ell)}_{c}$.

Proof.

Fix $W\in\mathbb{R}^{n_{\ell}\times m_{\ell}}$.

Part (i): By definitions (2.1) and (2.2), we have

[TABLE]

For orthogonality, compute

[TABLE]

where we used the fact that $\langle\mathbf{1},\mathbf{1}\rangle=n_{\ell}m_{\ell}$ and $\alpha^{(\ell)}(A)=\frac{1}{n_{\ell}m_{\ell}}\langle A,\mathbf{1}\rangle$ by definition.

Part (ii): Linearity of $P^{(\ell)}$ follows immediately from linearity of matrix operations and the scalar $\alpha^{(\ell)}(\cdot)$. For self-adjointness, let $A,B\in\mathbb{R}^{n_{\ell}\times m_{\ell}}$:

[TABLE]

For idempotence, note that for any matrix $W$,

[TABLE]

Therefore, $P^{(\ell)}(P^{(\ell)}(W))=P^{(\ell)}(W)-\alpha^{(\ell)}(P^{(\ell)}(W))\mathbf{1}=P^{(\ell)}(W)-0=P^{(\ell)}(W)$.

Part (iii): By the orthogonality established in part (i), we have

[TABLE]

Part (iv): For $A,B\in\mathcal{H}^{(\ell)}_{0}$, we have $\alpha^{(\ell)}(A)=\alpha^{(\ell)}(B)=0$. Then

[TABLE]

proving that $T_{c}$ is an isometry. Since $T_{c}(A)=A+c\mathbf{1}$ and $\alpha^{(\ell)}(A)=0$, we have $\alpha^{(\ell)}(T_{c}(A))=c$, so $T_{c}(A)\in\mathcal{H}^{(\ell)}_{c}$. This establishes the mapping $\mathcal{H}^{(\ell)}_{0}\to\mathcal{H}^{(\ell)}_{c}$. ∎

2.2 Smooth quantization and dequantization

Quantized weights in BitNet-like architectures are the signs of centered latent weights, with a scaling to preserve variance [12]. To make the forward map differentiable, we adopt smooth surrogates.

Definition 2.1** (Smooth sign, clip, and absolute value).**

For $\varepsilon\in(0,1]$, define

[TABLE]

Then $\mathrm{clip}_{\varepsilon}$ is $C^{\infty}$ and $1$-Lipschitz uniformly in $\varepsilon$.

Definition 2.2** (Smoothed BitLinear layer).**

Let $X\in\mathbb{R}^{m_{\ell}}$ be an input. The smoothed quantized weight is

[TABLE]

Define a smooth $L^{1}$-scale

[TABLE]

and a smooth absmax-like activation quantizer

[TABLE]

with fixed $b\in\mathbb{N}$, $Q_{b}=2^{b-1}$ and small $\delta\in(0,1)$. The layer map is

[TABLE]

Remark 2.1** (On STE consistency).**

The straight-through estimator (STE) is typically implemented by replacing $\partial\mathrm{sign}$ with an identity or bounded truncation in a margin [12]. Our $\mathrm{sgn}_{\varepsilon}$ provides a differentiable surrogate with uniformly bounded derivatives on compacta, making chain-rule gradients well-defined. In Section 4 we discuss stability of the mean-field limit as $\varepsilon\downarrow 0$.

2.3 Network, loss, and dynamics

Definition 2.3** (Layer dimensions and network architecture).**

Each layer $\ell\in\{1,\ldots,L\}$ defines a map $h^{(\ell)}:\mathbb{R}^{m_{\ell}}\to\mathbb{R}^{n_{\ell}}$ where:

• •

$m_{\ell}$* is the input dimension (fan-in) of layer $\ell$*

• •

$n_{\ell}$* is the output dimension (width) of layer $\ell$*

• •

For consistency: $n_{\ell-1}=m_{\ell}$ for $\ell\geq 2$

For regression tasks, we assume $n_{L}=1$ so that $f_{W}(x)\in\mathbb{R}$ is scalar-valued. For multi-class classification with $K$ classes, $n_{L}=K$ and $f_{W}(x)\in\mathbb{R}^{K}$.

The $L$-layer forward recursion is defined as:

[TABLE]

Here, $\sigma^{(\ell)}:\mathbb{R}^{n_{\ell}}\to\mathbb{R}^{n_{\ell}}$ denotes component-wise application of the activation function.

With smooth activations $\sigma^{(\ell)}$ (Assumption 3.1), for a data law $\pi$ on $\mathbb{R}^{d}\times\mathbb{R}^{n_{L}}$ and a $C^{2}$ loss $\ell:\mathbb{R}^{n_{L}}\times\mathbb{R}^{n_{L}}\to\mathbb{R}$, define the population risk

[TABLE]

We study discrete-time gradient descent

[TABLE]

with time-interpolation $t=k\eta$ and continuous-time limit $\eta\downarrow 0$.

3 Assumptions and basic estimates

Assumption 3.1** (Regularity and boundedness).**

Fix $T>0$ and constants $R,A_{\ell},C_{\ell},D_{\ell},L_{1},L_{2},M>0$.

• (R1)

Data: $\pi$ has compact support; $\|X\|_{\infty}\leq R$, $\|Y\|_{\infty}\leq R$ a.s.

• (R2)

Loss: $\ell\in C^{2}(\mathbb{R}^{n_{L}}\times\mathbb{R}^{n_{L}})$ with $\|\nabla^{2}\ell\|_{\text{op}}\leq L_{2}$ and $\|\nabla_{1}\ell(u,y)\|_{2}\leq L_{1}(1+\|u\|_{2})$.

• (R3)

Activations: $\sigma^{(\ell)}\in C^{2}(\mathbb{R})$ with $\|(\sigma^{(\ell)})^{\prime}\|_{\infty}\leq C_{\ell}$, $\|(\sigma^{(\ell)})^{\prime\prime}\|_{\infty}\leq D_{\ell}$, and $|\sigma^{(\ell)}(z)|\leq A_{\ell}(1+|z|)$.

• (R4)

Initialization and boundedness: $\sup_{n,\ell,i,j}\left\lvert W^{(\ell)}_{ij}(0)\right\rvert\leq M$, and there exists $M_{\star}=M_{\star}(T,L_{1},L_{2},\{C_{\ell},D_{\ell}\}_{\ell},M)<\infty$ such that all iterates satisfy $\|W^{(\ell)}(t)\|_{\infty}\leq M_{\star}$ for $t\in[0,T]$ through projection $\Pi_{\mathcal{B}_{M_{\star}}}$ where $\Pi_{\mathcal{B}_{M_{\star}}}(W)_{ij}=\mathrm{clip}(W_{ij};-M_{\star},M_{\star})$.

• (R5)

Smoothing parameters: Fix $\varepsilon\in(0,1]$, $b\in\mathbb{N}$, $\delta\in(0,1)$ throughout the analysis of the mean-field limit.

Lemma 3.1** (Lipschitzness of the forward map).**

Under Assumption 3.1, for fixed parameters $\varepsilon\in(0,1]$, $b\in\mathbb{N}$, $\delta\in(0,1)$, there exists a constant $L_{\mathrm{fwd}}=L_{\mathrm{fwd}}(\varepsilon,b,\delta,\{A_{\ell},C_{\ell},D_{\ell}\}_{\ell=1}^{L},M_{\star},R)<\infty$ such that for all $x$ in the support of $\pi$ and all weight configurations $W,\widehat{W}$ in the compact domain $K$,

[TABLE]

Moreover, $W\mapsto\mathcal{R}_{\varepsilon}(W)$ is $C^{1}$ with $\nabla\mathcal{R}_{\varepsilon}$ locally Lipschitz on $K$.

Proof.

By Assumption 3.1(R4), there exists $M_{\star}<\infty$ such that all iterates satisfy $\sup_{t\in[0,T]}\sup_{\ell,i,j}|W^{(\ell)}_{ij}(t)|\leq M_{\star}$, and we denote $K:=\{W:\|W^{(\ell)}\|_{\infty}\leq M_{\star},\forall\ell\}$.

Lipschitz Constants of Smooth Quantizers The smooth quantizers satisfy the following uniform Lipschitz properties:

1. (i)

$\mathrm{sgn}_{\varepsilon}(z)=\tanh(z/\varepsilon)$ is $\varepsilon^{-1}$-Lipschitz since

[TABLE] 2. (ii)

$|\cdot|_{\varepsilon}(z)=\sqrt{z^{2}+\varepsilon^{2}}$ is $1$-Lipschitz since

[TABLE] 3. (iii)

$\mathrm{clip}_{\varepsilon}(x;a,b)$ is $1$-Lipschitz uniformly in $\varepsilon$.

Analysis of $\gamma_{\varepsilon}$ Function For $\gamma_{\varepsilon}(x)=\max\{\varepsilon,\|x\|_{\infty}\}$, we have:

[TABLE]

Bounds for $\beta^{(\ell)}_{\varepsilon}$ Define $C_{\beta}:=\sqrt{(2M_{\star})^{2}+\varepsilon^{2}}$. Then:

1. (i)

Upper bound: $\beta^{(\ell)}_{\varepsilon}(W)\leq C_{\beta}$ since $|P^{(\ell)}(W)_{ij}|_{\varepsilon}\leq C_{\beta}$. 2. (ii)

Lipschitz property: $|\beta^{(\ell)}_{\varepsilon}(W)-\beta^{(\ell)}_{\varepsilon}(\widehat{W})|\leq\frac{1}{\sqrt{n_{\ell}m_{\ell}}}\|W^{(\ell)}-\widehat{W}^{(\ell)}\|_{F}$.

Proof of (ii): By linearity of $P^{(\ell)}$ and 1-Lipschitz property of $|\cdot|_{\varepsilon}$:

[TABLE]

Complete Inductive Proof We prove by strong induction that for each layer $\ell$:

[TABLE]

for appropriate constants $L^{(\ell)}$.

Base case ($\ell=0$): $h^{(0)}(x)=x$, so the inequality holds trivially with $L^{(0)}=0$.

Inductive step: Assume the statement holds for all layers $j<\ell$. The layer-$\ell$ map is:

[TABLE]

Let $u_{W}:=\beta^{(\ell)}_{\varepsilon}(W^{(\ell)})\widetilde{W}^{(\ell)}_{\varepsilon}\mathrm{Quant}^{(b)}_{\varepsilon}(h^{(\ell-1)}_{W}(x))$ and $u_{\widehat{W}}:=\beta^{(\ell)}_{\varepsilon}(\widehat{W}^{(\ell)})\widetilde{\widehat{W}}^{(\ell)}_{\varepsilon}\mathrm{Quant}^{(b)}_{\varepsilon}(h^{(\ell-1)}_{\widehat{W}}(x))$.

Since $\sigma^{(\ell)}$ is applied component-wise with Lipschitz constant $C_{\ell}$:

[TABLE]

For the matrix-vector product, we have:

[TABLE]

Using the triangle inequality and submultiplicativity of matrix norms:

[TABLE]

Using the bounds:

• •

$\|\widetilde{W}^{(\ell)}_{\varepsilon}\|_{2}\leq\sqrt{n_{\ell}m_{\ell}}$ (since each entry is bounded by 1)

• •

$\|\widetilde{W}^{(\ell)}_{\varepsilon}-\widetilde{\widehat{W}}^{(\ell)}_{\varepsilon}\|_{2}\leq\varepsilon^{-1}\|W^{(\ell)}-\widehat{W}^{(\ell)}\|_{F}$

• •

$\|\mathrm{Quant}^{(b)}_{\varepsilon}(u)-\mathrm{Quant}^{(b)}_{\varepsilon}(v)\|_{2}\leq\frac{Q_{b}}{\varepsilon}\|u-v\|_{2}$

• •

From Step 3: $|\beta^{(\ell)}_{\varepsilon}(W)-\beta^{(\ell)}_{\varepsilon}(\widehat{W})|\leq\frac{1}{\sqrt{n_{\ell}m_{\ell}}}\|W^{(\ell)}-\widehat{W}^{(\ell)}\|_{F}$

• •

Inductive hypothesis: $\|h^{(\ell-1)}_{W}(x)-h^{(\ell-1)}_{\widehat{W}}(x)\|_{2}\leq L^{(\ell-1)}\sum_{k=1}^{\ell-1}\|W^{(k)}-\widehat{W}^{(k)}\|_{F}$

Combining these estimates:

[TABLE]

This gives us the recursive relation:

[TABLE]

Setting $L_{\mathrm{fwd}}:=L^{(L)}$ completes the induction.

$C^{1}$** Regularity** The composition $\mathcal{F}(W)(x)=f_{W}(x)$ is $C^{1}$ on $K$ because:

1. (i)

Each smooth quantizer $\mathrm{sgn}_{\varepsilon}$, $|\cdot|_{\varepsilon}$, $\mathrm{clip}_{\varepsilon}$ is $C^{\infty}$. 2. (ii)

Matrix operations and function compositions preserve $C^{1}$ regularity. 3. (iii)

The chain rule applies on the bounded domain $K$.

Therefore, $\mathcal{R}_{\varepsilon}(W)=\mathbb{E}[\ell(f_{W}(X),Y)]$ is $C^{1}$ with locally Lipschitz gradient on $K$ by dominated convergence and the uniform bounds on $K$. ∎

4 Asymptotic Analysis as $\varepsilon\to 0$

This section studies the limiting behavior of the mean-field dynamics when the smoothing parameter $\varepsilon$ approaches zero. Our goal is to rigorously characterize the asymptotic properties of the solution $\bm{\mu}_{\varepsilon}=(\mu^{(1)}_{\varepsilon},\ldots,\mu^{(L)}_{\varepsilon})$ to the constrained continuity equations

[TABLE]

where

[TABLE]

with $\mathcal{R}^{(\ell)}_{\varepsilon}[\bm{\mu}]$ the functional derivative of the smoothed risk $\mathcal{R}_{\varepsilon}$.

4.1 Exponential decay and distributional analysis

Recall the smooth sign activation and its derivative:

[TABLE]

Lemma 4.1** (Exponential decay and uniform bounds).**

For any $\varepsilon\in(0,1]$ and $z\in\mathbb{R}$:

1. (i)

$\mathrm{sgn}^{\prime}_{\varepsilon}(z)=\frac{4}{\varepsilon}\frac{1}{(e^{|z|/\varepsilon}+e^{-|z|/\varepsilon})^{2}}\leq\frac{4}{\varepsilon}e^{-2|z|/\varepsilon}$* for $z\neq 0$.* 2. (ii)

$\int_{\mathbb{R}}\mathrm{sgn}^{\prime}_{\varepsilon}(z)\,dz=2$* for all $\varepsilon>0$.* 3. (iii)

For any bounded measurable function $\phi:\mathbb{R}\to\mathbb{R}$ with $\|\phi\|_{\infty}\leq M$:

[TABLE]

Proof.

Part (i): We have $\operatorname{sech}^{2}(u)=\frac{4}{(e^{u}+e^{-u})^{2}}$. For $u\neq 0$, the denominator $(e^{|u|}+e^{-|u|})^{2}\geq e^{2|u|}$, giving the stated bound.

Part (ii): By substitution $u=z/\varepsilon$:

[TABLE]

Part (iii): By the boundedness of $\phi$ and part (ii):

[TABLE]

∎

Lemma 4.2** (Distributional convergence).**

As $\varepsilon\downarrow 0$, we have $\mathrm{sgn}_{\varepsilon}(z)\to\mathrm{sign}(z)$ pointwise and

[TABLE]

where $\delta_{0}$ is the Dirac delta at zero.

Proof.

For any test function $\phi\in C_{c}^{\infty}(\mathbb{R})$:

[TABLE]

As $\varepsilon\to 0$, $\phi(\varepsilon u)\to\phi(0)$ uniformly on compact sets. Since $\int_{\mathbb{R}}\operatorname{sech}^{2}(u)\,du=2$, the dominated convergence theorem yields:

[TABLE]

∎

4.2 Uniform velocity field bounds via natural cancellation

The key insight is that the exponential decay of $\mathrm{sgn}^{\prime}_{\varepsilon}$ exactly compensates for the $\varepsilon^{-1}$ factor, yielding uniform bounds without requiring measure concentration.

Lemma 4.3** (Uniform bounds on singular integrals).**

Let $\mu$ be any probability measure on $\mathbb{R}^{m_{\ell}}$ with support in the compact set $\mathcal{K}:=\{w:\|w\|_{\infty}\leq M_{\star}\}$. For any bounded measurable function $\phi:\mathcal{K}\to\mathbb{R}$ and any $(i,j)\in\{1,\ldots,n_{\ell}\}\times\{1,\ldots,m_{\ell}\}$:

[TABLE]

uniformly in $\varepsilon\in(0,1]$.

Proof.

Since $P^{(\ell)}(w)_{ij}$ is a linear function of $w$ and $\mu$ is a probability measure, we can write this as an integral over $\mathbb{R}$ with respect to the pushforward measure $\nu:=(P^{(\ell)}(\cdot)_{ij})_{\#}\mu$:

[TABLE]

where $\tilde{\phi}(z)$ represents the conditional expectation of $\phi(w)$ given $P^{(\ell)}(w)_{ij}=z$, which satisfies $\|\tilde{\phi}\|_{\infty}\leq\|\phi\|_{\infty}$.

By Lemma 4.1(iii), since $\nu$ is a probability measure:

[TABLE]

∎

4.3 Main convergence theorem

Theorem 4.1** (Stability and convergence as $\varepsilon\to 0$).**

Suppose Assumptions 3.1 and A.1 hold uniformly in $\varepsilon\in(0,1]$. Let $\{\bm{\mu}_{\varepsilon}\}_{\varepsilon>0}$ be the unique solutions to the continuity equations (4.1) with velocities (4.2).

Then there exists a subsequence $\varepsilon_{k}\downarrow 0$ and a limit curve $\bm{\mu}_{0}\in C([0,T],\prod_{\ell=1}^{L}\mathcal{P}_{2}(\mathbb{R}^{m_{\ell}}))$ such that

[TABLE]

where $\mathcal{P}_{2}$ denotes the space of probability measures with finite second moment.

Furthermore, $\bm{\mu}_{0}$ solves a constrained transport equation of the form

[TABLE]

where $v^{(\ell)}_{0}$ is the limiting velocity field associated to the non-smoothed risk functional $\mathcal{R}_{0}$.

Proof.

The proof follows a compactness-identification strategy, utilizing the natural exponential decay of tanh derivatives to establish uniform bounds.

From the chain rule and gradient bounds (Lemma B.1), the velocity field has the structure:

[TABLE]

We analyze each term $\frac{\partial h^{(k)}}{\partial w}$ systematically. For $k=\ell$, we have:

[TABLE]

By the chain rule and product rule:

[TABLE]

The potentially problematic term is:

[TABLE]

Since $P^{(\ell)}(w)$ is componentwise linear in $w$, we have $\left\|\frac{\partial P^{(\ell)}(w)}{\partial w}\right\|_{\text{op}}\leq 1$.

The critical observation is that while $\mathrm{sgn}^{\prime}_{\varepsilon}(z)=\varepsilon^{-1}\operatorname{sech}^{2}(z/\varepsilon)$ contains the factor $\varepsilon^{-1}$, when this appears in the velocity field, it takes the form:

[TABLE]

for some bounded function $\phi$ arising from the loss and network architecture.

By Assumption 3.1(R4), we have:

• •

$\beta^{(\ell)}_{\varepsilon}(w)\leq C_{\beta}$ uniformly for some constant $C_{\beta}$ independent of $\varepsilon$

• •

$\|\phi\|_{\infty}\leq C_{\phi}$ for some constant $C_{\phi}$ from the boundedness of activations, loss derivatives, and network depth

Applying Lemma 4.3 with the bounded function $\psi(w):=\phi(w)\beta^{(\ell)}_{\varepsilon}(w)$, which satisfies $\|\psi\|_{\infty}\leq C_{\phi}C_{\beta}$:

[TABLE]

uniformly in $\varepsilon\in(0,1]$.

The other terms in $\frac{\partial h^{(\ell)}}{\partial w}$ are:

$\frac{\partial\beta^{(\ell)}_{\varepsilon}}{\partial w}\widetilde{W}^{(\ell)}_{\varepsilon}$: This is bounded since $\left|\frac{\partial\beta^{(\ell)}_{\varepsilon}}{\partial w}\right|\leq(n_{\ell}m_{\ell})^{-1}$ and $\|\widetilde{W}^{(\ell)}_{\varepsilon}\|_{\infty}\leq 1$. 2. 2.

Terms involving $\frac{\partial\mathrm{Quant}^{(b)}_{\varepsilon}}{\partial w}$: These have bounded derivatives by the smoothness of $\mathrm{clip}_{\varepsilon}$. 3. 3.

Terms involving $\frac{\partial h^{(\ell-1)}}{\partial w}$: These contribute through the recursive structure but do not introduce additional $\varepsilon^{-1}$ singularities beyond those already controlled.

By strong induction on layers $k=\ell,\ell+1,\ldots,L$, we can show that each $\frac{\partial h^{(k)}}{\partial w}$ satisfies a uniform bound independent of $\varepsilon$. The base case $k=\ell$ follows from the analysis above, and the inductive step follows by applying the same reasoning to the composition structure.

Combining all bounded terms in the expression for $v^{(\ell)}_{\varepsilon}(w,t)$:

[TABLE]

where:

• •

$L_{1}$ comes from Assumption 3.1(R2) bounding the loss derivative

• •

$C_{k}^{\text{Lip}}$ are the Lipschitz constants of the activations from Assumption 3.1(R3)

• •

$C_{k}^{\text{grad}}$ are the uniform bounds on $\left\|\frac{\partial h^{(k)}}{\partial w}\right\|$ established above

Since each of these constants is independent of $\varepsilon\in(0,1]$, we conclude that $C_{\text{uniform}}$ is independent of $\varepsilon$.

The uniform bound implies equicontinuity in the Wasserstein metric:

[TABLE]

for all $\varepsilon\in(0,1]$.

By constraint preservation (Lemma A.1) and boundedness assumptions (Assumption 3.1(R4)), all measures $\mu^{(\ell)}_{\varepsilon}(t)$ have support in the compact set $\mathcal{K}$ and satisfy:

[TABLE]

By the Arzelà-Ascoli theorem in $C([0,T],\mathcal{P}_{2}(\mathcal{K}))$, there exists a subsequence $\varepsilon_{k}\downarrow 0$ and a limit $\bm{\mu}_{0}$ such that:

[TABLE]

We now prove that the limit velocity field $v^{(\ell)}_{0}$ corresponds to the distributional gradient of the non-smoothed risk functional $\mathcal{R}_{0}$. The key is to show that integrals involving $\mathrm{sgn}^{\prime}_{\varepsilon_{k}}$ converge to the appropriate distributional limit.

Claim: For any test function $\varphi\in C_{c}^{1}(\mathbb{R}^{m_{\ell}})$ and any $(i,j)\in\{1,\ldots,n_{\ell}\}\times\{1,\ldots,m_{\ell}\}$:

[TABLE]

Proof of Claim: Let $\delta>0$ be arbitrary. We decompose the integration domain as:

[TABLE]

where:

[TABLE]

For $w\in\mathcal{K}_{\delta}^{+}$, we have $P^{(\ell)}(w)_{ij}>\delta$, so by Lemma 4.1(i):

[TABLE]

Since $\varphi$ is compactly supported with $\|\varphi\|_{\infty}\leq C_{\varphi}$ for some constant $C_{\varphi}$:

[TABLE]

As $k\to\infty$ (i.e., $\varepsilon_{k}\downarrow 0$), we have $\frac{1}{\varepsilon_{k}}e^{-2\delta/\varepsilon_{k}}\to 0$ exponentially fast. Similarly for $\mathcal{K}_{\delta}^{-}$.

Therefore:

[TABLE]

On $\mathcal{K}_{\delta}^{0}$, we have $|P^{(\ell)}(w)_{ij}|\leq\delta$. We use the change of variables $z=P^{(\ell)}(w)_{ij}$ and define the pushforward measure:

[TABLE]

Then:

[TABLE]

where $\widetilde{\varphi}_{\varepsilon_{k}}(z)$ is the conditional expectation of $\varphi(w)$ given $P^{(\ell)}(w)_{ij}=z$ and $w\in\mathcal{K}_{\delta}^{0}$.

By weak convergence $\mu^{(\ell)}_{\varepsilon_{k}}\rightharpoonup\mu^{(\ell)}_{0}$, we have $\nu_{k}^{\delta}\rightharpoonup\nu_{0}^{\delta}$ where $\nu_{0}^{\delta}:=(P^{(\ell)}(\cdot)_{ij})_{\#}(\mu^{(\ell)}_{0}|_{\mathcal{K}_{\delta}^{0}})$.

Since $\widetilde{\varphi}_{\varepsilon_{k}}(z)\to\widetilde{\varphi}_{0}(z)$ boundedly (by compactness), and by Lemma 4.2:

[TABLE]

Now we analyze what happens as $\delta\downarrow 0$. We have:

[TABLE]

as $\delta\downarrow 0$.

Since $\widetilde{\varphi}_{0}(0)$ is the conditional expectation of $\varphi(w)$ given $P^{(\ell)}(w)_{ij}=0$:

[TABLE]

Combining and taking $\delta\downarrow 0$:

[TABLE]

This shows that the limit velocity field $v^{(\ell)}_{0}$ corresponds to the distributional gradient of the non-smoothed risk functional $\mathcal{R}_{0}$, where the derivative of the sign function is replaced by twice the Dirac delta at zero.

For any test function $\varphi\in C_{c}^{1}(\mathbb{R}^{m_{\ell}})$ and $0\leq s<t\leq T$, the uniform bounds from Step 1 allow us to pass to the limit in:

[TABLE]

Differentiating with respect to $t$ yields the weak formulation of (4.4).

The zero-mean constraint is preserved in the limit since for any $t\in[0,T]$:

[TABLE]

by Lemma A.1 and weak convergence of measures. ∎

Remark 4.1** (Connection to Straight-Through Estimators).**

Theorem 4.1 provides a rigorous foundation for the straight-through estimator (STE) commonly used in BitNet training. The limiting dynamics correspond to gradient descent on the non-smoothed risk functional, where the singular gradients of the sign function are naturally regularized by the exponential decay inherent in the tanh smoothing. This validates the STE approximation as the mathematically correct limit of smooth quantization.

5 Conclusion

This work presents the first rigorous mean-field analysis of deep BitNet-like architectures under smooth quantization. By introducing differentiable surrogates for the sign and clipping functions, we establish well-posedness of the training dynamics in the space of probability measures and prove convergence of the empirical weight distributions to solutions of constrained transport equations as the smoothing parameter $\varepsilon\to 0$. Our key technical insight is that the natural exponential decay in the derivatives of $\tanh(z/\varepsilon)$ perfectly offsets the singular $\varepsilon^{-1}$ scaling, yielding uniform bounds on the velocity fields without requiring additional measure concentration arguments. Consequently, we rigorously justify the straight-through estimator as the correct limiting gradient flow for quantized networks. Future work includes extending this framework to true hard quantizers via differential inclusions and relaxing compactness assumptions on the weight domain.

Acknowledgments

The author thanks the maintainers of BitNet for inspiring the latent–quantized duality viewpoint.

Appendix A Empirical measures and mean-field limit

A.1 Row-wise empirical measures and constraint preservation

Recall from Definition 2.3 that each layer $\ell\in\{1,\ldots,L\}$ has weight matrix $W^{(\ell)}\in\mathbb{R}^{n_{\ell}\times m_{\ell}}$. For analysis purposes, we index the rows of $W^{(\ell)}$ by $i=1,\ldots,n_{\ell}$ and denote the $i$-th row as $w^{(\ell)}_{i}\in\mathbb{R}^{m_{\ell}}$. Thus:

[TABLE]

At discrete time $k$, define the empirical measure on $\mathbb{R}^{m_{\ell}}$:

[TABLE]

where $\delta_{x}$ denotes the Dirac measure at point $x\in\mathbb{R}^{m_{\ell}}$. For continuous-time analysis with interpolation $t=k\eta$, we define:

[TABLE]

Let $\bm{\widehat{\mu}}_{n}:=(\widehat{\mu}^{(1)}_{n_{1}},\ldots,\widehat{\mu}^{(L)}_{n_{L}})$ denote the collection of empirical measures across all layers.

The layer means $\alpha^{(\ell)}(k):=\Psi^{(\ell)}(W^{(\ell)}(k))$ from (2.1) satisfy the following evolution under gradient descent.

Lemma A.1** (Constraint preservation).**

Let the updates be (2.23). Then for all $k\geq 0$ and each $\ell$,

[TABLE]

so in continuous time with $t=k\eta$,

[TABLE]

Proof.

By linearity of $\Psi^{(\ell)}$ and the gradient descent update (2.23),

[TABLE]

Dividing by $\eta$ and passing to the limit $\eta\downarrow 0$ yields the differential form. ∎

A.2 Functional derivatives and velocity fields

For a collection of probability measures $\bm{\mu}=(\mu^{(1)},\ldots,\mu^{(L)})$ on $\mathbb{R}^{m_{1}}\times\cdots\times\mathbb{R}^{m_{L}}$, define the population risk functional:

[TABLE]

where $f_{\bm{\mu}}(x)$ represents the network output when layer weights are sampled according to the measures $\bm{\mu}$.

The functional derivative $\mathcal{R}^{(\ell)}_{\varepsilon}[\bm{\mu}]:\mathbb{R}^{m_{\ell}}\to\mathbb{R}$ is defined as the Gateaux derivative with respect to perturbations in $\mu^{(\ell)}$:

[TABLE]

Define the velocity fields $v^{(\ell)}:\mathbb{R}^{m_{\ell}}\times[0,T]\to\mathbb{R}^{m_{\ell}}$ by:

[TABLE]

where $\nabla_{w}$ denotes the gradient with respect to the row variable $w\in\mathbb{R}^{m_{\ell}}$.

A.3 Continuity equations and transport structure

The mean-field limit is characterized by the coupled system of continuity equations:

[TABLE]

in the sense of distributions on $\mathbb{R}^{m_{\ell}}\times(0,T)$.

Assumption A.1** (Regularity of velocity fields).**

There exists $L_{v}>0$ such that for all $\ell$ and all $\bm{\mu},\bm{\nu}$ with supports in a fixed compact set $\mathcal{K}\subset\mathbb{R}^{m_{\ell}}$ and satisfying $\int\|w\|^{2}d\mu^{(j)}(w),\int\|w\|^{2}d\nu^{(j)}(w)\leq M^{2}$ for all $j$ and some $M>0$:

1. (i)

Lipschitz dependence on measures:**

[TABLE] 2. (ii)

Spatial regularity:* For each fixed $\bm{\mu}$, the map $w\mapsto v^{(\ell)}(w;\bm{\mu})$ is globally Lipschitz with constant $L_{v}$ on $\mathcal{K}$.*

A.4 Mean-field convergence theorem

Theorem A.1** (Weak convergence to mean-field limit).**

Fix $\varepsilon\in(0,1]$, $b\in\mathbb{N}$, $\delta\in(0,1)$, and $T>0$. Under Assumptions 3.1 and A.1, let $n_{\ell}\to\infty$ for all $\ell$ with $n_{\ell}/n_{j}\to r_{\ell j}\in(0,\infty)$ and let $\eta\downarrow 0$ with $k\eta\to t\in[0,T]$.

Then the empirical process $\bm{\widehat{\mu}}_{n}$ converges weakly in $C([0,T],\prod_{\ell=1}^{L}\mathcal{P}(\mathbb{R}^{m_{\ell}}))$ to a unique $\bm{\mu}=(\mu^{(1)},\ldots,\mu^{(L)})$ that solves the coupled system (A.3) with velocity fields (A.2).

Moreover, for each $\ell$ and all $\varphi\in C_{c}^{1}(\mathbb{R}^{m_{\ell}})$:

[TABLE]

Proof.

The proof proceeds through four main steps: compactness, velocity field regularity verification, limit identification, and uniqueness.

Compactness of empirical measures.

By Assumption 3.1(R4), all row vectors $w_{i}^{(\ell)}(k)$ remain in the compact set $\mathcal{K}:=\{w\in\mathbb{R}^{m_{\ell}}:\|w\|_{\infty}\leq M_{\star}\}$ for $t\in[0,T]$.

From Lemma B.1, there exists $M_{\text{grad}}<\infty$ such that:

[TABLE]

This yields uniformly bounded increments:

[TABLE]

For equicontinuity, given $\epsilon>0$, choose $\delta=\epsilon/(2M_{\text{grad}})$ and $\eta<\epsilon/(6M_{\text{grad}})$. Then for $|t-s|<\delta$:

[TABLE]

By compactness of $\mathcal{P}(\mathcal{K})$ in the Wasserstein topology and the Arzelà-Ascoli theorem, $\{\bm{\widehat{\mu}}_{n}\}$ is relatively compact in $C([0,T],\prod_{\ell=1}^{L}\mathcal{P}(\mathcal{K}))$.

Velocity field regularity.

The functional derivative satisfies:

[TABLE]

where $f_{\bm{\mu},w}^{(\ell)}(x)$ denotes the network output when layer $\ell$ has an additional infinitesimal mass at position $w$.

By Lemma 3.1 and the chain rule, for $w,\tilde{w}\in\mathcal{K}$:

[TABLE]

This establishes Lipschitz continuity of the functional derivative, ensuring that $v^{(\ell)}(w,t)$ exists almost everywhere with:

[TABLE]

Limit identification.

Let $\bm{\mu}$ be any limit point of $\{\bm{\widehat{\mu}}_{n}\}$. For $\varphi\in C_{c}^{1}(\mathbb{R}^{m_{\ell}})$ and $0\leq s<t\leq T$, we perform discrete integration by parts:

[TABLE]

By Taylor expansion and uniform bounds:

[TABLE]

Substituting the gradient descent updates and taking limits:

[TABLE]

Differentiating with respect to $t$ yields the weak formulation of (A.3).

Uniqueness.

Let $\bm{\mu},\bm{\nu}$ be two solutions with identical initial conditions. Define:

[TABLE]

By Assumption A.1 and the contraction property of optimal transport:

[TABLE]

Summing over $\ell$ and applying Grönwall’s inequality with $d(0)=0$ yields $d(t)=0$ for all $t\in[0,T]$, establishing uniqueness. ∎

A.5 Interacting particle system interpretation

Remark A.1** (Connection to particle systems).**

The empirical measures (A.1) admit a natural interpretation in terms of interacting particle systems. Each row vector $w^{(\ell)}_{i}(k)\in\mathbb{R}^{m_{\ell}}$ can be viewed as the position of the $i$-th particle in layer $\ell$ at time $k$.

Under this interpretation:

• •

The empirical measure $\widehat{\mu}^{(\ell)}_{n_{\ell}}(k)=\frac{1}{n_{\ell}}\sum_{i=1}^{n_{\ell}}\delta_{w^{(\ell)}_{i}(k)}$ represents the spatial distribution of particles in layer $\ell$.

• •

The gradient descent update (2.23) becomes a system of interacting particles:

[TABLE]

where the force on particle $i$ depends on the positions of all particles across all layers.

• •

The mean-field limit corresponds to the thermodynamic limit where the number of particles $n_{\ell}\to\infty$ while their individual influence vanishes as $1/n_{\ell}$.

• •

The velocity field $v^{(\ell)}(w,t)$ in (A.2) represents the drift experienced by a test particle at position $w$ in the mean-field environment.

This particle system perspective provides intuitive insight into the dynamics, while the measure-theoretic formulation in the preceding subsections provides the rigorous mathematical foundation for the analysis.

Appendix B Gradients, chain rule, and bounds

B.1 Layerwise gradients with smooth quantization

Let $W\mapsto f_{W}$ be defined with smooth quantizers. Then

[TABLE]

with all Jacobians well-defined by the chain rule. The derivative $\partial\widetilde{W}^{(\ell)}_{\varepsilon}/\partial W^{(\ell)}$ exists and is bounded by $\varepsilon^{-1}$ entrywise; the derivative of $\beta^{(\ell)}_{\varepsilon}$ has entries

[TABLE]

bounded by $(n_{\ell}m_{\ell})^{-1}$.

Lemma B.1** (Gradient bound).**

Under Assumption 3.1, there exist constants $C_{\ell}=C_{\ell}(\varepsilon,b,\delta)$ such that for all entries $(i,j)$,

[TABLE]

and $\nabla_{W^{(\ell)}}\mathcal{R}_{\varepsilon}$ is locally Lipschitz on the compact domain.

Proof.

Apply (B.1) and bound each factor by (R2)–(R3) together with the Lipschitz constants of the smooth quantizers on the compact set of iterates (R4). The derivative of $\mathrm{sgn}_{\varepsilon}$ is bounded by $\varepsilon^{-1}$, the derivative of $\beta^{(\ell)}_{\varepsilon}$ is bounded by $(n_{\ell}m_{\ell})^{-1}$, and $\mathrm{Quant}^{(b)}_{\varepsilon}$ has bounded Jacobian for fixed $\varepsilon,b,\delta$. The expectation over compactly supported $(X,Y)$ preserves these bounds, yielding the stated inequality. Local Lipschitzness follows from boundedness of second derivatives on compacta. ∎