1 Introduction↩︎

Large language models (LLMs) and related foundation models have demonstrated remarkable capabilities across a wide range of applications, including natural language processing [1], [2], computer vision [3], [4], and scientific tasks [5], [6]. Their successful deployment to downstream tasks, especially in professional domains and specialized communities, often relies on fine-tuning to adapt pretrained models to task-specific data and objectives [7], [8]. While full fine-tuning can be effective, updating all model parameters is computationally expensive and memory intensive, particularly for modern large-scale scenarios. Parameter-efficient fine-tuning methods alleviate this difficulty by introducing a small number of trainable parameters while keeping most pretrained weights fixed, thereby reducing the computational and storage costs of model adaptation [9], [10]. Among these fine-tuning methods, low-rank adaptation (LoRA) [11] has become a widely used approach due to its simplicity, efficiency, and strong empirical performance. LoRA represents task-specific updates through trainable low-rank matrices inserted into selected model modules, reducing the number of trainable parameters while preserving sufficient adaptation capacity.

A central question in LoRA fine-tuning is how to allocate a limited adaptation budget across layers and modules [12], [13]. Different parts of a model may play different roles: some modules may encode general-purpose knowledge that should be largely preserved, whereas others may need stronger adaptation to capture domain-specific information. Since the rank of each LoRA matrix controls the adaptation capacity, an effective allocation should assign more resources to components with strong statistical evidence while avoiding unnecessary over-parameterization. This creates a high-dimensional resource allocation problem under stochastic training dynamics, where observed component scores serve as noisy measures of component contributions and are affected by mini-batch noise and optimizer-induced randomness. From a statistical perspective, adaptive rank allocation should not rely merely on the magnitude of empirical component scores but should also account for the uncertainty in estimating component contributions. Motivated by this viewpoint, this paper develops a statistically principled resource allocation method for LoRA based on hypothesis testing.

Due to the matrix factorization structure of LoRA, the adaptation in each module can be decomposed into a sum of rank-one components. Consequently, allocating the rank budget is equivalent to deciding which rank-one components should be retained or removed across layers and modules. This perspective naturally leads to a component selection problem: each candidate component represents a potential direction of adaptation, and its empirical magnitude reflects the extent to which the pretrained model is modified along that direction. However, during fine-tuning, these quantities are generated by stochastic optimization algorithms and are affected by mini-batch sampling, adaptive step sizes (learning rates), momentum, and other sources of training noise. Therefore, deciding whether a LoRA component is useful should not be viewed merely as a deterministic ranking problem. Rather, it can be formulated as a statistical inference problem in which one seeks to distinguish genuine adaptation signals from stochastic observations.

Several adaptive-rank methods have been proposed to address LoRA rank allocation. AdaLoRA [12] is among the first methods to explicitly identify the importance of adaptive resource allocation and develops a practical procedure based on component-wise importance scores. These scores are constructed from gradient-derived sensitivity measures and smoothed magnitude information, with gradients interpreted as indicators of how sensitive the training objective is to each component. More recently, IGU-LoRA [14] further refined this idea by incorporating uncertainty-aware and layer-wise grouping corrections into the component importance measure, aiming to stabilize rank allocation under noisy training dynamics, inspired by signal-to-noise ratio. Other automatic rank allocation methods include dynamic rank training [13], gating-based rank allocation [15], and meta-learning-based rank selection [16]. Although these methods are empirically effective, their importance scores are constructed through carefully designed combinations of gradient-based magnitude and sensitivity. Such combinations are intuitively motivated by stochasticity in training dynamics, but they are not derived from an explicit statistical model, an optimization criterion, or a sampling distribution for the component signals. In particular, they do not provide p-values, error control, or a formal criterion for determining whether an observed component signal is distinguishable from stochastic optimization noise. Bayesian approaches provide another route to uncertainty-aware fine-tuning by placing prior distributions on model parameters and using posterior inference to quantify uncertainty [17], [18]. Recent Bayesian variants of LoRA adapt this idea to parameter-efficient fine-tuning by treating LoRA parameters or low-rank adaptation components as random variables and approximating their posterior distributions [19], [20]. However, these methods typically modify the training procedure by introducing priors, variational approximations, and posterior sampling, which may increase computational cost and alter the optimization dynamics of standard LoRA fine-tuning. These limitations motivate the development of a framework that preserves the standard training procedure while providing statistical decisions for rank allocation.

1.1 Overview: Statistical Inference for Rank Allocation↩︎

In this paper, we address LoRA rank allocation through a statistical inference framework. The central question is to determine which rank-one LoRA components provide sufficient statistical evidence of contribution to downstream adaptation. Let \(s_{\ell,j}^{\star}\) denote the importance score of the \(j\)-th component in the \(\ell\)-th module, where \(\ell=1,\cdots,L\) and \(j=1,\cdots,r\). This score measures the contribution of the corresponding rank-one update, with larger values indicating stronger evidence for retaining the component. In practice, \(s_{\ell,j}^{\star}\) is not directly observed and must be estimated from the stochastic training trajectory. Since LoRA fine-tuning involves mini-batch sampling, momentum, and adaptive step sizes, the resulting empirical component scores are random quantities generated by the optimizer dynamics. This observation motivates a statistical formulation of rank allocation. Instead of treating empirical component scores as deterministic measures of importance, we view them as noisy estimates of underlying adaptation signals and make retain-or-prune decisions through hypothesis testing. Specifically, for a pre-specified threshold \(\Delta > 0\), we consider the one-sided hypothesis testing problem \[\mathcal{H}_{\ell,j}^{0}:~s_{\ell,j}^{\star} \geq \Delta \qquad \text{versus} \qquad \mathcal{H}_{\ell,j}^{1}:~s_{\ell,j}^{\star} < \Delta .\] Here, the null hypothesis represents the event that the component has a sufficient contribution to be retained, whereas the alternative represents that its contribution falls below the threshold. Under this formulation, components with strong evidence against \(\mathcal{H}_{\ell,j}^{0}\) are removed first until the prescribed rank budget is reached.

The key step in implementing this testing procedure is to characterize the distribution of score estimates for components along the training trajectory. Motivated by asymptotic theory for stochastic approximation and stochastic gradient descent [21], [22], we establish central limit theory for adaptive optimizers commonly used in LoRA fine-tuning, including AdamW [23], Adam [24], and Adafactor [25]. In particular, for the optimizer state or iterate sequence \(\{\boldsymbol{z}_t\}\), we show that, under suitable regularity conditions, \[\gamma_t^{-1/2}\left(\boldsymbol{z}_t-\boldsymbol{z}^{\star}\right) \overset{d}{\longrightarrow} \mathcal{N}\left(\boldsymbol{0},\boldsymbol{\Sigma}\right),\] and the Polyak–Ruppert average trajectory satisfies \[\sqrt{t}\left(\frac{1}{t}\sum_{\tau=1}^{t}\boldsymbol{z}_{\tau}-\boldsymbol{z}^{\star}\right) \overset{d}{\longrightarrow} \mathcal{N}\left(\boldsymbol{0},\boldsymbol{\Omega}\right),\] where \(\boldsymbol{z}^{\star}\) denotes the limiting point of the optimizer dynamics, \(\{\gamma_t\}\) denotes the step-size sequence, and \(\boldsymbol{\Sigma},\boldsymbol{\Omega}\succeq \boldsymbol{0}\) are covariance matrices determined by the local curvature, stochastic gradient noise, and optimizer-specific dynamics.

We define the empirical component score \(s_{\ell,j}(t)\), which is an estimator for \(s_{\ell,j}^{\star}\), as a function of the current iterate \(\boldsymbol{z}_t\) at time step \(t\), then the delta method further yields the weak convergence \[\label{eq95clt95s} \sqrt{t}\left( \frac{1}{t}\sum_{\tau=1}^{t}s_{\ell,j}(\tau) -s_{\ell,j}^{\star} \right) \overset{d}{\longrightarrow} \mathcal{N}\left(0,\bar\sigma_{\ell,j}^{2}\right),\tag{1}\] where \(\bar\sigma_{\ell,j}^{2}\) depends on the covariance \(\boldsymbol{\Omega}\) and the form of the component score function. Based on 1 , we construct the test statistic in the boundary case \[T_{\ell,j}(t) = \frac{ \sqrt{t} \left( \frac{1}{t}\sum_{\tau=1}^{t}s_{\ell,j}(\tau)-\Delta \right) }{\widehat{\bar\sigma}_{\ell,j}},\] where \(\widehat{\bar\sigma}_{\ell,j}\) is estimated from the historical score trajectory using the batch-means estimator. The corresponding one-sided p-value is given by \[p_{\ell,j}(t) = \Phi\left(T_{\ell,j}(t)\right),\] where \(\Phi(\cdot)\) denotes the cumulative distribution function of the standard normal distribution. A small p-value provides evidence against the null hypothesis, indicating that the component has insufficient statistical evidence for retention relative to the threshold \(\Delta\).

This leads to our proposed method, StatLoRA, which integrates statistical inference with adaptive LoRA rank allocation. Compared with existing methods, StatLoRA provides a statistically interpretable decision rule for component selection while preserving the standard LoRA training procedure. In particular, it does not require modifying the training loss, changing the model parameterization, or altering the optimizer dynamics. Instead, it uses the stochastic trajectory and score quantities already generated during training to estimate the uncertainty of component contributions and guide retain-or-prune decisions.

1.2 Related Literature and Contributions↩︎

Our theoretical development is closely related to the asymptotic theory of stochastic approximation and stochastic optimization [26], [27]. Classical results on Polyak–Ruppert averaging establish asymptotic normality and optimal covariance properties for averaged stochastic approximation recursions, with averaged stochastic gradient descent as a central example [21], [28]. Subsequent studies have developed central limit theory and asymptotic distributions for stochastic quasi-Newton [29], non-smooth [30], constrained [31], and inexact [32] methods. These distributional results provide a basis for statistical inference from optimization trajectories, including uncertainty quantification [33]–[35], confidence intervals [36]–[38], hypothesis testing [39], [40], and convergence diagnostics [41], [42]. However, most existing statistical inference results focus on standard stochastic gradient methods and their averaged variants. Modern deep learning, such as LoRA fine-tuning, in contrast, is typically performed with adaptive optimizers such as AdamW [23], Adam [24], and Adafactor [25], whose momentum terms, coordinate-wise scaling, and adaptive learning rates lead to more complex stochastic dynamics. Our work extends this literature by establishing central limit theory for iterates generated by adaptive optimizers and by using the resulting asymptotic distributions to construct hypothesis tests for LoRA rank allocation.

The main contributions of this paper are summarized as follows:

• We formulate adaptive rank allocation for LoRA as a component-wise statistical hypothesis testing problem, providing a statistical inference perspective on LoRA resource allocation under stochastic training dynamics.

• We establish central limit theory for the trajectories of adaptive optimizers commonly used in deep learning and large-scale LLM fine-tuning, including AdamW, Adam, and Adafactor, under suitable regularity conditions. These results provide asymptotic distributions for optimizer states and, through the delta method, for empirical LoRA component scores.

• We propose StatLoRA, a practical inference-based rank allocation algorithm that uses the derived asymptotic distributions to quantify uncertainty in component scores. Under a prescribed total rank budget, StatLoRA retains components with sufficient statistical evidence for contribution and prunes components with weak evidence for retention.

• We conduct experiments on LoRA fine-tuning with DeBERTaV3-base, BART-Large, and Qwen2.5-7B across natural language understanding, natural language generation, and question answering tasks. The results show that StatLoRA achieves competitive performance compared with vanilla LoRA, AdaLoRA, and IGU-LoRA under matched rank budgets. We further conduct sensitivity analyses for the statistical hyperparameters, examine the stability of the rank allocation, and provide empirical diagnostics for the asymptotic normality of the empirical component scores.

1.3 Organization and Notation↩︎

The remainder of this paper is organized as follows. Section 2 introduces the background on LoRA fine-tuning and formulates LoRA rank allocation as a hypothesis testing problem on each component. Section 3 develops asymptotic theory for adaptive optimizers commonly used in deep learning, with particular emphasis on AdamW. Section 4 derives central limit results for empirical component scores and presents the proposed StatLoRA algorithm, including practical variance estimation and p-value construction. Section 5 reports the experimental results. Section 6 concludes the paper and discusses possible directions for future work.

Throughout this paper, we use lowercase letters, bold lowercase letters, and bold uppercase letters, such as \(a\), \(\boldsymbol{a}\), and \(\boldsymbol{A}\), to denote scalars, vectors, and matrices, respectively. The notation \(\|\cdot\|\) represents the Euclidean norm for vectors and the spectral norm for matrices, while \(\|\cdot\|_{\mathrm{F}}\) denotes the Frobenius norm. For two deterministic positive sequences \(\{a_t\}\) and \(\{b_t\}\), we write \(a_t=o(b_t)\) if \(a_t/b_t\to 0\), and \(a_t=\mathcal{O}(b_t)\) if \(a_t/b_t\) is bounded. For random sequences \(\{X_t\}\) and positive deterministic sequences \(\{a_t\}\), we write \(X_t=o_p(a_t)\) if \(X_t/a_t\to 0\) in probability, and \(X_t=\mathcal{O}_p(a_t)\) if \(X_t/a_t\) is bounded in probability. For vector- or matrix-valued quantities, these asymptotic notations are understood with respect to the norm \(\|\cdot\|\). For vectors \(\boldsymbol{a},\boldsymbol{b}\in\mathbb{R}^d\), we use \(\boldsymbol{a}\odot\boldsymbol{b}\) to denote the Hadamard product, i.e., element-wise multiplication. In particular, \(\boldsymbol{a}^{\odot 2}:=\boldsymbol{a}\odot\boldsymbol{a}\) denotes the element-wise square of \(\boldsymbol{a}\). Other vector operations, including \(\sqrt{\boldsymbol{a}}\), \(\boldsymbol{a}/\boldsymbol{b}\), and \(\boldsymbol{a}+\epsilon\), are understood entry-wise whenever they appear. Moreover, for a vector \(\boldsymbol{a}\in\mathbb{R}^d\), \(\operatorname{diag}(\boldsymbol{a})\in\mathbb{R}^{d\times d}\) denotes the diagonal matrix whose diagonal entries are given by the entries of \(\boldsymbol{a}\). When stochastic processes are involved, all random variables are defined on a filtered probability space \((\Omega,\mathcal{F},\{\mathcal{F}_t\}_{t\ge 0},\mathbb{P}),\) where \(\{\mathcal{F}_t\}_{t\ge 0}\) is an increasing sequence of sub-\(\sigma\)-algebras of \(\mathcal{F}\). Here, \(\mathcal{F}_t\) represents the information available before the stochastic update at iteration \(t\). The notation \(\mathbb{E}[\cdot\mid\mathcal{F}_t]\) denotes conditional expectation with respect to \(\mathcal{F}_t\).

2 LoRA Rank Allocation as a Statistical Decision Problem↩︎

In this section, we first review the basic formulation of LoRA and the associated rank allocation problem. We then use the rank-one decomposition of LoRA updates to show that rank allocation can be represented as a component selection problem. Unlike existing methods that primarily rely on various designs of importance measures, we formulate component selection as a statistical decision problem and introduce a hypothesis testing framework for LoRA rank allocation.

2.1 Low-Rank Adaptation and Rank Allocation↩︎

In transfer learning, a fine-tuned parameter matrix \(\boldsymbol{W}_{\text{ft}}\) is obtained by adapting a pretrained parameter matrix \(\boldsymbol{W}_{\text{pt}}\): \[\boldsymbol{W}_{\text{ft}} = \boldsymbol{W}_{\text{pt}} + \Delta\boldsymbol{W},\] where \(\boldsymbol{W}_{\text{ft}}, \boldsymbol{W}_{\text{pt}},\Delta\boldsymbol{W}\in \mathbb{R}^{n_1 \times n_2}\). Standard full fine-tuning directly updates the entire adaptation matrix \(\Delta\boldsymbol{W}\), providing full flexibility on the adaptation matrix \(\Delta\boldsymbol{W}\), but incurring substantial computational, memory, and storage costs. Low-rank adaptation (LoRA) [11] addresses this issue by imposing a low-rank structure on the adaptation matrix. Specifically, the update is parameterized as \[\Delta\boldsymbol{W}= \boldsymbol{B}\boldsymbol{A},\] where \(\boldsymbol{A}\in \mathbb{R}^{r \times n_2}\), \(\boldsymbol{B}\in \mathbb{R}^{n_1 \times r}\), and \(r \ll \min\{n_1, n_2\}\). Instead of training a full matrix of size \(n_1 n_2\), LoRA only updates the factor matrices \(\boldsymbol{A}\) and \(\boldsymbol{B}\), reducing the number of trainable parameters to \(r(n_1 + n_2)\). Beyond computational efficiency, the low-rank parameterization restricts adaptation to a lower-dimensional subspace, which can serve as an implicit regularization mechanism and improve the stability and generalization of fine-tuning [43]–[45].

For LoRA-based fine-tuning of a multi-layer neural network \(\phi(\cdot;\Theta)\), such as a transformer-based large language model [46], let \(\Theta := \{\boldsymbol{W}_1, \boldsymbol{W}_2, \cdots, \boldsymbol{W}_{L}\}\) denote the collection of parameter matrices to which LoRA is applied. For the \(\ell\)-th parameter matrix, LoRA takes the form \[\label{eq95lora} \boldsymbol{W}_{\text{ft},\ell} = \boldsymbol{W}_{\text{pt},\ell} + \Delta\boldsymbol{W}_{\ell} = \boldsymbol{W}_{\text{pt},\ell} + \boldsymbol{B}_{\ell} \boldsymbol{A}_{\ell},\tag{2}\] where \(\boldsymbol{W}_{\text{ft},\ell}\), \(\boldsymbol{W}_{\text{pt},\ell}\), and \(\Delta\boldsymbol{W}_{\ell} \in \mathbb{R}^{n \times n}\) denote the fine-tuned parameter matrix, the pretrained parameter matrix, and the adaptation matrix, respectively. Here \(\boldsymbol{A}_{\ell} \in \mathbb{R}^{r \times n}\) and \(\boldsymbol{B}_{\ell} \in \mathbb{R}^{n \times r}\) are the LoRA factor matrices associated with the \(\ell\)-th parameter matrix. For notational simplicity, we present the formulation for square matrices of size \(n \times n\), and the same discussion extends directly to rectangular matrices.

Since different layers and modules may play different roles in the neural network, a uniform rank allocation across all parameter matrices may be inefficient [12]. Some modules may encode general-purpose knowledge that should be largely preserved, whereas others may require stronger adaptation to capture task-specific or domain-specific information. Consequently, a small rank may be sufficient for some parameter matrices, while a larger rank may be needed for others. On the other hand, assigning unnecessarily large ranks can increase computational cost and may introduce excessive flexibility, potentially harming stability and generalization. In many standard LoRA implementations, these ranks are treated as fixed hyperparameters, which may lead to suboptimal performance. This motivates rank allocation, which allows different layers or modules to use different ranks rather than imposing a uniform rank across the model [13], [14].

Several methods have been proposed for automatic rank allocation in LoRA fine-tuning [13], [15], [16]. AdaLoRA [12] is a pioneering work in this direction, introducing a budget allocation procedure based on importance scores for rank-one components in the LoRA update \(\boldsymbol{B}_{\ell}\boldsymbol{A}_{\ell}\). Its score construction uses gradient-derived sensitivity and magnitude information to identify components that are likely to contribute more to model adaptation. More recently, IGU-LoRA [14] further refines component scoring by incorporating integrated-gradient information within the layer and uncertainty-aware adjustments, with the goal of stabilizing importance estimation under noisy training dynamics. Although these methods are empirically useful, their importance scores are typically constructed using a carefully designed combination of gradient magnitude and sensitivity. Such scoring rules are intuitively motivated by stochasticity in training dynamics, but they are not derived from an explicit statistical model, a statistical inference criterion, or a sampling distribution for component signals. As a result, it remains unclear how to interpret the numerical value of such scores, how to compare them across layers and modules, or how to distinguish genuine adaptation signals from optimizer-induced stochastic noise.

2.2 Statistical Formulation of LoRA Rank Allocation↩︎

We now formulate the LoRA rank allocation as a statistical decision problem. Suppose that, for each layer or module \(\ell\), LoRA is initialized with a relatively large rank \(r\). The adaptation matrix can then be decomposed into a sum of rank-one components: \[\Delta\boldsymbol{W}_{\ell} = \boldsymbol{B}_{\ell} \boldsymbol{A}_{\ell} = \sum_{j=1}^{r} \boldsymbol{b}_{\ell,j} \boldsymbol{a}_{\ell,j}^{\top},\] where \(\boldsymbol{b}_{\ell,j} \in \mathbb{R}^{n}\) denotes the \(j\)-th column of \(\boldsymbol{B}_{\ell}\), and \(\boldsymbol{a}_{\ell,j}^{\top} \in \mathbb{R}^{1 \times n}\) denotes the \(j\)-th row of \(\boldsymbol{A}_{\ell}\). Thus, each product \(\boldsymbol{b}_{\ell,j} \boldsymbol{a}_{\ell,j}^{\top}\) represents a rank-one adaptation component.

Under this decomposition, rank allocation can be viewed as selecting which rank-one components should be retained or removed. Let \(r_{\ell} \leq r\) denote the number of retained components in the \(\ell\)-th module. Given a total rank budget \(R\), the selected ranks are required to satisfy \[\sum_{\ell=1}^{L} r_{\ell} \leq R.\] Equivalently, among the \(rL\) candidate rank-one components \(\boldsymbol{b}_{\ell,j} \boldsymbol{a}_{\ell,j}^{\top}\), for \(\ell=1,\cdots,L\) and \(j = 1, \cdots, r\), one needs to select at most \(R\) components to retain while removing the rest. Exhaustively searching over all possible subsets is computationally infeasible, since the search space grows combinatorially with \(rL\). Moreover, evaluating the downstream performance of many candidate allocations would require repeated fine-tuning or validation, which is computationally expensive for large language models.

To avoid such an exhaustive search, we propose to make component-wise decisions through statistical inference. Let \[\Delta\boldsymbol{W}_{\ell,j}^{\star} = \boldsymbol{b}_{\ell,j}^{\star}\boldsymbol{a}_{\ell,j}^{\star\top}\] denote the \(j\)-th rank-one component in the \(\ell\)-th module. Here, the superscript \(\star\) refers to the limiting value of the corresponding LoRA components under the stochastic fine-tuning dynamics, rather than to an observed finite-iteration estimate. In this sense, \(\Delta\boldsymbol{W}_{\ell,j}^{\star}\) plays the role of the population-level target component in the LoRA framework. We measure the strength of this component and define the population score by the squared Frobenius norm \[\label{eq95score95population} s_{\ell,j}^{\star} = \left\|\Delta\boldsymbol{W}_{\ell,j}^{\star} \right\|_{\mathrm{F}}^{2} = \left\|\boldsymbol{b}_{\ell,j}^{\star}\boldsymbol{a}_{\ell,j}^{\star\top} \right\|_{\mathrm{F}}^{2} = \left\| \boldsymbol{b}_{\ell,j}^{\star}\right\|^{2} \cdot \left\| \boldsymbol{a}_{\ell,j}^{\star}\right\|^{2}.\tag{3}\] The quantity \(s_{\ell,j}^{\star}\) represents the underlying adaptation signal of the corresponding component. In practice, however, \(s_{\ell,j}^{\star}\) is not directly observed. We only observe empirical estimates generated along the stochastic training trajectory, which are affected by mini-batch sampling, optimizer momentum, and adaptive step sizes.

We therefore formulate component selection through the following one-sided hypothesis test. For a pre-specified threshold \(\Delta>0\), consider \[\label{hypothesis} \mathcal{H}_{\ell,j}^{0}:~s_{\ell,j}^{\star} \geq \Delta \qquad \text{versus} \qquad \mathcal{H}_{\ell,j}^{1}:~s_{\ell,j}^{\star} < \Delta .\tag{4}\] Here, the null hypothesis represents that the component has a sufficiently large adaptation signal and should be retained, while the alternative represents that its contribution falls below the threshold. This choice of hypotheses is aligned with the pruning objective: a small p-value provides evidence against \(\mathcal{H}_{\ell,j}^{0}\), suggesting that the corresponding component has insufficient statistical evidence for retention relative to \(\Delta\). Given the total rank budget \(R\), we compute p-values for all candidate components and prune those with the strongest evidence against the null hypothesis until the budget constraint is satisfied.

The remaining key ingredient is the distribution of the empirical component scores, which will be defined later. Since the exact scores \(s_{\ell,j}^{\star}\) are inaccessible, they must be estimated from the iterates generated during LoRA fine-tuning. This is nontrivial because the training trajectory is a complex stochastic process affected by mini-batch sampling, momentum, coordinate-wise scaling, and adaptive step sizes. As a result, the empirical component scores, being functions of these iterates, inherit the uncertainty and temporal dependence induced by the training dynamics. If this uncertainty can be characterized through an asymptotic distribution of the empirical score process, then the hypothesis test in 4 can be implemented using estimated p-values. Thus, the central theoretical question is to characterize the asymptotic distribution of iterates generated by commonly used adaptive optimizers in LoRA fine-tuning, and then to transfer this distributional result to component score estimates. We address this question in the next section by first establishing central limit theory for stochastic optimizer trajectories and then applying the delta method to obtain asymptotic distributions for empirical component scores.

3 Asymptotic Theory for Stochastic Optimization↩︎

To implement the hypothesis testing framework in 4 , we need to characterize the distributional behavior of the quantities generated during LoRA fine-tuning. In this section, we study the asymptotic distribution of iterates produced by stochastic optimization algorithms commonly used in LoRA fine-tuning of LLMs, including AdamW, Adam, and Adafactor. We first develop a general central limit theory for stochastic approximation recursions. We then show that these adaptive optimizers can be represented within this framework by introducing suitable augmented state variables and optimizer-specific mean fields. These results provide the distributional foundation for the component score inference developed in the next section.

3.1 General Stochastic Recursion↩︎

We consider the stochastic optimization problem \[\label{eq0} \min_{\boldsymbol{x}\in \mathbb{R}^d} f(\boldsymbol{x}) := \mathbb{E}_{\zeta \sim \mathcal{P}}\left[\mathcal{L}(\boldsymbol{x};\zeta)\right],\tag{5}\] where \(\zeta\) denotes a random observation or mini-batch drawn from a sampling distribution \(\mathcal{P}\). The loss function \(\mathcal{L}(\boldsymbol{x};\zeta)\) is evaluated at the given sample \(\zeta\), and \(f:\mathbb{R}^{d}\to\mathbb{R}\) is the population objective, assumed to be continuously differentiable. The formulation in 5 covers a broad class of problems arising in statistical learning and machine learning. In deep learning, \(\boldsymbol{x}\) denotes the collection of trainable parameters of a neural network, \(\mathcal{L}\left(\boldsymbol{x};\zeta\right)\) is the loss function evaluated at mini-batches of samples, and \(f(\boldsymbol{x})\) is the corresponding population risk or expected training loss [47]. In statistical learning, \(\boldsymbol{x}\) may represent the parameters of a statistical model, \(\mathcal{L}\left(\boldsymbol{x};\zeta\right)\) may correspond to a negative log-likelihood or a regularized risk evaluated at a data sample, and \(f(\boldsymbol{x})\) is the population objective associated with an \(M\)-estimator or expected risk minimization [48].

We use \(\boldsymbol{z}_t\) to denote the algorithmic state at iteration \(t\) for a stochastic optimization algorithm applied to 5 . The parameter vector \(\boldsymbol{x}_t\) can be recovered from \(\boldsymbol{z}_t\), while \(\boldsymbol{z}_t\) may also include auxiliary variables needed to describe the algorithmic dynamics, such as momentum or adaptive second-moment estimates. We consider the following general stochastic approximation recursion: \[\label{eq95general95recursion} \boldsymbol{z}_{t+1} = \boldsymbol{z}_t + \gamma_t \left( F(\boldsymbol{z}_t) + \boldsymbol{\xi}_{t} + \boldsymbol{r}_{t} \right),\tag{6}\] where \(t \in \mathbb{N}\), \(F\) represents the mean field (or drift) of the algorithm, \(\boldsymbol{\xi}_{t}\) denotes the stochastic noise in the update direction, and \(\boldsymbol{r}_t\) represents a remainder term, accounting for approximation errors or additional algorithmic effects not captured by the leading mean field \(F\). The step-size sequence is chosen as \[\label{eq95step95size} \gamma_t = \gamma_0 t^{-\kappa}, \quad \kappa \in (1/2,1),\tag{7}\] where \(\gamma_0 > 0\). This corresponds to a polynomially decaying (or power-law) step size schedule.

To characterize the uncertainty and asymptotic behavior of the stochastic recursion 6 under the step-size schedule 7 , we impose the following regularity assumptions on the mean dynamics, stochastic noise, and remainder terms.

Remarks. Assumption 1 collects standard regularity conditions for deriving asymptotic normality of stochastic approximation recursions. Their roles are summarized as follows:

1. The convergence condition localizes the analysis around a limiting point \(\boldsymbol{z}^\star\), so that the local distributional behavior of the iterates can be studied after convergence.

2. The differentiability condition justifies a local linear approximation of the mean field \(F\) near \(\boldsymbol{z}^\star\). The Jacobian \(\boldsymbol{J}=\nabla F(\boldsymbol{z}^\star)\) determines the leading deterministic dynamics.

3. The martingale-difference and moment conditions control the stochastic noise. They are standard requirements for martingale central limit arguments and tail control.

4. The covariance condition ensures that the conditional covariance of the noise has a well-defined asymptotic limit, which determines the noise level in the limiting distribution.

5. The stability condition requires the linearized mean dynamics to be locally attracting, ensuring that the deterministic drift pulls the iterates toward the limiting point.

6. The remainder condition ensures that remainder terms are asymptotically negligible and do not affect the leading mean field and the limiting distribution.

Overall, these assumptions are standard in central limit theory for stochastic approximation and stochastic optimization [30]–[32], [37], [39].

Proof sketch. The detailed proof is provided in Appendix 7.3. We briefly summarize the main argument. Let \(\boldsymbol{e}_t=\boldsymbol{z}_t-\boldsymbol{z}^\star\). By the local expansion of \(F\) around \(\boldsymbol{z}^\star\), the recursion can be written as \[\boldsymbol{e}_{t+1} = (\boldsymbol{I}+\gamma_t\boldsymbol{J})\boldsymbol{e}_t +\gamma_t\boldsymbol{\xi}_t +\underbrace{\gamma_t \mathcal{O}(\|\boldsymbol{e}_t\|^2)}_{\text{Taylor remainder}} +\underbrace{\gamma_t\boldsymbol{r}_t}_{\text{original remainder}}.\] Under Assumptions (A1)–(A6), the Taylor remainder term and the original remainder term \(\boldsymbol{r}_t\) are asymptotically negligible. Hence, the leading term is governed by a stable linear stochastic recursion driven by martingale-difference noise. Applying a martingale central limit theorem to the resulting weighted martingale array gives the asymptotic normality of \(\gamma_t^{-1/2}(\boldsymbol{z}_t-\boldsymbol{z}^\star)\), with covariance characterized by the Lyapunov equation (see Lemma 4). For the Polyak–Ruppert average, summing the linearized recursion results in the asymptotic representation \[\sqrt{t}(\bar{\boldsymbol{z}}_t-\boldsymbol{z}^\star) = -\boldsymbol{J}^{-1}\frac{1}{\sqrt{t}}\sum_{\tau=1}^{t}\boldsymbol{\xi}_\tau +o_p(1).\] The martingale central limit theorem, together with Slutsky’s theorem, then implies the stated averaged central limit theorem.

1 provides the limiting distributions for the stochastic recursion 6 . Such results are classical in stochastic approximation theory and are closely related to Polyak–Ruppert averaging and asymptotic normality for stochastic recursive algorithms [21], [22], [27]. The theorem shows that, under standard local regularity and stability conditions, both the last iterate and the Polyak–Ruppert averaged iterate admit Gaussian limits after appropriate normalization. These results provide the theoretical foundation for studying the local asymptotic behavior of stochastic optimization algorithms. In the remainder of this section, we show how several adaptive optimizers used in LoRA fine-tuning can be represented in the form of 6 , and derive their corresponding asymptotic distributions.

3.2 Central Limit Theory for AdamW↩︎

We next specialize the general stochastic approximation framework to AdamW, a widely used adaptive optimizer in deep learning and LoRA fine-tuning [23]. Consider applying AdamW to the stochastic optimization problem in 5 , where the population gradient is observed through stochastic gradient estimates. The updates are given by \[\label{eq95adamw} \begin{align} \boldsymbol{m}_{t+1} &= (1-\alpha_t)\boldsymbol{m}_t + \alpha_t \boldsymbol{g}_t,\\ \boldsymbol{v}_{t+1} &= (1-\beta_t)\boldsymbol{v}_t + \beta_t \boldsymbol{g}_t^{\odot 2},\\ \boldsymbol{x}_{t+1} &= \boldsymbol{x}_t - \gamma_t\frac{\boldsymbol{m}_{t+1}}{\sqrt{\boldsymbol{v}_{t+1}}+\varepsilon} - \gamma_t\lambda \boldsymbol{x}_t, \end{align}\tag{8}\] where \(\boldsymbol{g}_t:=\nabla \mathcal{L}(\boldsymbol{x}_t;\zeta_t)\) denotes the stochastic gradient evaluated at \(\boldsymbol{x}_t\) with data sample \(\zeta_t\), \(\alpha_t=\alpha\gamma_t\) and \(\beta_t=\beta\gamma_t\) are step sizes for momentum terms with \(\alpha, \beta>0\), \(\lambda>0\) is the weight decay parameter, \(\gamma_t\) is the step size as defined in 7 , and \(\varepsilon>0\) is a numerical stabilization constant. Here, \(\boldsymbol{m}_t\) and \(\boldsymbol{v}_t\) denote the first- and second-moment variables, respectively. All vector operations in 8 , including division and square root, are understood element-wise.

We define the augmented state variable \[\label{eq95augmented95z} \boldsymbol{z}_t := (\boldsymbol{x}_t,\boldsymbol{m}_t,\boldsymbol{v}_t)\in\mathbb{R}^{3d}.\tag{9}\] Let \(\nabla f(\boldsymbol{x}_t)\) denote the population gradient at \(\boldsymbol{x}_t\). The stochastic gradient is written as \[\boldsymbol{g}_t = \nabla \mathcal{L}(\boldsymbol{x}_{t}; \zeta_t) = \nabla f(\boldsymbol{x}_t)+\boldsymbol{\xi}_t^{(g)},\] where \(\boldsymbol{\xi}_t^{(g)} := \nabla \mathcal{L}(\boldsymbol{x}_{t}; \zeta_t) - \nabla f(\boldsymbol{x}_t)\) is the stochastic gradient noise. We assume that the stochastic gradient is conditionally unbiased with respect to the pre-update filtration \(\mathcal{F}_t\), that is, \[\mathbb{E}\left[\boldsymbol{\xi}_t^{(g)}\mid\mathcal{F}_t\right]=\boldsymbol{0}, \qquad \mathbb{E}\left[\nabla \mathcal{L}(\boldsymbol{x}_t;\zeta_t)\mid\mathcal{F}_t\right] = \nabla f(\boldsymbol{x}_t).\] This corresponds to the standard mini-batch setting in stochastic optimization, where the mini-batch is sampled conditionally independently given the current iterate, such as neural network training and fine-tuning [47]. We define the second-moment function of the stochastic gradient by \[\label{eq95q} \boldsymbol{q}(\boldsymbol{x}) := \mathbb{E}\left[\nabla \mathcal{L}(\boldsymbol{x};\zeta)^{\odot 2}\right].\tag{10}\] Then the first- and second-moment updates in 8 can be written as \[\begin{align} & \boldsymbol{m}_{t+1} = \boldsymbol{m}_t+\alpha_t\bigl(\nabla f(\boldsymbol{x}_t)-\boldsymbol{m}_t\bigr)+\alpha_t\boldsymbol{\xi}_t^{(g)},\\ & \boldsymbol{v}_{t+1} = \boldsymbol{v}_t + \beta_t\big(\boldsymbol{q}(\boldsymbol{x}_t) - \boldsymbol{v}_t\big) + \beta_t\left(2\nabla f(\boldsymbol{x}_t)\odot \boldsymbol{\xi}_{t}^{(g)} + \boldsymbol{\xi}_{t}^{(g)\odot 2} - \mathbb{E}\left[\boldsymbol{\xi}_t^{(g)\odot 2}\mid \mathcal{F}_{t}\right]\right). \end{align}\] The parameter update can be rewritten as \[\boldsymbol{x}_{t+1} = \boldsymbol{x}_t + \gamma_t\left( -\frac{\boldsymbol{m}_t}{\sqrt{\boldsymbol{v}_t}+\varepsilon} -\lambda \boldsymbol{x}_t \right) + \gamma_t \boldsymbol{r}_t^{(x)},\] where \[\boldsymbol{r}_t^{(x)} = \frac{\boldsymbol{m}_t}{\sqrt{\boldsymbol{v}_t}+\varepsilon} - \frac{\boldsymbol{m}_{t+1}}{\sqrt{\boldsymbol{v}_{t+1}}+\varepsilon}\] collects the correction caused by replacing \((\boldsymbol{m}_{t+1},\boldsymbol{v}_{t+1})\) with \((\boldsymbol{m}_t,\boldsymbol{v}_t)\) in the leading drift.

Therefore, AdamW can be represented in the form of the general stochastic recursion 6 . Specifically, for \(\boldsymbol{z}=(\boldsymbol{x},\boldsymbol{m},\boldsymbol{v})\), the mean-field mapping \(F:\mathbb{R}^{3d}\to\mathbb{R}^{3d}\) is \[\label{eq95mean95field95adamw} F(\boldsymbol{z}) = \begin{pmatrix} -\dfrac{\boldsymbol{m}}{\sqrt{\boldsymbol{v}}+\varepsilon}-\lambda \boldsymbol{x}\\[2mm] \alpha\bigl(\nabla f(\boldsymbol{x})-\boldsymbol{m}\bigr)\\[1mm] \beta\bigl(\boldsymbol{q}(\boldsymbol{x})-\boldsymbol{v}\bigr) \end{pmatrix}.\tag{11}\] The martingale noise is given by \[\label{eq95noise95adamw} \boldsymbol{\xi}_t= \begin{pmatrix} \boldsymbol{0}\\ \alpha\boldsymbol{\xi}_t^{(g)}\\ \beta\Bigl( 2\nabla f(\boldsymbol{x}_t)\odot \boldsymbol{\xi}_t^{(g)} + \hat{\boldsymbol{\xi}}_t^{(g)} \Bigr) \end{pmatrix},\tag{12}\] where \[\label{eq95second95noise} \hat{\boldsymbol{\xi}}_t^{(g)} = \boldsymbol{\xi}_t^{(g)\odot 2} - \mathbb{E}\!\left[\boldsymbol{\xi}_t^{(g)\odot 2}\mid \mathcal{F}_t\right],\tag{13}\] and the remainder term is \[\label{eq95remainder95adamw} \boldsymbol{r}_t= \begin{pmatrix} \boldsymbol{r}_t^{(x)}\\ \boldsymbol{0}\\ \boldsymbol{0} \end{pmatrix}.\tag{14}\] Consequently, the AdamW iteration admits the stochastic approximation representation 6 , with mean field, martingale noise, and remainder term given by 11 , 12 , and 14 , respectively.

The following proposition verifies the convergence condition and local stability requirement in Assumption 1 for the AdamW mean-field dynamics. Due to the decoupled weight decay term in 8 , the limiting point of AdamW does not necessarily coincide with a stationary point of the population objective \(f\), which differs from optimizers without decoupled weight decay. We further characterize the Jacobian of the AdamW mean field at the limiting point and provide a sufficient condition under which the local stability condition in Assumption (A5) holds. The proof is available in Appendix 10.3.

We now specialize Assumption 1 to the AdamW recursion in 8 . The following conditions provide optimizer-specific regularity assumptions under which AdamW fits into the general stochastic approximation framework.

Remarks. The above conditions are sufficient to verify Assumption 1 for the AdamW recursion defined by 11 , 12 , and 14 . Their roles are analogous to those discussed after Assumption 1, and we therefore do not repeat the full discussion here. We only highlight two AdamW-specific points. First, the bounded \((4+\delta)\)-moment condition on the stochastic gradient noise is required because the second-moment update in 8 involves squared stochastic gradients. Second, the conditional covariance assumptions in (B4) require the noise structure, including the second-order gradient noise, to stabilize asymptotically. These conditions allow the local behavior of the AdamW iterates to be characterized through the central limit theory for the augmented stochastic recursion.

Proof sketch. The detailed proof is deferred to Appendix 7.4. The proof proceeds by verifying that the AdamW recursion can be written in the general stochastic approximation form 6 . Under Assumption 2, the augmented state \(\boldsymbol{z}_t=(\boldsymbol{x}_t,\boldsymbol{m}_t,\boldsymbol{v}_t)\) converges to a limiting point, the AdamW mean field is locally differentiable, the induced noise sequence is a martingale difference with a limiting conditional covariance, and the remainder term is asymptotically negligible. Hence, Assumptions (A1)–(A6) in Assumption 1 hold for the AdamW recursion. The last-iterate and Polyak–Ruppert averaged central limit theorems then follow directly from Theorem 1. The block covariance matrix \(\boldsymbol{Q}\) is obtained by computing the conditional covariance of the martingale noise vector in 12 .

Theorem 2 gives an explicit asymptotic distribution for the AdamW iterates and their Polyak–Ruppert averages. Compared with standard stochastic gradient methods, the analysis is more involved because AdamW contains first-moment averaging, second-moment adaptation, coordinate-wise scaling with momentum terms, and decoupled weight decay. These algorithmic components require an augmented-state formulation and lead to a nontrivial covariance structure. To the best of our knowledge, few existing results provide such an explicit central limit theorem for AdamW with both last-iterate and averaged-iterate limits. This result forms the main theoretical basis for the component score inference and hypothesis testing procedure developed in the next section.

3.3 Extensions to Adam and Adafactor↩︎

Adam [24] and Adafactor [25] are also commonly used optimizers in LoRA fine-tuning and large-scale model training. Both algorithms can be analyzed within the same stochastic approximation framework developed above, after introducing suitable augmented state variables and optimizer-specific mean-field mappings. Compared with AdamW, the main differences lie in the form of weight decay and the construction of the adaptive preconditioner. These differences lead to different limiting points, Jacobian matrices, and limiting covariance structures, but the underlying proof strategy remains similar: one verifies the corresponding stochastic approximation assumptions and then applies the general central limit theorem in Theorem 1.

To keep the main content focused, we defer the detailed formulations and proofs for Adam and Adafactor to Appendices 8 and 9. For completeness, we state the following general extension result.

The precise augmented states, mean-field mappings, Jacobian matrices, limiting covariance matrices, and proofs are provided in Appendices 8 and 9. The resulting covariance formulas differ across optimizers because they depend on the local linearization of the corresponding mean field and on the conditional covariance of the induced martingale noise.

The unified formulation also suggests that the same analysis can be extended to some other optimizers in deep learning, provided that they admit an augmented-state representation satisfying the regularity conditions in Assumption 1. Recent work has studied central limit theory for Adam-type methods using optimizer-specific arguments [49]–[51]. In contrast, our approach emphasizes a common stochastic approximation structure, which allows AdamW, Adam, and Adafactor to be treated within a unified framework.

4 Hypothesis Testing for LoRA Rank Allocation↩︎

In this section, we develop the statistical methodology and algorithm for LoRA rank allocation. The central objects of inference are the empirical component scores generated along stochastic optimizer trajectories. Building on the central limit theory for optimizer iterates established in the previous section, we first transfer the asymptotic analysis from iterates to LoRA component scores and derive the limiting distribution of averaged empirical scores. We then introduce a practical variance estimation procedure to approximate the corresponding long-run variance from the observed score trajectory. These ingredients allow us to construct one-sided hypothesis tests for evaluating the significance of individual rank-one components. Finally, we incorporate the resulting p-values into an adaptive pruning procedure, leading to the proposed StatLoRA algorithm for statistically guided LoRA rank allocation.

4.1 Central Limit Theory for Empirical Scores↩︎

The population component scores \(s_{\ell,j}^{\star}\), defined in 3 , are unknown in practice. We therefore estimate them from the training trajectory generated by stochastic optimization. For the \(j\)-th rank-one component in the \(\ell\)-th LoRA module, define the empirical score at iteration \(t\) by \[\label{eq95empirical95score} s_{\ell,j}(t) = \left\| \boldsymbol{b}_{\ell,j}(t)\right\|^2 \cdot \left\| \boldsymbol{a}_{\ell,j}(t)\right\|^2,\tag{15}\] where \(\boldsymbol{a}_{\ell,j}(t)\) and \(\boldsymbol{b}_{\ell,j}(t)\) denote the corresponding LoRA factors at iteration \(t\), after running the stochastic optimization algorithms (e.g., AdamW) for fine-tuning. In the LoRA fine-tuning setting, the state variable \(\boldsymbol{z}\) contains all trainable LoRA parameters \(\{\boldsymbol{A}_{\ell},\boldsymbol{B}_{\ell}\}_{\ell=1}^{L}\), possibly together with optimizer-specific auxiliary variables. Let \(\boldsymbol{P}_{\ell,j}\) denote the selection matrices that extract \(\boldsymbol{w}_{\ell,j} := (\boldsymbol{a}_{\ell,j},\boldsymbol{b}_{\ell,j})\) from the full state vector \(\boldsymbol{z}\), respectively. That is, \[\boldsymbol{w}_{\ell,j}(t)=\boldsymbol{P}_{\ell,j}\boldsymbol{z}_t .\] Define the induced joint covariance matrices \[\label{eq95cov95component} \boldsymbol{\Sigma}_{\ell,j}^{w}:=\boldsymbol{P}_{\ell,j}\boldsymbol{\Sigma}\boldsymbol{P}_{\ell,j}^{\top}, \qquad \bar{\boldsymbol{\Sigma}}_{\ell,j}^{w} := \boldsymbol{P}_{\ell,j}\boldsymbol{J}^{-1}\boldsymbol{Q}\boldsymbol{J}^{-\top}\boldsymbol{P}_{\ell,j}^{\top}.\tag{16}\] where \(\boldsymbol{\Sigma}\) is the limiting covariance matrix for last-iterate, while \(\boldsymbol{J}\) and \(\boldsymbol{Q}\) are Jacobian and noise covariance matrices appearing in 2. By the central limit theory for iterates established in 2, these component-wise LoRA factors inherit asymptotic normality from the full optimizer state. In particular, \[\label{eq95asymp95a} \gamma_{t}^{-1/2}\left(\boldsymbol{w}_{\ell,j}(t) - \boldsymbol{w}_{\ell,j}^{\star}\right) \;\xrightarrow{d}\; \mathcal{N}\left(\boldsymbol{0},\boldsymbol{\Sigma}_{\ell,j}^{w}\right),\tag{17}\] and \[\label{eq95asymp95avg95w} \sqrt{t}\left(\frac{1}{t}\sum_{\tau=1}^{t}\boldsymbol{w}_{\ell,j}(\tau) - \boldsymbol{w}_{\ell,j}^{\star}\right) \;\xrightarrow{d}\; \mathcal{N}\left(\boldsymbol{0},\bar{\boldsymbol{\Sigma}}_{\ell,j}^{w}\right),\tag{18}\] where \(\boldsymbol{w}^{\star}_{\ell,j} := (\boldsymbol{a}^{\star}_{\ell,j},\boldsymbol{b}^{\star}_{\ell,j})\). These relations show that the joint LoRA component vector inherits an asymptotic distribution from the full optimizer state. We next transfer these results to the empirical component score \(s_{\ell,j}(t)\).

Since the empirical score \(s_{\ell,j}(t)\) is a smooth function of the LoRA factors \(\boldsymbol{a}_{\ell,j}(t)\) and \(\boldsymbol{b}_{\ell,j}(t)\), the central limit theory for optimizer iterates can be transferred to the score process through the delta method. The following theorem characterizes the asymptotic distribution of both the last-iterate score and the averaged empirical score.

The first-order delta method may degenerate at the zero-score boundary \(s_{\ell,j}^{\star}=0\), where one of the two LoRA factors may vanish. Since our null hypothesis is formulated as \(s_{\ell,j}^{\star}\geq \Delta\) with \(\Delta>0\), the relevant testing regime is away from this degenerate boundary. In this regime, both \(\boldsymbol{a}_{\ell,j}^{\star}\) and \(\boldsymbol{b}_{\ell,j}^{\star}\) are nonzero, and the first-order Taylor expansion of the score function provides the leading asymptotic term.

To express the limiting variances compactly, define the score function \[h(\boldsymbol{w}_{\ell,j}) :=\|\boldsymbol{a}_{\ell,j}\|^{2}\|\boldsymbol{b}_{\ell,j}\|^{2}.\] Then the variances in 3 are given by \[\sigma_{\ell,j}^{2} = \nabla h(\boldsymbol{w}_{\ell,j}^{\star})^{\top} \boldsymbol{\Sigma}_{\ell,j}^{w} \nabla h(\boldsymbol{w}_{\ell,j}^{\star}), \qquad \bar{\sigma}_{\ell,j}^{2} = \nabla h(\boldsymbol{w}_{\ell,j}^{\star})^{\top} \bar{\boldsymbol{\Sigma}}_{\ell,j}^{w} \nabla h(\boldsymbol{w}_{\ell,j}^{\star}),\] where \[\nabla h(\boldsymbol{w}_{\ell,j}^{\star}) = \begin{pmatrix} 2\|\boldsymbol{b}_{\ell,j}^{\star}\|^{2}\boldsymbol{a}_{\ell,j}^{\star}\\ 2\|\boldsymbol{a}_{\ell,j}^{\star}\|^{2}\boldsymbol{b}_{\ell,j}^{\star} \end{pmatrix},\] and covariance matrices are defined in 16 .

Proof sketch. The detailed proof is deferred to Appendix 7.5. The proof follows from the delta method. Since \(s_{\ell,j}(t) = h(\boldsymbol{w}_{\ell,j}(t)),\) a first-order Taylor expansion around \(\boldsymbol{w}_{\ell,j}^{\star}\) yields \[s_{\ell,j}(t)-s_{\ell,j}^{\star} = \nabla h(\boldsymbol{w}_{\ell,j}^{\star})^{\top} \left( \boldsymbol{w}_{\ell,j}(t)-\boldsymbol{w}_{\ell,j}^{\star} \right) + o_p\left( \|\boldsymbol{w}_{\ell,j}(t)-\boldsymbol{w}_{\ell,j}^{\star}\| \right).\] The condition \(s_{\ell,j}^{\star}\geq\Delta>0\) excludes the zero-score boundary at which the first-order expansion may degenerate. Combining the above expansion with the component-wise central limit theory inherited from the optimizer state gives ?? . The averaged-score result ?? follows by averaging the Taylor expansion along the trajectory and showing that the averaged second-order remainder is \(o_p(t^{-1/2})\).

3 extends the optimizer-level central limit theory to the empirical component scores used for LoRA rank allocation. Thus, the empirical score is treated not merely as a deterministic importance measure, but as an estimator with an asymptotic distribution. This provides the basis for uncertainty quantification and for the hypothesis testing procedure developed below.

4.2 Hypothesis Testing for LoRA Components↩︎

The asymptotic distributions in 3 provide the basis for hypothesis testing for each LoRA component. Our goal is to use the empirical score trajectory to evaluate whether a population score \(s_{\ell,j}^{\star}\) is sufficiently large, and hence whether the corresponding LoRA component should be retained or pruned. To implement such tests, however, one must estimate the asymptotic variance appearing in the central limit theorem for empirical scores.

A direct approach is to use the explicit variance formula in 3. This leads to a plug-in estimator [32], [39], [52], where the optimizer-level covariance matrices, Jacobians, and score derivatives are estimated separately and substituted into the theoretical expression. Although conceptually straightforward, this approach is often impractical in large-scale fine-tuning, since it requires estimating high-dimensional covariance matrices and local derivative quantities that may be unstable or expensive to compute.

As an alternative, we use a batch-means estimator [53], [54] to estimate the long-run variance of the scalar score process directly from its observed trajectory. Batch-means estimators are standard tools for variance estimation in central limit theorems for dependent stochastic processes and are particularly natural for averaged estimators. This motivates us to base the testing procedure on the Polyak–Ruppert averaged empirical score rather than the last-iterate score. Specifically, we use the averaged-score central limit theorem \[\sqrt{t} \left( \frac{1}{t}\sum_{\tau=1}^{t}s_{\ell,j}(\tau) - s_{\ell,j}^{\star} \right) \xrightarrow{d} \mathcal{N} \left( 0, \bar{\sigma}_{\ell,j}^{2} \right),\] and estimate the long-run variance \(\bar{\sigma}_{\ell,j}^{2}\) from the score trajectory. The following proposition states the consistency of the batch-means estimator under standard regularity conditions.

2 provides a practical way to estimate the unknown long-run variance in the averaged-score central limit theorem. Instead of estimating the full optimizer-state covariance matrix and the derivatives appearing in the delta-method formula, the batch-means estimator uses only the observed scalar score trajectory. In our setting, temporal dependence arises because the empirical scores are generated along a single stochastic optimizer trajectory. Under the usual functional central limit theorem, moment, weak-dependence, and batch-size conditions for batch-means consistency, the estimator is consistent. We defer the detailed regularity conditions and proof to Appendix 10.4.

The conditions \(b\to\infty\), \(M\to\infty\), and \(b/t\to0\) are asymptotic requirements for consistency of the batch-means estimator. With the variance estimate \(\widehat{\bar{\sigma}}_{\ell,j}^{2}\), we define the test statistic \[T_{\ell,j}(t) = \frac{ \sqrt{t} \left( \bar{s}_{\ell,j}(t) - \Delta \right) }{ \widehat{\bar{\sigma}}_{\ell,j} }, \qquad \bar{s}_{\ell,j}(t) = \frac{1}{t} \sum_{\tau=1}^{t} s_{\ell,j}(\tau).\] Here, \(\widehat{\bar{\sigma}}_{\ell,j}\) denotes the batch-means estimate of \(\bar{\sigma}_{\ell,j}\). Under the null hypothesis \[\label{eq95hypothesis95testing} \mathcal{H}_{\ell,j}^{0}:~s_{\ell,j}^{\star}\geq \Delta \qquad \text{versus} \qquad \mathcal{H}_{\ell,j}^{1}:~s_{\ell,j}^{\star}<\Delta,\tag{19}\] the corresponding one-sided p-value is \[p_{\ell,j}(t) = \Phi\left(T_{\ell,j}(t)\right),\] where \(\Phi(\cdot)\) denotes the cumulative distribution function of the standard normal distribution. A small p-value provides evidence against the null hypothesis, indicating that the component has insufficient population score relative to the threshold \(\Delta\). Such components are therefore viewed as having weak statistical evidence for retention and are prioritized for pruning.

4.3 The StatLoRA Algorithm↩︎

We now describe the proposed StatLoRA algorithm, a statistical inference-based algorithm for LoRA rank allocation. StatLoRA starts from an over-parameterized LoRA model with an initial rank \(r\) for each selected module and gradually removes rank-one components that have weak statistical evidence for retention during fine-tuning. The method uses the empirical score trajectory of each component to estimate its signal strength and uncertainty, and then performs hypothesis testing based on the p-values.

2 provides a batch-means estimator for the long-run variance of the averaged empirical score. The consistency result is asymptotic and requires both the batch size \(b\) and the number of batches \(M\) to increase. In large-scale model fine-tuning, however, it is not practical to let \(b\) and \(M\) grow indefinitely. We therefore use a fixed batch size \(b\) in finite training runs, and the number of batches is determined by the current trajectory length, \(M=t/b\) for notational simplicity. This provides a computationally efficient finite-sample approximation to the long-run variance.

At a pruning step \(t\), let \(\mathcal{A}_t\) denote the set of active LoRA components, namely, the components that have not been pruned before iteration \(t\). For each \((\ell,j)\in\mathcal{A}_t\), the score trajectory is partitioned into \(M\) consecutive batches, each of length \(b\), so that \(t=Mb\). We prune a component when the null hypothesis in 19 is rejected, providing sufficient evidence that its limiting score is below the pre-specified threshold \(\Delta\). Specifically, define the batch means by

\[\bar{s}_{\ell,j}^{(m)}(t) = \frac{1}{b} \sum_{\tau=(m-1)b+1}^{mb} s_{\ell,j}(\tau), \qquad m=1,\ldots,M.\] Then, \[\label{eq95mean} \bar{s}_{\ell,j}(t) = \frac{1}{t} \sum_{\tau=1}^{t} s_{\ell,j}(\tau) = \frac{1}{M} \sum_{m=1}^{M} \bar{s}_{\ell,j}^{(m)}(t).\tag{20}\] Therefore, the long-run variance is estimated by \[\label{eq95var95est95batch95means} \widehat{\bar{\sigma}}_{\ell,j}^{2}(t) = \frac{b}{M-1} \sum_{m=1}^{M} \left( \bar{s}_{\ell,j}^{(m)}(t) - \bar{s}_{\ell,j}(t) \right)^2 .\tag{21}\]

The choice of the threshold \(\Delta\) in 19 is another important practical issue. Rather than fixing an absolute threshold, we choose it adaptively from the empirical score distribution. Therefore, we use the notation \(\Delta^{(t)}\) to denote the threshold in hypothesis testing at different time steps. Specifically, for a pre-specified quantile level \(q\in(0,1)\), we set \[\Delta^{(t)} = \operatorname{Quantile}_{q} \left( \{\bar{s}_{\ell,j}(t): (\ell,j)\in\mathcal{A}_t\} \right),\] where \(\mathcal{A}_t\) denotes the set of active components at iteration \(t\) and \(\operatorname{Quantile}_q\) denotes the empirical \(q\)-quantile of the finite set of averaged component scores. For example, \(q=0.3\) sets \(\Delta^{(t)}\) to the lower 30% empirical quantile of the averaged component scores among active components.

Given \(\bar{s}_{\ell,j}(t)\), \(\widehat{\bar{\sigma}}_{\ell,j}(t)\), and \(\Delta^{(t)}\), the algorithm uses the statistic: \[T_{\ell,j}(t) = \frac{ \sqrt{t}\left(\bar s_{\ell,j}(t)-\Delta^{(t)}\right) }{ \widehat{\bar\sigma}_{\ell,j}(t) }.\] The corresponding p-value is \[p_{\ell,j}(t)=\Phi(T_{\ell,j}(t)).\] Since the null hypothesis is \(\mathcal{H}^{0}_{\ell,j}:s_{\ell,j}^{\star}\geq \Delta^{(t)}\), a small p-value provides evidence that the component has insufficient population-level signal. StatLoRA therefore prunes components with the smallest p-values according to the prescribed budget schedule. The full algorithm is summarized in 1.

StatLoRA introduces little additional computational overhead. The component scores are computed from vector norms of existing LoRA factors, and the averaged scores, batch-means variances, test statistics, and p-values are scalar quantities updated during training. Thus, the dominant computational cost remains the standard LoRA fine-tuning procedure. In our experiments, the additional computation is small relative to the cost of model training.

5 Experiments↩︎

In this section, we evaluate the proposed StatLoRA method for LoRA fine-tuning across a range of language tasks, datasets, and pretrained models. The experiments are designed to validate both downstream task performance and the statistical behavior of the proposed rank allocation method. We first compare StatLoRA with existing LoRA rank allocation baselines on benchmark tasks, including natural language understanding, natural language generation, and question answering. We then conduct sensitivity analyses to examine the effect of the hyperparameters in StatLoRA and study the stability of rank allocation strategies. Finally, we provide empirical validations for the derived asymptotic normality for empirical scores.

5.1 Experimental Setups↩︎

We evaluate the proposed StatLoRA method on LoRA fine-tuning across three categories of language tasks: natural language understanding, natural language generation, and question answering. These experiments are designed to assess whether the proposed statistical rank allocation method is effective across different tasks, datasets, pretrained language models, and rank budgets. The task-specific datasets, evaluation metrics, and pretrained models are described in the corresponding subsections.

We compare StatLoRA with vanilla LoRA [11], AdaLoRA [12], and IGU-LoRA [14]. Vanilla LoRA uses a fixed rank allocation, where the same rank is assigned to each selected module. AdaLoRA adaptively allocates ranks by constructing importance scores for rank-one components based on gradient-derived sensitivity estimates. IGU-LoRA further builds on AdaLoRA by modifying the importance-score construction through gradient uncertainty and layer-wise information. These methods represent fixed-rank and adaptive-rank allocation strategies for LoRA, and therefore provide a natural comparison for the proposed hypothesis-testing-based allocation rule.

For fair comparison, we follow the training schedule used in AdaLoRA [12], consisting of an initial warm-up phase, a rank-reduction phase, and a final fine-tuning phase after pruning. Across methods, we keep the pretrained model, dataset, LoRA target modules, total rank budget, optimizer, training steps, and evaluation protocol fixed whenever applicable. We use AdamW [23] as the optimizer throughout the experiments. In StatLoRA, following the theoretical setup, the step size is chosen in the polynomially decaying form \(\gamma_t=\gamma_0 t^{-\kappa}\). We set \(\kappa=0.51\), which satisfies the condition \(\kappa\in(1/2,1)\) used in the asymptotic theory, and choose \(\gamma_0\) so that the resulting learning rate scale is comparable to that used in the AdaLoRA implementation. The momentum parameters \(\alpha\) and \(\beta\) are selected analogously so that the effective update magnitudes are comparable across methods. With these choices, the main difference among these methods lies in the criterion used to prune or retain rank-one components during the rank-reduction phase.

The final rank budget is set to one half of the initial total LoRA rank across all selected modules. For StatLoRA, the method-specific hyperparameters are the number of batches \(M\), the batch size \(b\) used in the batch-means variance estimator, and the threshold quantile \(q\). Since \(t=Mb\), we mainly choose \(b\) according to the length of the warm-up and pruning phases, typically in the range of \(50\) to \(150\). The threshold quantile is set to \(q=0.3\) in the main experiments, and its sensitivity is examined later.

Each experiment follows the same three-stage training procedure. During the warm-up phase, no pruning is performed, and the learning rate gradually increases and stabilizes. During the rank-reduction phase, rank-one components are pruned according to the specific allocation rule. After the target budget is reached, the selected ranks are fixed, and the model continues fine-tuning in the final phase with a gradually decreasing learning rate. Unless otherwise specified, experiments are repeated over five random seeds, and we report the mean performance together with the empirical standard deviation. All experiments are conducted on a single NVIDIA A100 GPU.

5.2 Natural Language Understanding↩︎

We start with the natural language understanding (NLU) tasks. NLU verifies whether a model can understand and compare natural language inputs, including sentence classification, textual entailment, semantic similarity, paraphrase detection, and linguistic acceptability. These tasks provide a standard testbed for examining whether a rank allocation method can maintain or improve downstream performance under a fixed rank budget.

We conduct experiments on seven tasks from the GLUE benchmark [57]: MNLI, SST-2, CoLA, STS-B, QNLI, RTE, and MRPC. MNLI and RTE are natural language inference tasks, where the model predicts the logical relationship between a premise and a hypothesis. SST-2 is a binary sentiment classification task. CoLA evaluates linguistic acceptability by asking whether a sentence is grammatically acceptable. STS-B measures the semantic similarity between two sentences. QNLI is a question-answering-derived entailment task. MRPC is a paraphrase identification task, where the model determines whether two sentences are semantically equivalent.

We use DeBERTaV3-base [58] as the pretrained language model and apply LoRA-based fine-tuning to the same set of modules for all compared methods. The baselines include vanilla LoRA with a fixed rank allocation, AdaLoRA, and IGU-LoRA. Following the standard GLUE evaluation protocol [57], we report accuracy for MNLI, SST-2, QNLI, and RTE; Matthews correlation for CoLA; Pearson correlation for STS-B; and F1 score for MRPC. All methods are evaluated on the validation set under the same training schedule and final rank budget.

For all NLU experiments, each LoRA module is initialized with rank \(r=2\). For rank allocation methods, the final rank budget is chosen so that the average retained rank is one per module. Thus, the total rank budget is one-half of the initial total LoRA rank across the model. Under this budget, some modules may be assigned rank zero, whereas others may retain one or more rank-one components. This setting tests whether the allocation method can identify where the limited adaptation capacity should be used. For StatLoRA, we use the batch-means variance estimator with batch size \(b=120\), \(M=t/b\) batches, and set the threshold quantile to \(q=0.3\). All methods are trained with AdamW and follow the same warm-up, rank-reduction, and final fine-tuning schedule.

As shown in Table 1, StatLoRA achieves the best performance on four of the seven reported tasks and obtains the second-best performance on two additional tasks. Compared with the fixed-rank LoRA baseline, StatLoRA improves the mean performance on all tasks under the same total rank budget. This suggests that adaptive allocation of rank-one components can improve parameter utilization when the number of adaptation components is limited. Compared with AdaLoRA and IGU-LoRA, StatLoRA remains competitive across the GLUE tasks, with favorable results on MNLI, SST-2, CoLA, and MRPC. The gains are not uniform across all tasks, which is expected because the GLUE tasks differ in sample size, task structure, and sensitivity to the available adaptation capacity.

Overall, the NLU results show that the proposed hypothesis-testing-based allocation rule provides a reliable alternative to heuristic importance-score criteria for LoRA rank allocation. By incorporating uncertainty estimates into component selection, StatLoRA can allocate the limited rank budget in a way that is at least competitive with existing adaptive-rank methods and often improves upon the fixed-rank baseline.

5.3 Natural Language Generation↩︎

We next evaluate StatLoRA on natural language generation (NLG) tasks. In contrast to NLU tasks, which are typically formulated as classification or matching problems over input texts, NLG tasks require the model to generate fluent and semantically faithful output sequences. We focus on abstractive summarization, where the goal is to produce a concise summary that captures the main information in a source document. This setting provides a useful testbed for examining whether an adaptive rank allocation method can preserve generation quality under a fine-tuning budget.

We conduct experiments on two widely used summarization datasets: XSum [59] and CNN/DailyMail [60], [61]. XSum is an extreme summarization dataset in which each input document is paired with a short, one-sentence summary. It therefore requires the model to generate highly compressed summaries and is sensitive to whether the model can identify salient information. CNN/DailyMail consists of news articles paired with multi-sentence summaries, and evaluates the model’s ability to generate longer summaries while preserving the main semantic content.

We use BART-Large [62] as the pretrained language model and apply LoRA-based fine-tuning to the same set of selected modules for all compared methods. The baselines include vanilla LoRA, AdaLoRA, and IGU-LoRA. All methods use the same pretrained model, training schedule, optimizer, target modules, and final rank budget. As in the NLU experiments, each LoRA module is initialized with rank \(r=2\), and the final rank budget is set to one-half of the initial total LoRA rank. For StatLoRA, we use the batch-means variance estimator with batch size \(b=100\), \(M=t/b\) batches, and set the threshold quantile to \(q=0.3\). Rank-one components are pruned during the rank-reduction phase according to the method-specific allocation rule.

Following standard evaluation practice for summarization, we report ROUGE-1, ROUGE-2, and ROUGE-L scores [63]. ROUGE-1 measures unigram overlap between the generated and reference summaries, ROUGE-2 measures bigram overlap, and ROUGE-L is based on the longest common subsequence. These metrics provide complementary measures of lexical overlap between generated summaries and reference summaries.

As shown in Table 2, StatLoRA achieves competitive performance on both summarization datasets under the same rank budget. On XSum, StatLoRA obtains the highest mean score across all three ROUGE metrics. The improvements over vanilla LoRA are modest, especially relative to the empirical standard deviations, but the results indicate that the proposed statistical allocation rule can retain useful LoRA components in the highly compressed summarization setting of XSum.

On CNN/DailyMail, StatLoRA obtains the highest mean ROUGE-2 and ROUGE-L scores and the second-highest mean ROUGE-1 score. Compared with vanilla LoRA, StatLoRA improves all three reported metrics under the same rank budget. IGU-LoRA gives a slightly higher ROUGE-1 score, whereas StatLoRA provides the best overall performance across ROUGE-2 and ROUGE-L. This suggests that StatLoRA remains competitive on longer-document summarization tasks, where the model must generate more detailed summaries while preserving semantic content.

Overall, the NLG results show that the proposed rank allocation rule is not limited to classification-style NLU tasks. Across both summarization datasets, StatLoRA matches or improves upon existing LoRA variants under the same rank budget. This suggests that incorporating uncertainty estimates into component selection can provide a reliable alternative to heuristic importance-score rules for adaptive LoRA rank allocation.

5.4 Question Answering↩︎

We further evaluate StatLoRA on question answering tasks using a larger pretrained language model. Compared with the previous experiments based on medium-scale pretrained models, this setting is closer to practical large language model adaptation, where parameter-efficient fine-tuning is important due to the computational and memory costs of updating all model parameters. Question answering tasks require the model to select correct answers from natural language questions, and may involve reading comprehension, factual knowledge, scientific reasoning, or commonsense reasoning. These tasks, therefore, provide a complementary setting for evaluating whether the proposed statistically guided rank allocation method remains effective for larger models and reasoning-oriented tasks.

We conduct experiments on four question answering datasets: BoolQ [64], ARC-Easy [65], OpenBookQA [66], and CommonsenseQA [67]. BoolQ is a binary question answering dataset in which the model answers yes/no questions based on a short passage. ARC-Easy consists of grade-school science questions with multiple-choice answers and evaluates basic scientific reasoning. OpenBookQA is a multiple-choice benchmark that requires combining elementary science facts with commonsense reasoning. CommonsenseQA evaluates commonsense reasoning through multiple-choice questions constructed from commonsense knowledge relations. Together, these datasets cover reading comprehension, scientific question answering, and commonsense reasoning.

Qwen2.5-7B [68] is used as the pretrained language model, and we apply LoRA-based fine-tuning to the same set of selected modules for all compared methods. The baselines include vanilla LoRA, AdaLoRA, and IGU-LoRA. For these large-model experiments, each LoRA module is initialized with rank \(r=16\), and the final rank budget is set to one-half of the initial total LoRA rank. For StatLoRA, we use the batch-means variance estimator with batch size \(b=100\), and set the threshold quantile to \(q=0.3\). All methods are trained with the same optimizer, learning-rate schedule, warm-up phase, rank-reduction phase, and final fine-tuning phase. Following standard evaluation practice for binary and multiple-choice question answering tasks, we report accuracy as the evaluation metric.

As shown in Table 3, StatLoRA achieves the highest mean accuracy on three of the four question answering tasks and the second-highest mean accuracy on BoolQ. On ARC-Easy and OpenBookQA, StatLoRA improves over both the fixed-rank LoRA and the adaptive-rank methods. On CommonsenseQA, StatLoRA also attains the highest mean accuracy, suggesting that the proposed allocation rule remains competitive on commonsense reasoning tasks. On BoolQ, AdaLoRA gives the highest mean accuracy, while StatLoRA obtains a close second-best result.

Overall, these results indicate that StatLoRA remains competitive when applied to a larger pretrained language model. Compared with the medium-scale experiments above, this setting is more representative of practical large language model fine-tuning, where full-parameter updates are often computationally expensive and parameter-efficient adaptation is desirable. The empirical pattern across ARC-Easy, OpenBookQA, and CommonsenseQA suggests that uncertainty-aware component selection can identify useful adaptation directions even in larger models and reasoning-oriented tasks.

5.5 Statistical Validation and Sensitivity Analysis↩︎

In this subsection, we examine the statistical behavior of the proposed StatLoRA method beyond its performance over other baseline methods. We first study the sensitivity of StatLoRA to its statistical hyperparameters, including the batch size \(b\) used in the batch-means variance estimator and the quantile level \(q\) used to construct the threshold in hypothesis testing. We then investigate the stability of the rank allocation across repeated runs and use the resulting allocation patterns to study how different layers and modules contribute to fine-tuning on specific datasets. Finally, we provide empirical diagnostics for the asymptotic theory of component scores in 3.

5.5.1 Sensitivity to Hyperparameters↩︎

We examine the sensitivity of StatLoRA to hyperparameters used in the statistical inference for LoRA components. In particular, StatLoRA involves two hyperparameters that directly enter the hypothesis-testing step: the batch size \(b\) in the batch-means variance estimator and the quantile level \(q\) used to determine the adaptive threshold. Since \(Mb\) determines the number of past score values used for estimating the long-run variance, we vary \(b\) to study its effect.

The batch size \(b\) controls the length of each batch used to estimate the long-run variance of the averaged score process. A smaller \(b\) uses more local information but may lead to noisier variance estimates, whereas a larger \(b\) produces more stable averages but reduces the number of batches. The quantile level \(q\) determines the threshold \(\Delta^{(t)}\) in the hypothesis testing by taking the \(q\)-quantile of the averaged component scores among active components. Thus, \(q\) controls the relative stringency of the component’s hypothesis testing criterion.

We conduct this ablation study on SST-2 from the GLUE datasets, using the same DeBERTaV3-base model and training protocol as in the main NLU experiments. We vary \(b\) and \(q\) separately while keeping all other hyperparameters fixed, including the initial LoRA rank, final rank budget, optimizer, learning-rate schedule, and rank-reduction schedule. This design isolates the effect of the variance-estimation and threshold-selection parameters on the final performance of StatLoRA.

As shown in 2, StatLoRA is stable over the tested ranges of \(b\) and \(q\). [fig:ablation95b] reports the sensitivity to the batch size \(b\) in the batch-means variance estimator. When \(b\) is small, the block averages may not adequately capture the temporal dependence of the score trajectory, which can lead to unstable variance estimates. When \(b\) is relatively large, it can still maintain satisfactory performances. In this experiment, \(b=120\) gives the highest mean accuracy and the smallest empirical standard deviation among the tested values, suggesting a favorable batch size given a fixed number of total iterations.

[fig:ablation95q] reports the sensitivity to the threshold quantile \(q\). The performance remains stable across a range of values, indicating that StatLoRA is not highly sensitive to the precise choice of \(q\). Very small or very large quantile levels lead to slightly lower mean accuracy or larger variability. This behavior is consistent with the role of \(q\) in determining the adaptive threshold: a smaller value of \(q\) yields a more permissive pruning criterion, whereas a larger value yields a more stringent criterion. The default value \(q=0.3\) gives the highest mean accuracy with relatively small variability in this experiment.

Overall, these results suggest that StatLoRA is robust to moderate changes in the batch means and thresholding parameters. Based on this sensitivity analysis, we use \(b=120\), and \(q=0.3\) as the default configuration in the main NLU experiments.

5.5.2 Stability of Rank Allocation↩︎

Beyond performance comparisons with existing methods, we further examine the stability of the rank allocation produced by StatLoRA. Adaptive rank allocation is inherently data- and optimization-dependent, and different random seeds may lead to different selected components even under the same total rank budget. For a statistical inference-based pruning procedure, it is therefore important to evaluate whether the resulting allocation patterns are reproducible and stable across independent runs, rather than being driven primarily by stochastic variation from random initialization, data sampling, or optimization noise.

We conduct this analysis on the OpenBookQA dataset in the question answering task using Qwen2.5-7B as the pretrained language model. The experiment is repeated over five independent random seeds, while keeping the pretrained model, optimizer, training schedule, initial rank, final rank budget, and StatLoRA hyperparameters fixed. For each run, we record the final retained rank of each LoRA module after the rank-reduction phase. We then compute the mean and empirical standard deviation of the retained ranks across runs. The mean heatmap in 3 (a) summarizes the typical allocation pattern selected by StatLoRA, whereas the standard-deviation heatmap in 3 (b) measures the stability of these allocation decisions across independent training trajectories.

This analysis complements the accuracy results in Table 3. While the performance measures the final task outcome, the rank allocation heatmaps provide a more direct view of the component-selection behavior of StatLoRA. Stable allocation patterns indicate that the hypothesis-testing-based pruning rule tends to select similar modules across repeated stochastic training runs, whereas large variability would indicate sensitivity to randomness during training and pruning.

As shown in 3, the rank allocation produced by StatLoRA exhibits a structured and stable pattern. The mean-rank heatmap in 3 (a) shows that the retained ranks are not uniformly distributed across layers and module types. Larger ranks are assigned mainly to middle layers, particularly from layers 7 to 19, and later layers (layers 22-24) tend to receive less adaptation. The corresponding standard deviation heatmap in 3 (b) shows that the variability is generally moderate relative to the scale of the retained ranks. This suggests that the proposed StatLoRA method is robust to randomness during fine-tuning.

The heatmaps also reveal module-wise heterogeneity. In particular, the MLP-related modules, especially \(\boldsymbol{W}_{\mathrm{up}}\) and \(\boldsymbol{W}_{\mathrm{gate}}\), tend to receive larger retained ranks than many of the attention projection matrices. This pattern suggests that, for OpenBookQA fine-tuning with Qwen2.5-7B, the feed-forward transformation blocks may play an important role in task-specific adaptation. This observation is consistent with the functional role of these modules: \(\boldsymbol{W}_{\mathrm{up}}\) expands the hidden representation into a higher-dimensional intermediate space, while \(\boldsymbol{W}_{\mathrm{gate}}\) controls the gated nonlinear transformation. We also observe that the key and value projection matrices, \(\boldsymbol{W}_{\mathrm{k}}\) and \(\boldsymbol{W}_{\mathrm{v}}\), receive relatively large rank allocations, indicating that the attention mechanism also contributes substantially to task adaptation. We regard these patterns as descriptive evidence derived from the learned rank allocations.

The standard-deviation heatmap provides further information about the uncertainty of the allocation. Modules with consistently large retained ranks, such as several \(\boldsymbol{W}_{\mathrm{up}}\) modules in earlier and middle layers, tend to have relatively small variability across runs, indicating that these allocation decisions are stable under different random seeds. Larger standard deviations appear mainly in transition regions where the mean retained rank is neither clearly high nor clearly low. This suggests that the main source of allocation variability is concentrated around statistically ambiguous modules, where the statistical evidence for retaining rank-one components is less decisive.

Overall, the rank-allocation analysis indicates that StatLoRA produces structured and stable allocation patterns across layers and module types. These observations provide additional evidence that the proposed hypothesis-testing-based pruning strategy is not merely optimizing evaluation accuracy, but also providing interpretable component selection behavior under a fixed rank budget.

5.5.3 P-value Diagnostics for Pruned Components↩︎

We next examine the p-values associated with the components removed by StatLoRA. This analysis is intended as a diagnostic check of the proposed hypothesis-testing interpretation. Recall that StatLoRA is based on the null hypothesis \[\mathcal{H}^{0}_{\ell,j}: s_{\ell,j}^{\star}\geq \Delta,\] where \(s_{\ell,j}^{\star}\) denotes the population score of the \((\ell,j)\)-th rank-one component and \(\Delta\) is the threshold. A small one-sided p-value provides evidence against this null hypothesis, and hence indicates that the corresponding component has insufficient statistical support for being retained. During the rank-reduction phase, StatLoRA removes components with the smallest p-values. We collect the p-values of the removed components and visualize their empirical distribution.

As shown in 4, the p-values of the pruned components are concentrated near zero. In particular, \(76.9\%\) of the removed components have p-values below \(0.05\). This pattern is consistent with the intended statistical interpretation of StatLoRA: the components selected for removal are typically those for which the null hypothesis is not well supported by the empirical score trajectory. In other words, the pruning rule is not based only on the magnitude of the empirical score, but also incorporates the estimated uncertainty of that score through the standardized test statistic.

This analysis should be interpreted as a diagnostic. The purpose of this diagnostic is to verify that the empirical pruning behavior is consistent with the proposed hypothesis testing criterion. The observed concentration near zero suggests that, in this experiment, most removed components have weak statistical evidence for retention relative to the threshold.

5.5.4 Empirical Evidence for Asymptotic Normality↩︎

We further examine the asymptotic normality underlying the hypothesis testing for component scores in StatLoRA. The central limit theorem in 3 is asymptotic and is derived for each empirical component score. Therefore, an empirical validation that closely reflects the asymptotic theorem requires repeated training runs under the same and suitable experimental setting.

To provide such a diagnostic, we conduct an additional experiment on OpenBookQA with Qwen2.5-7B. We repeat the training procedure over 200 independent random seeds without pruning, and randomly select a LoRA component for evaluation. For each run, we compute the averaged empirical score and construct the corresponding standardized statistic. Since the population score \(s_{\ell,j}^{\star}\) is unknown, we approximate it by the average score over the final 100 iterations. The variance is estimated using the same batch-means estimator as in StatLoRA.

As shown in 5, the empirical distribution of the standardized statistics is broadly consistent with the asymptotic normality result in 3. The histogram is centered near zero and approximately follows the general shape of the standard normal density, while the Q-Q plot shows that most empirical quantiles lie close to the diagonal reference line. Deviations are still visible, especially in the tails and around the central region. Such discrepancies are expected because the theoretical result is asymptotic, whereas the experiment is based on a finite training horizon of approximately 6000 iterations. In addition, \(s_{\ell,j}^{\star}\) is approximated by a late-stage average, and the variance is estimated by batch means, both of which introduce finite-sample approximation error.

Overall, this experiment provides empirical support for the asymptotic normality of the averaged score statistic derived in 3. We interpret the result as an empirical diagnostic for the asymptotic theory.

6 Conclusion and Discussion↩︎

In this paper, we developed StatLoRA, a statistically guided method for LoRA rank allocation. The key idea is to formulate component selection as a hypothesis testing problem, where rank-one LoRA components are retained or pruned according to their estimated contribution and the associated uncertainty. This provides a statistical interpretation of rank allocation, in contrast to existing adaptive rank methods based primarily on carefully designed importance scores. To support this hypothesis-testing-based procedure, we established asymptotic theory for stochastic optimization algorithms commonly used in deep learning, with particular emphasis on AdamW. Within the LoRA fine-tuning setting, we further derived central limit theory for empirical LoRA scores using the delta method. To make the inference procedure practically implementable, we used a batch-means estimator to estimate the long-run variance of the averaged score statistic and established its consistency under suitable conditions. These theoretical ingredients lead to a practical p-value-based pruning rule, which is integrated into the StatLoRA fine-tuning procedure. Comprehensive experiments across natural language understanding, natural language generation, and question answering tasks demonstrate that StatLoRA is competitive with, and often improves upon, existing fixed-rank and adaptive-rank LoRA methods under matched rank budgets. Beyond performance comparisons, we conducted sensitivity analyses for the statistical hyperparameters, examined the stability of the selected rank allocation across repeated runs, and provided empirical diagnostics for the asymptotic normality of component scores. These results support the practical usefulness of the proposed uncertainty-aware and hypothesis-testing-based rank allocation method.

The proposed framework also has several limitations. The p-values used by StatLoRA should be interpreted as evidence of pruning for each component, rather than as formal guarantees for simultaneous multiple testing. In addition, the testing procedure relies on asymptotic normality and finite-sample variance estimation from stochastic training trajectories. Although our empirical results suggest that the resulting method is stable and practically useful, the calibration of the tests may depend on the training dynamics, batch-means estimation, and finite-sample behavior of large-scale fine-tuning. Thus, StatLoRA should be viewed as a statistically guided rank allocation method.

Several directions remain for future work. First, it would be natural to incorporate multiple-testing correction or false discovery rate control into the component selection rule, especially when the number of LoRA components is large. The current framework performs hypothesis testing for each individual component, whereas a joint or multiple testing formulation may provide a more principled approach to rank allocation by accounting for dependence and interactions among components. Second, finite-sample calibration and inference methods would be valuable. Possible directions include self-normalized statistics, bootstrap calibration, and refined concentration bounds. Such results would strengthen the statistical guarantees beyond the asymptotic regime considered in this paper. Third, future work may consider adaptive procedures that both prune and reallocate components during training, allowing the rank budget to move more flexibly across layers and modules.

7 Proofs of Main Theorems↩︎

7.1 Preparation Lemmas↩︎

We first present several auxiliary lemmas that are used in the proofs of the main theorems. Their proofs are deferred to Appendix 11.

7.2 Main Lemmas↩︎

We present several key lemmas that are used in the proof of the asymptotic theory for the general stochastic recursion. Their detailed proofs are deferred to 10.

7.3 Proof of 1↩︎

Define the error process \[\boldsymbol{e}_t := \boldsymbol{z}_t - \boldsymbol{z}^\star.\] Since \(F(\boldsymbol{z}^\star)=\boldsymbol{0}\), the recursion \[\boldsymbol{z}_{t+1} = \boldsymbol{z}_t + \gamma_t \big(F(\boldsymbol{z}_t)+\boldsymbol{\xi}_{t}+\boldsymbol{r}_{t}\big)\] can be rewritten as \[\boldsymbol{e}_{t+1} = \boldsymbol{e}_t + \gamma_t \big(F(\boldsymbol{z}^\star+\boldsymbol{e}_t)-F(\boldsymbol{z}^\star)+\boldsymbol{\xi}_{t}+\boldsymbol{r}_{t}\big).\] By (A2) in 1, \(F\) is continuously differentiable in a neighborhood of \(\boldsymbol{z}^\star\), and hence admits the local expansion \[F(\boldsymbol{z}^\star+\boldsymbol{e}_t) = F(\boldsymbol{z}^\star)+\boldsymbol{J}\boldsymbol{e}_t+\rho(\boldsymbol{e}_t) = \boldsymbol{J}\boldsymbol{e}_t+\rho(\boldsymbol{e}_t),\] where \(\|\rho(\boldsymbol{e}_t)\| = \mathcal{O}(\|\boldsymbol{e}_t\|^2)\) as \(\boldsymbol{e}_t\to \boldsymbol{0}\). Therefore, \[\label{eq2} \boldsymbol{e}_{t+1} = (\boldsymbol{I}+\gamma_t \boldsymbol{J})\boldsymbol{e}_t + \gamma_t \boldsymbol{\xi}_{t} + \gamma_t\big(\rho(\boldsymbol{e}_t)+\boldsymbol{r}_{t}\big).\tag{22}\]

Let \[\Phi_{t,k}:=\prod_{j=k}^t (I+\gamma_j \boldsymbol{J}).\] By repeated substitution of 22 , we obtain the variation-of-constants representation \[\gamma_t^{-1/2}\boldsymbol{e}_{t+1} = \gamma_t^{-1/2}\Phi_{t,1}\boldsymbol{e}_1 + \sum_{k=1}^t \Phi_{t,k+1}\frac{\gamma_k}{\sqrt{\gamma_t}}\boldsymbol{\xi}_k + \gamma_t^{-1/2}\sum_{k=1}^t \Phi_{t,k+1}\gamma_k\big(\rho(\boldsymbol{e}_{k})+\boldsymbol{r}_k\big).\]

For the first term, it follows from ?? that \[\begin{align} \left\| \gamma_t^{-1/2}\Phi_{t,1}\boldsymbol{e}_1 \right\| & \leq C_0 \gamma_t^{-1/2} \prod_{j=1}^{t}\left(1 - c\gamma_j\right) \left\| \boldsymbol{e}_1 \right\| \\ & \leq C_0 \gamma_t^{-1/2} \exp\left(1-c\sum_{j=1}^{t} \gamma_{j}\right) \left\| \boldsymbol{e}_1 \right\| \to 0. \end{align}\] We then consider the third term, by 5, we have \[\gamma_t^{-1/2}\sum_{k=1}^t \Phi_{t,k+1}\gamma_k\big(\rho(\boldsymbol{e}_{k})+\boldsymbol{r}_k\big) \xrightarrow{p} 0.\]

It therefore remains to analyze the martingale term \[\mathcal{J}_t := \sum_{k=1}^t \Phi_{t,k+1}\frac{\gamma_k}{\sqrt{\gamma_t}}\boldsymbol{\xi}_k.\] By 4, under 1, we have \[\mathcal{J}_t \xrightarrow{d} \mathcal{N}(\boldsymbol{0},\boldsymbol{\Sigma}),\] where \(\boldsymbol{\Sigma}\) is the unique positive semidefinite solution to the corresponding Lyapunov equation. Hence, by Slutsky’s theorem, \[\gamma_t^{-1/2}(\boldsymbol{z}_t-\boldsymbol{z}^\star) = \gamma_t^{-1/2}\boldsymbol{e}_t \xrightarrow{d} \mathcal{N}(\boldsymbol{0},\boldsymbol{\Sigma}).\]

Now, we proceed to prove the asymptotic normality for the Polyak-Ruppert average. We continue from \[\boldsymbol{e}_{t+1} = \boldsymbol{e}_t + \gamma_t \big(\boldsymbol{J}\boldsymbol{e}_{t}+\boldsymbol{\xi}_{t}+\rho(\boldsymbol{e}_t)+\boldsymbol{r}_{t}\big).\] For simplicity, we denote \(\boldsymbol{\delta}_t = \rho(\boldsymbol{e}_t)+\boldsymbol{r}_{t}\). Rearranging gives \[\label{eq:telescoping-identity} \boldsymbol{J}\boldsymbol{e}_t = \frac{\boldsymbol{e}_{t+1}-\boldsymbol{e}_t}{\gamma_t} - \boldsymbol{\xi}_{t} - \boldsymbol{\delta}_t.\tag{23}\] Summing 23 from \(t=1\) to \(T\) yields \[\label{eq:sum-identity} \boldsymbol{J}\sum_{t=1}^T \boldsymbol{e}_t = \sum_{t=1}^T \frac{\boldsymbol{e}_{t+1}-\boldsymbol{e}_t}{\gamma_t} - \sum_{t=1}^T \boldsymbol{\xi}_{t} - \sum_{t=1}^T \boldsymbol{\delta}_t.\tag{24}\] Hence \[\label{eq:avg-decomposition} \sqrt{T}\,(\bar \boldsymbol{z}_T-\boldsymbol{z}^\star) = -\boldsymbol{J}^{-1}\frac{1}{\sqrt{T}}\sum_{t=1}^T \boldsymbol{\xi}_{t} + \boldsymbol{J}^{-1}\boldsymbol{R}_{T,1} - \boldsymbol{J}^{-1}\boldsymbol{R}_{T,2},\tag{25}\] where \[\boldsymbol{R}_{T,1}:=\frac{1}{\sqrt{T}}\sum_{t=1}^T \frac{\boldsymbol{e}_{t+1}-\boldsymbol{e}_t}{\gamma_t}, \qquad \boldsymbol{R}_{T,2}:=\frac{1}{\sqrt{T}}\sum_{t=1}^T \boldsymbol{\delta}_t.\] We now show that both remainder terms are negligible.

From the last-iterate result proved above, we have \(\boldsymbol{e}_t=O_p(\gamma_t^{1/2})\). Then, Assumption 1 gives \[\rho(\boldsymbol{e}_t)=O_p(\left\| \boldsymbol{e}_t\right\|^2)=O_p(\gamma_t).\] The negligible-remainder condition from Theorem 1 yields \(\boldsymbol{r}_t=O_p(\gamma_t)\), hence \[\boldsymbol{\delta}_t=O_p(\gamma_t).\] Therefore, \[\left\|\boldsymbol{R}_{T,2}\right\| \le \frac{1}{\sqrt{T}}\sum_{t=1}^T O_p(\gamma_t) = O_p\!\left(T^{-1/2}\sum_{t=1}^T t^{-\kappa}\right) = O_p\!\left(T^{1/2-\kappa}\right) = o_{p}(1),\] since \(\kappa>1/2\).

For \(\boldsymbol{R}_{T,1}\), summation by parts gives \[\sum_{t=1}^T \frac{\boldsymbol{e}_{t+1}-\boldsymbol{e}_t}{\gamma_t} = \frac{\boldsymbol{e}_{T+1}}{\gamma_T} - \frac{\boldsymbol{e}_1}{\gamma_{1}} + \sum_{t=2}^T \boldsymbol{e}_t\left(\frac{1}{\gamma_{t-1}}-\frac{1}{\gamma_t}\right).\] Since \(\boldsymbol{e}_t=O_p(\gamma_t^{1/2})\) and \(\gamma_t=\gamma_0 t^{-\kappa}\), \[\frac{\left\|\boldsymbol{e}_{T+1}\right\|}{\sqrt{T}\gamma_T} = O_p\!\left(T^{\kappa/2-1/2}\right) = o_{p}(1),\] and \[\frac{1}{\sqrt{T}}\sum_{t=2}^T \left\| \boldsymbol{e}_t\right\|\left(\frac{1}{\gamma_t}-\frac{1}{\gamma_{t-1}}\right) = O_p\!\left( \frac{1}{\sqrt{T}} \sum_{t=2}^T t^{-\kappa/2} t^{\kappa-1} \right) = O_p\!\left(T^{\kappa/2-1/2}\right) = o_{p}(1).\] Thus \(R_{T,1}= o_{p}(1)\).

Finally, by the martingale central limit theorem and the conditional covariance assumption, \[\frac{1}{\sqrt{T}}\sum_{t=1}^T \boldsymbol{\xi}_{t}\xrightarrow{d} \mathcal{N}(\boldsymbol{0},\boldsymbol{Q}).\] Applying Slutsky’s theorem to 25 proves \[\sqrt{T}\,(\bar \boldsymbol{z}_T-\boldsymbol{z}^\star)\xrightarrow{d} \mathcal{N}(\boldsymbol{0},\boldsymbol{J}^{-1}\boldsymbol{Q}\boldsymbol{J}^{-\top}).\]

7.4 Proof of 2↩︎

We verify Assumptions (A1)–(A6) in 1 one by one.

7.4.0.1 Verification of (A1).

By Assumption (B1), the augmented iterate \(\boldsymbol{z}_t\) converges almost surely to a limiting point \[\boldsymbol{z}^\star=(\boldsymbol{x}^\star,\boldsymbol{m}^\star,\boldsymbol{v}^\star),\] satisfying ?? and ?? . In particular, \[F(\boldsymbol{z}^\star)=\boldsymbol{0},\] and \[\boldsymbol{z}_t\to \boldsymbol{z}^\star, \quad \text{a.s. as}\quad t\to\infty.\] Hence Assumption (A1) holds.

7.4.0.2 Verification of (A2).

By Assumption (B2), the objective function \(f\) is twice continuously differentiable and the function \(\boldsymbol{q}(\boldsymbol{x})\) is continuously differentiable in a neighborhood of \(\boldsymbol{x}^\star\). Therefore, the mappings \[\boldsymbol{x}\mapsto \nabla f(\boldsymbol{x}), \qquad \boldsymbol{x}\mapsto \boldsymbol{q}(\boldsymbol{x})\] are continuously differentiable in a neighborhood of \(\boldsymbol{x}^\star\). Moreover, since \(\varepsilon>0\) , the map \[(\boldsymbol{m},\boldsymbol{v})\mapsto \frac{\boldsymbol{m}}{\sqrt{\boldsymbol{v}}+\varepsilon}\] is continuously differentiable on any neighborhood where the denominator remains positive componentwise. It follows that the mean-field mapping \[F(\boldsymbol{z}) = \begin{pmatrix} -\dfrac{\boldsymbol{m}}{\sqrt{\boldsymbol{v}}+\varepsilon}-\lambda \boldsymbol{x}\\[1mm] \alpha\bigl(\nabla f(\boldsymbol{x})-\boldsymbol{m}\bigr)\\[1mm] \beta\bigl(\boldsymbol{q}(\boldsymbol{x})-\boldsymbol{v}\bigr) \end{pmatrix}\] is continuously differentiable in a neighborhood of \(\boldsymbol{z}^\star\). Moreover, due to the local Lipschitzness of \(\nabla^2 f\) and \(\nabla \boldsymbol{q}\), \(\nabla F\) is also locally Lipschitz continuous. Hence, there exists a Jacobian matrix \[\boldsymbol{J}:=\nabla F(\boldsymbol{z}^\star),\] such that \[F(\boldsymbol{z}^\star+\boldsymbol{e}) = F(\boldsymbol{z}^\star)+\boldsymbol{J}\boldsymbol{e}+\rho(\boldsymbol{e}), \qquad \|\rho(\boldsymbol{e})\|=\mathcal{O}(\|\boldsymbol{e}\|^2), \quad \boldsymbol{e}\to\boldsymbol{0}.\] Thus Assumption (A2) holds.

7.4.0.3 Verification of (A3).

Recall that the augmented noise is \[\boldsymbol{\xi}_t= \begin{pmatrix} \boldsymbol{0}\\[1mm] \alpha \boldsymbol{\xi}_t^{(g)}\\[1mm] \beta\Bigl(2\nabla f(\boldsymbol{x}_t)\odot \boldsymbol{\xi}_t^{(g)}+\hat{\boldsymbol{\xi}}_t^{(g)}\Bigr) \end{pmatrix},\] where \[\hat{\boldsymbol{\xi}}_t^{(g)} = \boldsymbol{\xi}_t^{(g)\odot 2} - \mathbb{E}\!\left[\boldsymbol{\xi}_t^{(g)\odot 2}\mid \mathcal{F}_{t}\right].\] By Assumption (B3), \[\mathbb{E}[\boldsymbol{\xi}_t^{(g)}\mid \mathcal{F}_{t}]=\boldsymbol{0}.\] Also, \[\mathbb{E}[\hat{\boldsymbol{\xi}}_t^{(g)}\mid \mathcal{F}_{t}] = \mathbb{E}[\boldsymbol{\xi}_t^{(g)\odot 2}\mid \mathcal{F}_{t}] - \mathbb{E}[\boldsymbol{\xi}_t^{(g)\odot 2}\mid \mathcal{F}_{t}] = \boldsymbol{0}.\] Therefore, \[\mathbb{E}[\boldsymbol{\xi}_t\mid \mathcal{F}_{t}] = \boldsymbol{0},\] so that \(\{\boldsymbol{\xi}_t\}\) is a martingale difference sequence.

Next, we verify the bounded \((2+\delta)\)-moment. Since \(\boldsymbol{z}_t\to \boldsymbol{z}^\star\) almost surely, the sequence \(\{\boldsymbol{x}_t\}\) is almost surely bounded for all sufficiently large \(t\), and hence \(\{\nabla f(\boldsymbol{x}_t)\}\) is also almost surely bounded. Therefore, \[\|\boldsymbol{\xi}_t\| \le C\Bigl(\|\boldsymbol{\xi}_t^{(g)}\|+\|\boldsymbol{\xi}_t^{(g)}\|^2\Bigr)\] for some deterministic constant \(C>0\). By Assumption (B3), \[\sup_{t\ge 1}\mathbb{E}\!\left[\|\boldsymbol{\xi}_t^{(g)}\|^{4+\delta}\mid \mathcal{F}_{t}\right]<\infty, \qquad \text{a.s.}\] It follows that, \[\sup_{t\ge 1}\mathbb{E}\!\left[\|\boldsymbol{\xi}_t\|^{2+\delta/2}\mid \mathcal{F}_{t}\right]<\infty, \qquad \text{a.s.}\] Hence, Assumption (A3) holds.

7.4.0.4 Verification of (A4).

By definition, \[\boldsymbol{\xi}_t= \begin{pmatrix} \boldsymbol{0}\\[1mm] \alpha \boldsymbol{\xi}_t^{(g)}\\[1mm] \beta\Bigl(2\nabla f(\boldsymbol{x}_t)\odot \boldsymbol{\xi}_t^{(g)}+\hat{\boldsymbol{\xi}}_t^{(g)}\Bigr) \end{pmatrix},\] where \[\hat{\boldsymbol{\xi}}_t^{(g)} = \boldsymbol{\xi}_t^{(g)\odot 2} - \mathbb{E}\!\left[\boldsymbol{\xi}_t^{(g)\odot 2}\mid \mathcal{F}_{t}\right].\]

Since \(\boldsymbol{x}_t\to \boldsymbol{x}^\star\) almost surely and \(\nabla f\) is continuous, we have \[\nabla f(\boldsymbol{x}_t)\to \nabla f(\boldsymbol{x}^\star), \qquad \text{a.s.}\] Hence, \[2\nabla f(\boldsymbol{x}_t)\odot \boldsymbol{\xi}_t^{(g)} = 2\,\operatorname{diag}\!\bigl(\nabla f(\boldsymbol{x}_t)\bigr)\boldsymbol{\xi}_t^{(g)} \to 2\operatorname{diag}\!\bigl(\nabla f(\boldsymbol{x}^{\star})\bigr) \boldsymbol{\xi}_t^{(g)} = 2\operatorname{diag}\!\bigl(\boldsymbol{m}^\star\bigr) \boldsymbol{\xi}_t^{(g)}\] in the conditional second moment. Therefore, \[\mathbb{E}[\boldsymbol{\xi}_t\boldsymbol{\xi}_t^\top\mid \mathcal{F}_{t}] \to \boldsymbol{Q}, \qquad \text{a.s.},\] where \[\label{eq95Q95adamw} \boldsymbol{Q}= \begin{pmatrix} \boldsymbol{0}& \boldsymbol{0}& \boldsymbol{0}\\ \boldsymbol{0}& \boldsymbol{Q}_{22} & \boldsymbol{Q}_{23}\\ \boldsymbol{0}& \boldsymbol{Q}_{32} & \boldsymbol{Q}_{33} \end{pmatrix},\tag{26}\] with \[\begin{align} \boldsymbol{Q}_{22} &= \alpha^2 \boldsymbol{Q}^{(g)},\\[1mm] \boldsymbol{Q}_{23} &= \alpha\beta\bigl(2\boldsymbol{Q}^{(g)}\operatorname{diag}\!\bigl(\boldsymbol{m}^{\star}\bigr)+\boldsymbol{R}^{(g)}\bigr),\\[1mm] \boldsymbol{Q}_{32} &= \alpha\beta\bigl(2\operatorname{diag}\!\bigl(\boldsymbol{m}^{\star}\bigr)\boldsymbol{Q}^{(g)}+\boldsymbol{R}^{(g)\top}\bigr),\\[1mm] \boldsymbol{Q}_{33} &=\beta^2\Bigl( 4\operatorname{diag}\!\bigl(\boldsymbol{m}^{\star}\bigr)\boldsymbol{Q}^{(g)}\operatorname{diag}\!\bigl(\boldsymbol{m}^{\star}\bigr) +2\operatorname{diag}\!\bigl(\boldsymbol{m}^{\star}\bigr)\boldsymbol{R}^{(g)} +2\boldsymbol{R}^{(g)\top}\operatorname{diag}\!\bigl(\boldsymbol{m}^{\star}\bigr) +\hat{\boldsymbol{Q}}^{(g)} \Bigr). \end{align}\] Thus, Assumption (A4) holds.

7.4.0.5 Verification of (A5).

Assumption (B5) states exactly that the Jacobian matrix \(\boldsymbol{J}\), defined in ?? , is Hurwitz. Hence Assumption (A5) holds.

7.4.0.6 Verification of (A6).

The remainder term is \[\boldsymbol{r}_t= \begin{pmatrix} \boldsymbol{r}_t^{(x)}\\ \boldsymbol{0}\\ \boldsymbol{0} \end{pmatrix}, \qquad \boldsymbol{r}_t^{(x)} = \frac{\boldsymbol{m}_t}{\sqrt{\boldsymbol{v}_t}+\varepsilon} - \frac{\boldsymbol{m}_{t+1}}{\sqrt{\boldsymbol{v}_{t+1}}+\varepsilon}.\] Since \[\boldsymbol{m}_{t+1}-\boldsymbol{m}_t = \alpha_t\bigl(\nabla f(\boldsymbol{x}_t)-\boldsymbol{m}_t\bigr)+\alpha_t\boldsymbol{\xi}_t^{(g)},\] and \[\boldsymbol{v}_{t+1}-\boldsymbol{v}_t = \beta_t\bigl(\boldsymbol{q}(\boldsymbol{x}_t)-\boldsymbol{v}_t\bigr) + \beta_t\Bigl(2\nabla f(\boldsymbol{x}_t)\odot \boldsymbol{\xi}_t^{(g)}+\hat{\boldsymbol{\xi}}_t^{(g)}\Bigr),\] we obtain \[\mathbb{E}\!\left[\|\boldsymbol{m}_{t+1}-\boldsymbol{m}_t\|^2\mid \mathcal{F}_{t}\right] = O(\gamma_t^2),\] and \[\mathbb{E}\!\left[\|\boldsymbol{v}_{t+1}-\boldsymbol{v}_t\|^2\mid \mathcal{F}_{t}\right] = O(\gamma_t^2), \qquad \text{a.s.}\] Since the map \[(\boldsymbol{m},\boldsymbol{v})\mapsto \frac{\boldsymbol{m}}{\sqrt{\boldsymbol{v}}+\varepsilon}\] is locally Lipschitz in a neighborhood of \((\boldsymbol{m}^\star,\boldsymbol{v}^\star)\), there exists \(C>0\) such that \[\|\boldsymbol{r}_t^{(x)}\| \le C\Bigl(\|\boldsymbol{m}_{t+1}-\boldsymbol{m}_t\|+\|\boldsymbol{v}_{t+1}-\boldsymbol{v}_t\|\Bigr), \qquad \text{a.s.}\] Hence \[\|\boldsymbol{r}_t\|^2 \le 2C^2\|\boldsymbol{m}_{t+1}-\boldsymbol{m}_t\|^2 + 2C^2\|\boldsymbol{v}_{t+1}-\boldsymbol{v}_t\|^2, \qquad \text{a.s.}\] and therefore \[\mathbb{E}\!\left[\|\boldsymbol{r}_t\|^2\mid \mathcal{F}_{t}\right] = O(\gamma_t^2), \qquad \text{a.s.}\] which implies the required negligibility of the remainder term. Thus Assumption (A6) holds.

Combining the above verifications, we conclude that Assumptions (A1)–(A6) in 1 are satisfied. Therefore, 1 applies directly to the AdamW iterates, and \[\gamma_t^{-1/2}(\boldsymbol{z}_t-\boldsymbol{z}^\star)\xrightarrow{d}\mathcal{N}(\boldsymbol{0},\boldsymbol{\Sigma}),\] \[\sqrt{t}\left(\Bar{\boldsymbol{z}}_t - \boldsymbol{z}^{\star}\right)\;\xrightarrow{d}\; \mathcal{N}(\boldsymbol{0},\boldsymbol{J}^{-1}\boldsymbol{Q}\boldsymbol{J}^{-\top}),\] where \(\boldsymbol{\Sigma}\) is the unique positive semidefinite solution to \[\boldsymbol{J}\boldsymbol{\Sigma}+\boldsymbol{\Sigma}\boldsymbol{J}^\top+\boldsymbol{Q}=\boldsymbol{0}.\] This completes the proof.

7.5 Proof of 3↩︎

Fix a LoRA component \((\ell,j)\). Recall that the full optimizer state is denoted by \(\boldsymbol{z}_t\), and let \(\boldsymbol{P}_{\ell,j}\) be the selection matrix such that \[\boldsymbol{w}_{\ell,j}(t) := \begin{pmatrix} \boldsymbol{a}_{\ell,j}(t)\\ \boldsymbol{b}_{\ell,j}(t) \end{pmatrix} = \boldsymbol{P}_{\ell,j}\boldsymbol{z}_t .\] Similarly, \[\boldsymbol{w}_{\ell,j}^{\star} = \boldsymbol{P}_{\ell,j}\boldsymbol{z}^{\star}.\] By the central limit theorem for the optimizer state under Assumption 2, we have \[\gamma_t^{-1/2}(\boldsymbol{z}_t-\boldsymbol{z}^\star) \xrightarrow{d} \mathcal{N}(\boldsymbol{0},\boldsymbol{\Sigma}),\] and \[\sqrt t(\bar\boldsymbol{z}_t-\boldsymbol{z}^\star) \xrightarrow{d} \mathcal{N}(0,\boldsymbol{J}^{-1}\boldsymbol{Q}\boldsymbol{J}^{-\top}), \qquad \bar\boldsymbol{z}_t:=\frac{1}{t}\sum_{\tau=1}^t \boldsymbol{z}_\tau .\] Applying the continuous mapping theorem to the linear map \(\boldsymbol{P}_{\ell,j}\), we obtain \[\gamma_t^{-1/2} \left( \boldsymbol{w}_{\ell,j}(t)-\boldsymbol{w}_{\ell,j}^{\star} \right) \xrightarrow{d} \mathcal{N}(0,\boldsymbol{\Sigma}_{\ell,j}^{w}),\] where \[\boldsymbol{\Sigma}_{\ell,j}^{w} := \boldsymbol{P}_{\ell,j}\boldsymbol{\Sigma}\boldsymbol{P}_{\ell,j}^{\top}.\] Similarly, \[\sqrt t \left( \bar\boldsymbol{w}_{\ell,j,t}-\boldsymbol{w}_{\ell,j}^{\star} \right) \xrightarrow{d} \mathcal{N}(0,\bar{\boldsymbol{\Sigma}}_{\ell,j}^{w}),\] where \[\bar\boldsymbol{w}_{\ell,j,t} := \frac{1}{t}\sum_{\tau=1}^t \boldsymbol{w}_{\ell,j}(\tau), \qquad \bar{\boldsymbol{\Sigma}}_{\ell,j}^{w} := \boldsymbol{P}_{\ell,j}\boldsymbol{J}^{-1}\boldsymbol{Q}\boldsymbol{J}^{-\top}\boldsymbol{P}_{\ell,j}^{\top}.\]

Define \[h(\boldsymbol{w}_{\ell,j}) := \|\boldsymbol{a}_{\ell,j}\|^2\|\boldsymbol{b}_{\ell,j}\|^2 .\] Then \[s_{\ell,j}(t)=h(\boldsymbol{w}_{\ell,j}(t)), \qquad s_{\ell,j}^{\star}=h(\boldsymbol{w}_{\ell,j}^{\star}).\] Since \[s_{\ell,j}^{\star} = \|\boldsymbol{a}_{\ell,j}^{\star}\|^2 \|\boldsymbol{b}_{\ell,j}^{\star}\|^2 \geq \Delta>0,\] both \(\boldsymbol{a}_{\ell,j}^{\star}\) and \(\boldsymbol{b}_{\ell,j}^{\star}\) are nonzero. Hence \(h\) is continuously differentiable in a neighborhood of \(\boldsymbol{w}_{\ell,j}^{\star}\), and its gradient is \[\nabla h(\boldsymbol{w}_{\ell,j}^{\star}) = \begin{pmatrix} 2\|\boldsymbol{b}_{\ell,j}^{\star}\|^{2}\boldsymbol{a}_{\ell,j}^{\star}\\ 2\|\boldsymbol{a}_{\ell,j}^{\star}\|^{2}\boldsymbol{b}_{\ell,j}^{\star} \end{pmatrix}.\] By a first-order Taylor expansion, \[h(\boldsymbol{w}_{\ell,j}(t))-h(\boldsymbol{w}_{\ell,j}^{\star}) = \nabla h(\boldsymbol{w}_{\ell,j}^{\star})^\top \left( \boldsymbol{w}_{\ell,j}(t)-\boldsymbol{w}_{\ell,j}^{\star} \right) + r_t,\] where \[r_t = o_p\left( \|\boldsymbol{w}_{\ell,j}(t)-\boldsymbol{w}_{\ell,j}^{\star}\| \right).\] Since \[\boldsymbol{w}_{\ell,j}(t)-\boldsymbol{w}_{\ell,j}^{\star} = O_p(\gamma_t^{1/2}),\] we have \[\gamma_t^{-1/2}r_t=o_p(1).\] Therefore, \[\begin{align} \gamma_t^{-1/2} \left( s_{\ell,j}(t)-s_{\ell,j}^{\star} \right) &= \nabla h(\boldsymbol{w}_{\ell,j}^{\star})^\top \gamma_t^{-1/2} \left( \boldsymbol{w}_{\ell,j}(t)-\boldsymbol{w}_{\ell,j}^{\star} \right) +o_p(1). \end{align}\] Combining this expansion with the component-wise CLT and Slutsky’s theorem yields \[\gamma_t^{-1/2} \left( s_{\ell,j}(t)-s_{\ell,j}^{\star} \right) \xrightarrow{d} \mathcal{N}(0,\sigma_{\ell,j}^2),\] where \[\sigma_{\ell,j}^{2} = \nabla h(\boldsymbol{w}_{\ell,j}^{\star})^\top \boldsymbol{\Sigma}_{\ell,j}^{w} \nabla h(\boldsymbol{w}_{\ell,j}^{\star}).\] This proves ?? .

We next prove the averaged-score result. For each \(\tau\), the Taylor expansion gives \[s_{\ell,j}(\tau)-s_{\ell,j}^{\star} = \nabla h(\boldsymbol{w}_{\ell,j}^{\star})^\top \left( \boldsymbol{w}_{\ell,j}(\tau)-\boldsymbol{w}_{\ell,j}^{\star} \right) + r_\tau,\] where \[r_\tau = O_p\left( \|\boldsymbol{w}_{\ell,j}(\tau)-\boldsymbol{w}_{\ell,j}^{\star}\|^2 \right)\] under the local smoothness of \(h\). Since the optimizer-state central limit theory implies \[\boldsymbol{w}_{\ell,j}(\tau)-\boldsymbol{w}_{\ell,j}^{\star} = O_p(\gamma_\tau^{1/2}),\] we have \[r_\tau=O_p(\gamma_\tau).\] Therefore, \[\frac{1}{\sqrt t}\sum_{\tau=1}^t r_\tau = O_p\left( \frac{1}{\sqrt t}\sum_{\tau=1}^t \gamma_\tau \right).\] Because \(\gamma_\tau=\gamma_0\tau^{-\kappa}\) with \(\kappa\in(1/2,1)\), \[\frac{1}{\sqrt t}\sum_{\tau=1}^t \gamma_\tau = O\left(t^{1/2-\kappa}\right) = o(1).\] Hence \[\frac{1}{\sqrt t}\sum_{\tau=1}^t r_\tau=o_p(1).\] It follows that \[\begin{align} \sqrt t \left( \frac{1}{t}\sum_{\tau=1}^t s_{\ell,j}(\tau) - s_{\ell,j}^{\star} \right) &= \nabla h(\boldsymbol{w}_{\ell,j}^{\star})^\top \sqrt t \left( \frac{1}{t}\sum_{\tau=1}^t \boldsymbol{w}_{\ell,j}(\tau) - \boldsymbol{w}_{\ell,j}^{\star} \right) +o_p(1). \end{align}\] Using the averaged-state central limit theory and Slutsky’s theorem, we obtain \[\sqrt t \left( \frac{1}{t}\sum_{\tau=1}^t s_{\ell,j}(\tau) - s_{\ell,j}^{\star} \right) \xrightarrow{d} \mathcal{N}(0,\bar\sigma_{\ell,j}^2),\] where \[\bar\sigma_{\ell,j}^{2} = \nabla h(\boldsymbol{w}_{\ell,j}^{\star})^\top \bar{\boldsymbol{\Sigma}}_{\ell,j}^{w} \nabla h(\boldsymbol{w}_{\ell,j}^{\star}).\] This proves ?? .

Finally, since \(\boldsymbol{\Sigma}_{\ell,j}^{w}\) and \(\bar{\boldsymbol{\Sigma}}_{\ell,j}^{w}\) are positive semidefinite covariance matrices, both limiting variances are nonnegative: \[\sigma_{\ell,j}^{2}\geq 0, \qquad \bar\sigma_{\ell,j}^{2}\geq 0.\] The limiting normal distribution is understood to be degenerate when the corresponding variance is zero.

8 Adam Optimizer↩︎

8.1 Main Results↩︎

We consider the Adam optimizer [24] applied to the stochastic optimization problem 5 , where only noisy gradients are available.

Given a sequence of stochastic gradients \(\boldsymbol{g}_t := \nabla \mathcal{L}(\boldsymbol{x}_t;\zeta_t)\), Adam updates are given by \[\label{eq95adam} \begin{align} \boldsymbol{m}_{t+1} &= (1-\alpha_t) \boldsymbol{m}_t + \alpha_t\boldsymbol{g}_{t}, \\ \boldsymbol{v}_{t+1} &= (1-\beta_t) \boldsymbol{v}_t + \beta_t\boldsymbol{g}_{t}^{\odot 2}, \\ \boldsymbol{x}_{t+1} &= \boldsymbol{x}_t - \gamma_t \frac{\boldsymbol{m}_{t+1}}{\sqrt{\boldsymbol{v}_{t+1}} + \varepsilon}, \end{align}\tag{27}\] where \(\alpha_t = \alpha\gamma_t,\beta_t = \beta \gamma_t\) are step sizes, \(\varepsilon>0\) is a small number, and the division is understood elementwise. We define the augmented state \[\label{eq95state95adam} \boldsymbol{z}_t := (\boldsymbol{x}_t, \boldsymbol{m}_t, \boldsymbol{v}_t) \in \mathbb{R}^{3d},\tag{28}\] which are considered and updated during Adam’s iteration. Let \(\nabla f(\boldsymbol{x}_t)\) denote the true gradient, and write the stochastic gradient as \[\boldsymbol{g}_{t} = \nabla f(\boldsymbol{x}_t) + \boldsymbol{\xi}_{t}^{(g)},\] where \(\{\boldsymbol{\xi}_{t}^{(g)}\}\) is a martingale difference sequence for gradients. We denote \[\boldsymbol{q}(\boldsymbol{x}) = \mathbb{E}\left[\nabla \mathcal{L}(\boldsymbol{x};\zeta)^{\odot 2}\right].\] Then the updates can be rewritten as \[\begin{align} & \boldsymbol{m}_{t+1} = \boldsymbol{m}_t + \alpha_t\big(\nabla f(\boldsymbol{x}_t) - \boldsymbol{m}_t\big) + \alpha_t\boldsymbol{\xi}_{t}^{(g)}, \\ & \boldsymbol{v}_{t+1} = \boldsymbol{v}_t + \beta_t\big(\boldsymbol{q}(\boldsymbol{x}_t) - \boldsymbol{v}_t\big) + \beta_t\left(2\nabla f(\boldsymbol{x}_t)\odot \boldsymbol{\xi}_{t}^{(g)} + \boldsymbol{\xi}_{t}^{(g)\odot 2} - \mathbb{E}\left[\boldsymbol{\xi}_t^{(g)\odot 2}\mid \mathcal{F}_{t}\right]\right). \end{align}\] The update of \(\boldsymbol{x}_t\) can be written as \[\boldsymbol{x}_{t+1} = \boldsymbol{x}_t - \gamma_t \frac{\boldsymbol{m}_t}{\sqrt{\boldsymbol{v}_t}+\varepsilon} + \gamma_t\boldsymbol{r}^{(x)}_{t},\] where \(\boldsymbol{r}^{(x)}_{t} = \frac{\boldsymbol{m}_t}{\sqrt{\boldsymbol{v}_t}+\varepsilon} - \frac{\boldsymbol{m}_{t+1}}{\sqrt{\boldsymbol{v}_{t+1}}+\varepsilon}\) collects the higher-order terms arising from the difference between \((\boldsymbol{m}_{t+1},\boldsymbol{v}_{t+1})\) and \((\boldsymbol{m}_t,\boldsymbol{v}_t)\).

Based on the analysis and reformulation of the Adam algorithm, we define the mean-field mapping \(F:\mathbb{R}^{3d}\to\mathbb{R}^{3d}\) as \[F(\boldsymbol{z}) = \begin{pmatrix} -\dfrac{\boldsymbol{m}}{\sqrt{\boldsymbol{v}}+\varepsilon} \\ \alpha\big(\nabla f(\boldsymbol{x})-\boldsymbol{m}\big) \\ \beta\big(\boldsymbol{q}(\boldsymbol{x})-\boldsymbol{v}\big) \end{pmatrix}.\] Then, the corresponding martingale noise is given by \[\boldsymbol{\xi}_{t} = \begin{pmatrix} \boldsymbol{0}\\ \alpha\boldsymbol{\xi}_{t}^{(g)} \\ \beta\big(2\nabla f(\boldsymbol{x}_t)\odot \boldsymbol{\xi}_{t}^{(g)} + \boldsymbol{\xi}_{t}^{(g)\odot 2} - \mathbb{E}[\boldsymbol{\xi}_{t}^{(g)\odot 2}|\mathcal{F}_t]\big) \end{pmatrix}.\] The remainder term \(\boldsymbol{r}_{t}\) collects all higher-order corrections, defined as \[\boldsymbol{r}_t = \begin{pmatrix} \boldsymbol{r}_t^{(x)}\\ \boldsymbol{0}\\ \boldsymbol{0} \end{pmatrix} = \begin{pmatrix} \frac{\boldsymbol{m}_t}{\sqrt{\boldsymbol{v}_t}+\varepsilon} - \frac{\boldsymbol{m}_{t+1}}{\sqrt{\boldsymbol{v}_{t+1}}+\varepsilon}\\ \boldsymbol{0}\\ \boldsymbol{0} \end{pmatrix}.\] With the above definitions, the Adam iteration can be written in the stochastic approximation form \[\boldsymbol{z}_{t+1} = \boldsymbol{z}_t + \gamma_t\big(F(\boldsymbol{z}_t)+\boldsymbol{\xi}_{t}+\boldsymbol{r}_{t}\big).\]

8.2 Proof of 4↩︎

We verify (A1)–(A6) in 1 one by one.

8.2.0.1 Verification of (A1).

Assumption (C1) states that the augmented iterate \(\boldsymbol{z}_t=(\boldsymbol{x}_t,\boldsymbol{m}_t,\boldsymbol{v}_t)\) converges almost surely to \[\boldsymbol{z}^\star=(\boldsymbol{x}^\star,\boldsymbol{0},\boldsymbol{q}^{\star}),\] where \(\nabla f(\boldsymbol{x}^\star)=\boldsymbol{0}\) and \(\boldsymbol{q}^{\star} = \boldsymbol{q}(\boldsymbol{x}^{\star})\). By the definition of the Adam mean-field mapping \[F(\boldsymbol{z})= \begin{pmatrix} -\dfrac{\boldsymbol{m}}{\sqrt{\boldsymbol{v}}+\varepsilon}\\[1mm] \alpha(\nabla f(\boldsymbol{x})-\boldsymbol{m})\\[1mm] \beta(\boldsymbol{q}(\boldsymbol{x})-\boldsymbol{v}) \end{pmatrix},\] we have \[F(\boldsymbol{z}^\star) = \begin{pmatrix} \boldsymbol{0}\\ \alpha(\nabla f(\boldsymbol{x}^\star)-\boldsymbol{0})\\ \beta(\boldsymbol{q}(\boldsymbol{x}^{\star})-\boldsymbol{q}^{\star}) \end{pmatrix} = \boldsymbol{0}.\] Hence Assumption (A1) holds.

8.2.0.2 Verification of (A2).

By Assumption (C2), the objective function \(f\) is twice continuously differentiable in a neighborhood of \(\boldsymbol{x}^\star\). Therefore, the mappings \[\boldsymbol{x}\mapsto \nabla f(\boldsymbol{x}), \qquad \boldsymbol{x}\mapsto \nabla f(\boldsymbol{x})^{\odot 2}\] are continuously differentiable in a neighborhood of \(\boldsymbol{x}^\star\). Since the map \[(\boldsymbol{m},\boldsymbol{v})\mapsto \frac{\boldsymbol{m}}{\sqrt{\boldsymbol{v}}+\varepsilon}\] is continuously differentiable for \(\varepsilon>0\) and \(\boldsymbol{v}> 0\), and the function \(\boldsymbol{q}(\boldsymbol{x})\) is continuously differentiable in a neighborhood of \(\boldsymbol{x}^{\star}\), it follows that the mean-field mapping \(F\) is continuously differentiable in a neighborhood of \(\boldsymbol{z}^\star\). Moreover, the Lipschitzness of \(\nabla^2 f\) and \(\nabla \boldsymbol{q}\) implies that the Jacobian \(\nabla F\) is Lipschitz continuous in a neighborhood of \(\boldsymbol{z}^\star\). Consequently, its first-order expansion \[F(\boldsymbol{z}^\star+\boldsymbol{e}) = F(\boldsymbol{z}^\star)+\boldsymbol{J}\boldsymbol{e}+\rho(\boldsymbol{e}), \qquad \|\rho(\boldsymbol{e})\|=\mathcal{O}(\|\boldsymbol{e}\|^2),\] holds, and Assumption (A2) is satisfied.

8.2.0.3 Verification of (A3).

By definition, \[\boldsymbol{\xi}_t= \begin{pmatrix} \boldsymbol{0}\\[1mm] \alpha\boldsymbol{\xi}_t^{(g)}\\[1mm] \beta\bigl(2\nabla f(\boldsymbol{x}_t)\odot \boldsymbol{\xi}_t^{(g)}+\hat{\boldsymbol{\xi}}_t^{(g)}\bigr) \end{pmatrix}, \qquad \hat{\boldsymbol{\xi}}_t^{(g)} = \boldsymbol{\xi}_t^{(g)\odot 2} - \mathbb{E}[\boldsymbol{\xi}_t^{(g)\odot 2}\mid \mathcal{F}_{t}].\] Since \(\{\boldsymbol{\xi}_t^{(g)}\}\) is a martingale difference sequence, \[\mathbb{E}[\boldsymbol{\xi}_t^{(g)}\mid \mathcal{F}_{t}] = \boldsymbol{0}.\] Moreover, \[\mathbb{E}[\hat{\boldsymbol{\xi}}_t^{(g)}\mid \mathcal{F}_{t}] = \mathbb{E}[\boldsymbol{\xi}_t^{(g)\odot 2}\mid \mathcal{F}_{t}] - \mathbb{E}[\boldsymbol{\xi}_t^{(g)\odot 2}\mid \mathcal{F}_{t}] = \boldsymbol{0}.\] Therefore, \[\mathbb{E}[\boldsymbol{\xi}_t\mid \mathcal{F}_{t}] = \boldsymbol{0},\] so that \(\{\boldsymbol{\xi}_t\}\) is a martingale difference sequence.

It remains to verify the bounded \((2+\delta)\)-moment. By Assumption (C3), \[\sup_{t\ge1}\mathbb{E}\!\left[\|\boldsymbol{\xi}_t^{(g)}\|^{4+\delta}\mid \mathcal{F}_{t}\right] < \infty, \qquad \text{a.s.}\] Since \(\boldsymbol{x}_t\to \boldsymbol{x}^\star\) a.s. and the iterates are locally bounded, \(\nabla f(\boldsymbol{x}_t)\) is almost surely bounded for all sufficiently large \(t\). Hence \[\|2\nabla f(\boldsymbol{x}_t)\odot \boldsymbol{\xi}_t^{(g)}+\hat{\boldsymbol{\xi}}_t^{(g)}\| \le C\bigl(\|\boldsymbol{\xi}_t^{(g)}\|+\|\boldsymbol{\xi}_t^{(g)}\|^2\bigr)\] for some deterministic constant \(C>0\). It follows that \[\sup_{t\ge1}\mathbb{E}\!\left[\|\boldsymbol{\xi}_t\|^{2+\delta}\mid \mathcal{F}_{t}\right] < \infty, \qquad \text{a.s.},\] possibly after reducing \(\delta\) if necessary. Therefore Assumption (A3) holds.

8.2.0.4 Verification of (A4).

By definition, the noise term admits \[\boldsymbol{\xi}_t= \begin{pmatrix} \boldsymbol{0}\\ \alpha \boldsymbol{\xi}_t^{(g)}\\ \beta\bigl(2\nabla f(\boldsymbol{x}_t)\odot \boldsymbol{\xi}_t^{(g)}+\hat{\boldsymbol{\xi}}_t^{(g)}\bigr) \end{pmatrix}.\] Hence, \[\boldsymbol{\xi}_t\boldsymbol{\xi}_t^\top= \begin{pmatrix} \boldsymbol{0}& \boldsymbol{0}& \boldsymbol{0}\\ \boldsymbol{0}& \alpha^2\boldsymbol{\xi}_t^{(g)}\boldsymbol{\xi}_t^{(g)\top} & \alpha\beta\,\boldsymbol{\xi}_t^{(g)} \bigl(2\nabla f(\boldsymbol{x}_t)\odot \boldsymbol{\xi}_t^{(g)}+\hat{\boldsymbol{\xi}}_t^{(g)}\bigr)^\top \\[1mm] \boldsymbol{0}& \alpha\beta\bigl(2\nabla f(\boldsymbol{x}_t)\odot \boldsymbol{\xi}_t^{(g)}+\hat{\boldsymbol{\xi}}_t^{(g)}\bigr)\boldsymbol{\xi}_t^{(g)\top} & \beta^2 \bigl(2\nabla f(\boldsymbol{x}_t)\odot \boldsymbol{\xi}_t^{(g)}+\hat{\boldsymbol{\xi}}_t^{(g)}\bigr) \bigl(2\nabla f(\boldsymbol{x}_t)\odot \boldsymbol{\xi}_t^{(g)}+\hat{\boldsymbol{\xi}}_t^{(g)}\bigr)^\top \end{pmatrix}.\] Since \(\boldsymbol{x}_t\to \boldsymbol{x}^\star\) almost surely and \(\nabla f(\boldsymbol{x}^\star)=\boldsymbol{0}\), we have \[\nabla f(\boldsymbol{x}_t)\to \boldsymbol{0}, \qquad \text{a.s.}\] Therefore, every term involving the factor \(\nabla f(\boldsymbol{x}_t)\) vanishes asymptotically. Taking conditional expectation with respect to \(\mathcal{F}_{t}\) gives \[\mathbb{E}[\boldsymbol{\xi}_t\boldsymbol{\xi}_t^\top\mid \mathcal{F}_{t}] = \begin{pmatrix} \boldsymbol{0}& \boldsymbol{0}& \boldsymbol{0}\\ \boldsymbol{0}& \alpha^2\mathbb{E}[\boldsymbol{\xi}_t^{(g)}\boldsymbol{\xi}_t^{(g)\top}\mid \mathcal{F}_{t}] & o(1) + \alpha\beta\mathbb{E}\!\left[ \boldsymbol{\xi}_t^{(g)}\hat{\boldsymbol{\xi}}_t^{(g)\top} \mid \mathcal{F}_{t} \right] \\ \boldsymbol{0}& o(1) + \alpha\beta\mathbb{E}\!\left[ \hat{\boldsymbol{\xi}}_t^{(g)}\boldsymbol{\xi}_t^{(g)\top} \mid \mathcal{F}_{t-1} \right] & \beta^2\mathbb{E}[\hat{\boldsymbol{\xi}}_t^{(g)}\hat{\boldsymbol{\xi}}_t^{(g)\top}\mid \mathcal{F}_{t}] + o(1) \end{pmatrix}, \qquad \text{a.s.}\] Using the assumed limits, \[\mathbb{E}[\boldsymbol{\xi}_t^{(g)}\boldsymbol{\xi}_t^{(g)\top}\mid \mathcal{F}_{t}] \to \boldsymbol{Q}^{(g)}, \qquad \mathbb{E}[\hat{\boldsymbol{\xi}}_t^{(g)}\hat{\boldsymbol{\xi}}_t^{(g)\top}\mid \mathcal{F}_{t}] \to \hat{\boldsymbol{Q}}^{(g)},\] we conclude that \[\label{eq9} \mathbb{E}[\boldsymbol{\xi}_t\boldsymbol{\xi}_t^\top\mid \mathcal{F}_{t}] \to \begin{pmatrix} \boldsymbol{0}& \boldsymbol{0}& \boldsymbol{0}\\ \boldsymbol{0}& \alpha^2 \boldsymbol{Q}^{(g)} & \alpha\beta \boldsymbol{R}^{(g)}\\ \boldsymbol{0}& \alpha\beta \boldsymbol{R}^{(g)\top} & \beta^2 \hat{\boldsymbol{Q}}^{(g)} \end{pmatrix} = \boldsymbol{Q}, \qquad \text{a.s.}\tag{29}\]

8.2.0.5 Verification of (A5).

Since \(\nabla^2 f(\boldsymbol{x}^\star)\succ \boldsymbol{0}\) by Assumption (C1), the matrix \(\boldsymbol{J}\) is Hurwitz, according to 3. Hence, Assumption (A5) is satisfied.

8.2.0.6 Verification of (A6).

The remainder term is given by \[\boldsymbol{r}_t= \begin{pmatrix} \boldsymbol{r}_t^{(x)}\\ \boldsymbol{0}\\ \boldsymbol{0} \end{pmatrix}, \qquad \boldsymbol{r}_t^{(x)} = \frac{\boldsymbol{m}_t}{\sqrt{\boldsymbol{v}_t}+\varepsilon} - \frac{\boldsymbol{m}_{t+1}}{\sqrt{\boldsymbol{v}_{t+1}}+\varepsilon}.\] Since \[\boldsymbol{m}_{t+1}-\boldsymbol{m}_t = \alpha_t(\nabla f(\boldsymbol{x}_t)-\boldsymbol{m}_t)+\alpha_t\boldsymbol{\xi}_t^{(g)},\] and \[\boldsymbol{v}_{t+1}-\boldsymbol{v}_t = \beta_t(\boldsymbol{q}(\boldsymbol{x}_t)-\boldsymbol{v}_t) + \beta_t\bigl(2\nabla f(\boldsymbol{x}_t)\odot \boldsymbol{\xi}_t^{(g)}+\hat{\boldsymbol{\xi}}_t^{(g)}\bigr),\] we have \[\mathbb{E}\left[\|\boldsymbol{m}_{t+1}-\boldsymbol{m}_t\|^2 \mid \mathcal{F}_{t}\right]=\mathcal{O}(\gamma_t^2), \quad \mathbb{E}\left[\|\boldsymbol{v}_{t+1}-\boldsymbol{v}_t\|^2 \mid \mathcal{F}_{t}\right]=\mathcal{O}(\gamma_t^2), \quad \text{a.s.},\] under Assumptions (C1)–(C3). Since the map \[(\boldsymbol{m},\boldsymbol{v})\mapsto \frac{\boldsymbol{m}}{\sqrt{\boldsymbol{v}}+\varepsilon}\] is locally Lipschitz, it follows that \[\|\boldsymbol{r}_t^{(x)}\| \le C\bigl(\|\boldsymbol{m}_{t+1}-\boldsymbol{m}_t\|+\|\boldsymbol{v}_{t+1}-\boldsymbol{v}_t\|\bigr), \quad \text{a.s.}\] Therefore, \[\mathbb{E}\left[\|\boldsymbol{r}_t\|^2\mid \mathcal{F}_{t}\right]=\mathcal{O}(\gamma_t^2), \qquad \text{a.s.},\] and Assumption (A6) holds.

Combining the above verifications, we conclude that Assumptions (A1)–(A6) in 1 are satisfied under 3. Therefore, 1 applies directly to the Adam optimizer, and the corresponding asymptotic covariance matrix \(\boldsymbol{Q}\) is given by 29 .

9 Adafactor Optimizer↩︎

9.1 Main Results↩︎

We consider the AdaFactor optimizer applied to the stochastic optimization problem 5 , where only noisy gradients are available.

Given a sequence of stochastic gradients \(\boldsymbol{g}_t := \nabla \mathcal{L}(\boldsymbol{x}_t;\zeta_t)\), AdaFactor updates are given by \[\label{eq95adafactor} \begin{align} \boldsymbol{v}_{t+1} &= (1-\beta_t)\boldsymbol{v}_t+\beta_t \boldsymbol{g}_t^{\odot 2},\\ \boldsymbol{x}_{t+1} &= \boldsymbol{x}_t-\gamma_t \frac{\boldsymbol{g}_t}{\sqrt{\boldsymbol{v}_{t+1}}+\varepsilon}, \end{align}\tag{30}\] where \(\beta_t=\beta\gamma_t\), \(\varepsilon>0\) is a small constant, and the division is understood elementwise. We define the augmented state \[\boldsymbol{z}_t:=(\boldsymbol{x}_t,\boldsymbol{v}_t)\in\mathbb{R}^{2d}.\] Let \(\nabla f(\boldsymbol{x}_t)\) denote the true gradient, and write \[\boldsymbol{g}_t=\nabla f(\boldsymbol{x}_t)+\boldsymbol{\xi}_t^{(g)},\] where \(\{\boldsymbol{\xi}_t^{(g)}\}\) is a martingale difference sequence. Then, \[\begin{align} & \boldsymbol{v}_{t+1} = \boldsymbol{v}_t+\beta_t\bigl(\boldsymbol{q}(\boldsymbol{x}_t)-\boldsymbol{v}_t\bigr) +\beta_t\Bigl( 2\nabla f(\boldsymbol{x}_t)\odot \boldsymbol{\xi}_t^{(g)} +\boldsymbol{\xi}_t^{(g)\odot 2} - \mathbb{E}\left[\boldsymbol{\xi}_t^{(g)\odot 2}\mid \mathcal{F}_{t}\right] \Bigr),\\ & \boldsymbol{x}_{t+1} = \boldsymbol{x}_t - \gamma_t \frac{\nabla f(\boldsymbol{x}_t)}{\sqrt{\boldsymbol{v}_t}+\varepsilon} - \gamma_t \frac{\boldsymbol{\xi}_t^{(g)}}{\sqrt{\boldsymbol{v}_t}+\varepsilon} + \gamma_t \boldsymbol{r}_t^{(x)}, \end{align}\] where \[\boldsymbol{r}_t^{(x)} = \frac{\boldsymbol{g}_t}{\sqrt{\boldsymbol{v}_t}+\varepsilon} - \frac{\boldsymbol{g}_t}{\sqrt{\boldsymbol{v}_{t+1}}+\varepsilon}.\]

We define the mean-field mapping \(F:\mathbb{R}^{2d}\to\mathbb{R}^{2d}\) by \[F(\boldsymbol{z}) = \begin{pmatrix} -\dfrac{\nabla f(\boldsymbol{x})}{\sqrt{\boldsymbol{v}}+\varepsilon}\\[2mm] \beta\bigl(\boldsymbol{q}(\boldsymbol{x})-\boldsymbol{v}\bigr) \end{pmatrix}.\] The corresponding martingale noise is \[\boldsymbol{\xi}_t= \begin{pmatrix} -\frac{\boldsymbol{\xi}_t^{(g)}}{\sqrt{\boldsymbol{v}_t}+\varepsilon}\\[1mm] \beta\Bigl( 2\nabla f(\boldsymbol{x}_t)\odot \boldsymbol{\xi}_t^{(g)} +\hat{\boldsymbol{\xi}}_t^{(g)} \Bigr) \end{pmatrix},\] where \[\hat{\boldsymbol{\xi}}_t^{(g)} = \boldsymbol{\xi}_t^{(g)\odot 2} - \mathbb{E}\!\left[\boldsymbol{\xi}_t^{(g)\odot 2}\mid \mathcal{F}_{t}\right].\] The remainder term is \[\boldsymbol{r}_t= \begin{pmatrix} \boldsymbol{r}_t^{(x)}\\ \boldsymbol{0} \end{pmatrix}.\] Hence, the AdaFactor iteration can be written in the stochastic approximation form \[\boldsymbol{z}_{t+1} = \boldsymbol{z}_t+\gamma_t\bigl(F(\boldsymbol{z}_t)+\boldsymbol{\xi}_t+\boldsymbol{r}_t\bigr).\]

9.2 Proof of 5↩︎

We verify Assumptions (A1)–(A6) in 1 one by one.

9.2.0.1 Verification of (A1).

By Assumption (D1), the augmented iterate \(\boldsymbol{z}_t=(\boldsymbol{x}_t,\boldsymbol{v}_t)\) converges almost surely to a limiting point \[\boldsymbol{z}^\star=(\boldsymbol{x}^\star,\boldsymbol{v}^\star),\] where \(\nabla f(\boldsymbol{x}^\star)=\boldsymbol{0}\). By the definition of the AdaFactor mean-field mapping \[F(\boldsymbol{z})= \begin{pmatrix} -\dfrac{\nabla f(\boldsymbol{x})}{\sqrt{\boldsymbol{v}}+\varepsilon}\\[2mm] \beta\bigl(\boldsymbol{q}(\boldsymbol{x})-\boldsymbol{v}\bigr) \end{pmatrix},\] we have \[F(\boldsymbol{z}^\star) = \begin{pmatrix} -\dfrac{\nabla f(\boldsymbol{x}^\star)}{\sqrt{\boldsymbol{v}^\star} + \varepsilon}\\[2mm] \beta\bigl(\boldsymbol{q}(\boldsymbol{x}^\star)-\boldsymbol{v}^\star\bigr) \end{pmatrix} = \boldsymbol{0}.\] Hence Assumption (A1) holds.

9.2.0.2 Verification of (A2).

By Assumption (D2), the objective function \(f\) is twice continuously differentiable and the function \(\boldsymbol{q}(\boldsymbol{x})\) is continuously differentiable in a neighborhood of \(\boldsymbol{x}^\star\). Therefore, the mappings \[\boldsymbol{x}\mapsto \nabla f(\boldsymbol{x}), \qquad \boldsymbol{x}\mapsto \boldsymbol{q}(\boldsymbol{x})\] are continuously differentiable in a neighborhood of \(\boldsymbol{x}^\star\). Since \(\varepsilon>0\) and \(\boldsymbol{v}> 0\), the map \[\boldsymbol{v}\mapsto \frac{1}{\sqrt{\boldsymbol{v}}+\varepsilon}\] is also continuously differentiable. Hence, the mean-field mapping \(F\) is continuously differentiable in a neighborhood of \(\boldsymbol{z}^\star\). Moreover, the Lipschitzness of \(\nabla^2 f\) and \(\nabla \boldsymbol{q}\) implies that the Jacobian \(\nabla F\) is Lipschitz continuous in a neighborhood of \(\boldsymbol{z}^\star\). Consequently, its first-order expansion admits \[F(\boldsymbol{z}^\star+\boldsymbol{e}) = F(\boldsymbol{z}^\star)+\boldsymbol{J}\boldsymbol{e}+\rho(\boldsymbol{e}), \qquad \|\rho(\boldsymbol{e})\|=\mathcal{O}(\|\boldsymbol{e}\|^2), \quad \boldsymbol{e}\to\boldsymbol{0}.\] Thus Assumption (A2) holds.

9.2.0.3 Verification of (A3).

Recall that the augmented noise is \[\boldsymbol{\xi}_t= \begin{pmatrix} -\frac{\boldsymbol{\xi}_t^{(g)}}{\sqrt{\boldsymbol{v}_t}+\varepsilon}\\[1mm] \beta\Bigl(2\nabla f(\boldsymbol{x}_t)\odot \boldsymbol{\xi}_t^{(g)}+\hat{\boldsymbol{\xi}}_t^{(g)}\Bigr) \end{pmatrix},\] where \[\hat{\boldsymbol{\xi}}_t^{(g)} = \boldsymbol{\xi}_t^{(g)\odot 2} - \mathbb{E}\!\left[\boldsymbol{\xi}_t^{(g)\odot 2}\mid \mathcal{F}_t\right].\] By Assumption (D3), \[\mathbb{E}[\boldsymbol{\xi}_t^{(g)}\mid \mathcal{F}_t]=\boldsymbol{0}.\] Moreover, \[\mathbb{E}[\hat{\boldsymbol{\xi}}_t^{(g)}\mid \mathcal{F}_t] = \mathbb{E}[\boldsymbol{\xi}_t^{(g)\odot 2}\mid \mathcal{F}_t] - \mathbb{E}[\boldsymbol{\xi}_t^{(g)\odot 2}\mid \mathcal{F}_t] = \boldsymbol{0}.\] Hence, \[\mathbb{E}[\boldsymbol{\xi}_t\mid \mathcal{F}_t]=\boldsymbol{0},\] so that \((\boldsymbol{\xi}_t)\) is a martingale difference sequence.

Next, since \(\boldsymbol{x}_t\to \boldsymbol{x}^\star\) almost surely and \(\nabla f\) is continuous, the sequence of \(\nabla f(\boldsymbol{x}_t)\) is almost surely bounded for all sufficiently large \(t\). Therefore, \[\|\boldsymbol{\xi}_t\| \le C\Bigl(\|\boldsymbol{\xi}_t^{(g)}\|+\|\boldsymbol{\xi}_t^{(g)}\|^2\Bigr)\] for some deterministic constant \(C>0\). By Assumption (D3), \[\sup_{t\ge1} \mathbb{E}\!\left[\|\boldsymbol{\xi}_t^{(g)}\|^{4+\delta}\mid \mathcal{F}_t\right] <\infty, \qquad \text{a.s.}\] It follows that \[\sup_{t\ge1} \mathbb{E}\!\left[\|\boldsymbol{\xi}_t\|^{2+\delta/2}\mid \mathcal{F}_t\right] <\infty, \qquad \text{a.s.}\] Hence Assumption (A3) holds.

9.2.0.4 Verification of (A4).

By definition, \[\boldsymbol{\xi}_t= \begin{pmatrix} -\frac{\boldsymbol{\xi}_t^{(g)}}{\sqrt{\boldsymbol{v}_t}+\varepsilon}\\[1mm] \beta\Bigl(2\nabla f(\boldsymbol{x}_t)\odot \boldsymbol{\xi}_t^{(g)}+\hat{\boldsymbol{\xi}}_t^{(g)}\Bigr) \end{pmatrix}.\] Since \(\boldsymbol{x}_t\to \boldsymbol{x}^\star\) almost surely and \(\nabla f(\boldsymbol{x}^\star)=\boldsymbol{0}\), we have \[\nabla f(\boldsymbol{x}_t)\to \boldsymbol{0}, \qquad \text{a.s.}\] Therefore, the term \(2\nabla f(\boldsymbol{x}_t)\odot \boldsymbol{\xi}_t^{(g)}\) vanishes asymptotically in conditional second moment, and the conditional covariance matrix converges to \[\boldsymbol{Q}= \begin{pmatrix} \text{diag}\left(\frac{1}{\sqrt{\boldsymbol{v}^\star} + \varepsilon} \right)\boldsymbol{Q}^{(g)}\text{diag}\left(\frac{1}{\sqrt{\boldsymbol{v}^\star} + \varepsilon} \right) & -\text{diag}\left(\frac{\beta}{\sqrt{\boldsymbol{v}^\star} + \varepsilon} \right)\boldsymbol{R}^{(g)}\\[1mm] -\text{diag}\left(\frac{\beta}{\sqrt{\boldsymbol{v}^\star} + \varepsilon} \right)\boldsymbol{R}^{(g)\top} & \beta^2 \hat{\boldsymbol{Q}}^{(g)} \end{pmatrix}.\] Hence \[\mathbb{E}[\boldsymbol{\xi}_t\boldsymbol{\xi}_t^\top\mid \mathcal{F}_t] \to \boldsymbol{Q}, \qquad \text{a.s.}\] Thus Assumption (A4) holds.

9.2.0.5 Verification of (A5).

By the proposition for AdaFactor, the Jacobian matrix at the limiting point is \[\boldsymbol{J}= \begin{pmatrix} -\operatorname{diag}\!\left(\frac{1}{\sqrt{\boldsymbol{v}^\star}+\varepsilon}\right) \nabla^2 f(\boldsymbol{x}^\star) & \boldsymbol{0}\\[2mm] \beta \nabla \boldsymbol{q}(\boldsymbol{x}^\star) & -\beta \boldsymbol{I}_d \end{pmatrix}.\] Since \(\nabla^2 f(\boldsymbol{x}^\star)\succ \boldsymbol{0}\) by Assumption (D1), the matrix \(\boldsymbol{J}\) is Hurwitz, by 4. Therefore, Assumption (A5) holds.

9.2.0.6 Verification of (A6).

The remainder term is \[\boldsymbol{r}_t= \begin{pmatrix} \boldsymbol{r}_t^{(x)}\\ \boldsymbol{0} \end{pmatrix}, \qquad \boldsymbol{r}_t^{(x)} = \frac{\boldsymbol{g}_t}{\sqrt{\boldsymbol{v}_t}+\varepsilon} - \frac{\boldsymbol{g}_t}{\sqrt{\boldsymbol{v}_{t+1}}+\varepsilon}.\] We decompose \[\boldsymbol{r}_t^{(x)} = \left( \frac{\nabla f(\boldsymbol{x}_t)}{\sqrt{\boldsymbol{v}_t}+\varepsilon} - \frac{\nabla f(\boldsymbol{x}_t)}{\sqrt{\boldsymbol{v}_{t+1}}+\varepsilon} \right) + \left(\frac{\boldsymbol{\xi}_t^{(g)}}{\sqrt{\boldsymbol{v}_{t}}+\varepsilon} - \frac{\boldsymbol{\xi}_t^{(g)}}{\sqrt{\boldsymbol{v}_{t+1}}+\varepsilon}\right).\] Using the local Lipschitz continuity of \(\boldsymbol{v}\mapsto (\sqrt{\boldsymbol{v}}+\varepsilon)^{-1}\), together with \[\boldsymbol{v}_{t+1}-\boldsymbol{v}_t = \beta_t\bigl(\boldsymbol{q}(\boldsymbol{x}_t)-\boldsymbol{v}_t\bigr) + \beta_t\Bigl(2\nabla f(\boldsymbol{x}_t)\odot \boldsymbol{\xi}_t^{(g)}+\hat{\boldsymbol{\xi}}_t^{(g)}\Bigr),\] one obtains \[\mathbb{E}\!\left[\|\boldsymbol{v}_{t+1}-\boldsymbol{v}_t\|^2\mid \mathcal{F}_t\right] = O(\gamma_t^2), \qquad \text{a.s.}\] Hence the first term in the above decomposition of \(\boldsymbol{r}_t^{(x)}\) is of order \(O(\gamma_t)\) in conditional \(L^2\). The second term is controlled by the bounded fourth moments of \(\boldsymbol{\xi}_t^{(g)}\). Therefore, \[\mathbb{E}\!\left[\|\boldsymbol{r}_t\|^2\mid \mathcal{F}_t\right] = O(\gamma_t^2), \qquad \text{a.s.}\] which implies the required negligibility of the remainder term. Thus Assumption (A6) holds.

Combining the above verifications, we conclude that Assumptions (A1)–(A6) in 1 are satisfied. Therefore, 1 applies directly to the AdaFactor iterates, and \[\gamma_t^{-1/2}(\boldsymbol{z}_t-\boldsymbol{z}^\star) \;\xrightarrow{d}\; \mathcal{N}(\boldsymbol{0},\boldsymbol{\Sigma}),\] \[\sqrt{t}\left(\Bar{\boldsymbol{z}}_t - \boldsymbol{z}^{\star}\right)\;\xrightarrow{d}\; \mathcal{N}(\boldsymbol{0},\boldsymbol{J}^{-1}\boldsymbol{Q}\boldsymbol{J}^{-\top}),\] where \(\boldsymbol{\Sigma}\) is the unique positive semidefinite solution to \[\boldsymbol{J}\boldsymbol{\Sigma}+\boldsymbol{\Sigma}\boldsymbol{J}^\top+\boldsymbol{Q}=\boldsymbol{0}.\]

10 Proofs of Main Lemmas↩︎

10.1 Proof of 4↩︎

We aim to analyze the fluctuation term \[\mathcal{J}_t := \sum_{k=1}^{t}\Phi_{t,k+1}\frac{\gamma_k}{\sqrt{\gamma_t}}\boldsymbol{\xi}_{k},\] where \[\Phi_{t,k}:=\prod_{j=k}^{t}(\boldsymbol{I}+\gamma_j\boldsymbol{J}).\]

To prove the asymptotic normality of \(\mathcal{J}_t\), we apply the Cramér–Wold theorem. Fix any \(\boldsymbol{d}\in \mathbb{R}^d\), and define \[\label{eq3} \Delta_{t,k} := \boldsymbol{d}^\top \Phi_{t,k+1}\frac{\gamma_k}{\sqrt{\gamma_t}}\boldsymbol{\xi}_{k}, \qquad 1\le k\le t.\tag{31}\] Then \[\boldsymbol{d}^\top \mathcal{J}_t = \sum_{k=1}^{t}\Delta_{t,k}.\]

Since \(\Phi_{t,k+1}\gamma_k/\sqrt{\gamma_t}\) is \(\mathcal{F}_k\)-measurable and \[\mathbb{E}[\boldsymbol{\xi}_{k}\mid\mathcal{F}_k]=\boldsymbol{0},\] it follows that \[\mathbb{E}[\Delta_{t,k}\mid \mathcal{F}_k]=0.\] Hence, for each fixed \(t\), \((\Delta_{t,k})_{1\le k\le t}\) is a martingale difference array.

We now verify the two standard conditions of the martingale central limit theorem.

Step 1: conditional variance convergence. Define \[V_t(\boldsymbol{d}) := \sum_{k=1}^{t}\mathbb{E}\big[\Delta_{t,k}^2\mid\mathcal{F}_k\big].\] Using the definition of \(\Delta_{t,k}\) in 31 , we obtain \[V_t(\boldsymbol{d}) = \sum_{k=1}^{t} \frac{\gamma_k^2}{\gamma_t}\, \boldsymbol{d}^\top \Phi_{t,k+1} \mathbb{E}\big[\boldsymbol{\xi}_{k}\boldsymbol{\xi}_{k}^\top\mid\mathcal{F}_k\big] \Phi_{t,k+1}^\top \boldsymbol{d}.\] Note that \[\begin{align} & \left\|V_t(\boldsymbol{d}) - \sum_{k=1}^{t} \frac{\gamma_k^2}{\gamma_t}\, \boldsymbol{d}^\top \Phi_{t,k+1}\boldsymbol{Q}\Phi_{t,k+1}^\top \boldsymbol{d}\right\| \\ & \leq \gamma_t^{-1} \left\| \boldsymbol{d}\right\|^2 \sum_{k=1}^{t} \left\|\Phi_{t,k+1} \right\|^2\gamma_k^2 \left\| \mathbb{E}\big[\boldsymbol{\xi}_{k}\boldsymbol{\xi}_{k}^\top\mid\mathcal{F}_k\big] - \boldsymbol{Q}\right\|\\ & \leq C_0^2 \gamma_t^{-1} \left\| \boldsymbol{d}\right\|^2 \sum_{k=1}^{t} \prod_{j=k+1}^{t}\left(1-c\gamma_j\right)^2\gamma_k^2 \left\| \mathbb{E}\big[\boldsymbol{\xi}_{k}\boldsymbol{\xi}_{k}^\top\mid\mathcal{F}_k\big] - \boldsymbol{Q}\right\| \to 0, \end{align}\] where the last inequality comes from ?? and the convergence is implied by 1 and 1 (A4). Hence, \[V_t(\boldsymbol{d}) - \sum_{k=1}^{t} \frac{\gamma_k^2}{\gamma_t}\, \boldsymbol{d}^\top \Phi_{t,k+1}\boldsymbol{Q}\Phi_{t,k+1}^\top \boldsymbol{d} \;\xrightarrow[]{}\; 0,\]

Let \[\boldsymbol{U}_t := \sum_{k=1}^{t} \frac{\gamma_k^2}{\gamma_t}\, \Phi_{t,k+1}\boldsymbol{Q}\Phi_{t,k+1}^\top.\] Then, by the standard discrete Lyapunov argument in 3, \[\boldsymbol{U}_t \to \boldsymbol{\Sigma},\] where \(\boldsymbol{\Sigma}\) is the unique positive semidefinite solution of \[\boldsymbol{J}\boldsymbol{\Sigma}+ \boldsymbol{\Sigma}\boldsymbol{J}^\top + \boldsymbol{Q}= \boldsymbol{0}.\] Therefore, \[V_t(\boldsymbol{d}) \to \boldsymbol{d}^\top \boldsymbol{\Sigma}\boldsymbol{d}.\]

Step 2: Lindeberg condition. For any \(\bar\varepsilon>0\), consider \[L_t(\bar\varepsilon,\boldsymbol{d}) := \sum_{k=1}^{t} \mathbb{E}\Big[ \Delta_{t,k}^2 \mathbf{1}_{\{|\Delta_{t,k}|>\bar\varepsilon\}} \;\Big|\; \mathcal{F}_k \Big].\] By Hölder’s inequality and the elementary bound \[\mathbf{1}_{\{|x|>\bar\varepsilon\}} \le \frac{|x|^\delta}{\bar\varepsilon^\delta}, \qquad \delta>0,\] we have \[\Delta_{t,k}^2\mathbf{1}_{\{|\Delta_{t,k}|>\bar\varepsilon\}} \le \bar\varepsilon^{-\delta}|\Delta_{t,k}|^{2+\delta}.\] Therefore, \[L_t(\bar\varepsilon,\boldsymbol{d}) \le \bar\varepsilon^{-\delta} \sum_{k=1}^{t} \mathbb{E}\big[ |\Delta_{t,k}|^{2+\delta}\mid\mathcal{F}_k \big].\] Since \[|\Delta_{t,k}| \le \|\boldsymbol{d}\|\, \left\|\Phi_{t,k+1}\right\| \frac{\gamma_k}{\sqrt{\gamma_t}} \|\boldsymbol{\xi}_{k}\|,\] we obtain \[\mathbb{E}\big[ |\Delta_{t,k}|^{2+\delta}\mid\mathcal{F}_k \big] \le \left\| \boldsymbol{d}\right\|^{2+\delta} \left\|\Phi_{t,k+1}\right\|^{2+\delta} \left(\frac{\gamma_k}{\sqrt{\gamma_t}}\right)^{2+\delta} \mathbb{E}\big[\|\boldsymbol{\xi}_{k}\|^{2+\delta}\mid\mathcal{F}_k\big].\] By 1 (A4), \[\sup_k \mathbb{E}\big[\|\boldsymbol{\xi}_{k}\|^{2+\delta}\mid\mathcal{F}_k\big] < \infty \qquad \text{a.s.}\] Hence \[L_t(\bar\varepsilon,\boldsymbol{d}) \le C \sum_{k=1}^{t} \left\|\Phi_{t,k+1}\right\|^{2+\delta} \left(\frac{\gamma_k}{\sqrt{\gamma_t}}\right)^{2+\delta},\] for some constants \(C>0\). Using the exponential stability estimate \[\|\Phi_{t,k+1}\| \le C_0\!\prod_{j=k+1}^{t}\left(1-c\gamma_j\right),\] and 1, it follows that \[L_t(\bar\varepsilon,\boldsymbol{d}) = \mathcal{O}\left(\gamma_{t}^{\delta/2}\right)\to 0.\] Thus, the Lindeberg condition holds.

By the martingale central limit theorem for triangular arrays, we conclude that \[\boldsymbol{d}^\top \mathcal{J}_t \xrightarrow{d} \mathcal{N}(0,\boldsymbol{d}^\top \boldsymbol{\Sigma}\boldsymbol{d}).\] Since \(\boldsymbol{d}\in\mathbb{R}^d\) is arbitrary, the Cramér–Wold theorem yields \[\mathcal{J}_t \xrightarrow{d} \mathcal{N}(\boldsymbol{0},\boldsymbol{\Sigma}).\]

10.2 Proof of 5↩︎

Since \(\boldsymbol{J}\) is Hurwitz, there exists a symmetric positive definite matrix \(\boldsymbol{P}\) such that \[\boldsymbol{J}^\top \boldsymbol{P}+ \boldsymbol{P}\boldsymbol{J}= -\boldsymbol{I}.\] Define the function \[V_t := \boldsymbol{e}_t^\top \boldsymbol{P}\boldsymbol{e}_t.\] Since \(\boldsymbol{P}\succ \boldsymbol{0}\), there exist constants \(0 < c_1 \le c_2 < \infty\) such that \[c_1 \|\boldsymbol{e}_t\|^2 \le V_t \le c_2 \|\boldsymbol{e}_t\|^2.\]

By the definition of \(\rho(\boldsymbol{e})\) and the mean-value theorem, there exists matrix \(\boldsymbol{H}_t\) such that \(\rho(\boldsymbol{e}_t) = \boldsymbol{H}_t\boldsymbol{e}_t\) and \(\boldsymbol{H}_t \to \boldsymbol{0}\) as \(t \to \infty\), almost surely. We first introduce the following conditions:

1. \(\|\boldsymbol{z}_{t}\| \leq M\),

2. \(\left\| \boldsymbol{H}_t\right\| \leq M\),

for some constants \(M > 0\). According to 1, these conditions hold for all sufficiently large \(t\), almost surely. In other words, we define the stopping time \[\tau_{t_0} = \inf_{j}\{j \geq t_0:~\text{any of the above conditions does not hold at j-th iteration}\}.\] Then there exists a (potentially random) finite time \(t_0 < \infty\) such that \(\tau_{t_0} = \infty\), almost surely. We denote by \(\boldsymbol{1}_{\mathcal{A}}\) the indicator function of an event \(\mathcal{A}\), which equals \(1\) if the event holds and \(0\) otherwise.

We fix a sufficiently large \(t_0 >0\). We define \(\boldsymbol{G}_t = \boldsymbol{J}+ \boldsymbol{H}_t\), and consider the recursion \[\boldsymbol{e}_{t+1} = \boldsymbol{e}_t + \gamma_t \boldsymbol{G}_t \boldsymbol{e}_t + \gamma_t \boldsymbol{\xi}_{t} + \gamma_t \boldsymbol{r}_t,\] we have \[\begin{align} V_{t+1} &= (\boldsymbol{e}_t + \gamma_t \boldsymbol{G}_t \boldsymbol{e}_t + \gamma_t \boldsymbol{\xi}_{t} + \gamma_t \boldsymbol{r}_t)^\top \boldsymbol{P} (\boldsymbol{e}_t + \gamma_t \boldsymbol{G}_t \boldsymbol{e}_t + \gamma_t \boldsymbol{\xi}_{t} + \gamma_t \boldsymbol{r}_t) \\ &= \boldsymbol{e}_t^\top \boldsymbol{P}\boldsymbol{e}_t + 2\gamma_t \boldsymbol{e}_t^\top \boldsymbol{P}\boldsymbol{G}_t \boldsymbol{e}_t + 2\gamma_t \boldsymbol{e}_t^\top \boldsymbol{P}\boldsymbol{\xi}_{t} + 2 \gamma_t \boldsymbol{e}_t^\top \boldsymbol{P}\boldsymbol{r}_{t}\\ &+ \gamma_t^2 (\boldsymbol{G}_t \boldsymbol{e}_t + \boldsymbol{\xi}_{t} + \boldsymbol{r}_t)^\top \boldsymbol{P}(\boldsymbol{G}_t \boldsymbol{e}_t + \boldsymbol{\xi}_{t} + \boldsymbol{r}_t). \end{align}\] Taking conditional expectation with respect to \(\mathcal{F}_t\), and using \(\mathbb{E}[\boldsymbol{\xi}_{t}\mid \mathcal{F}_t] = \boldsymbol{0}\), we get \[\begin{align} \mathbb{E}[V_{t+1}\boldsymbol{1}_{\tau_{t_0} \geq t+1}\mid \mathcal{F}_t] & \leq V_t \boldsymbol{1}_{\tau_{t_0} \geq t} + 2\gamma_t \boldsymbol{e}_t^\top \boldsymbol{P}\boldsymbol{G}_t \boldsymbol{e}_t \boldsymbol{1}_{\tau_{t_0} \geq t} + 2\gamma_t \boldsymbol{e}_t^\top \boldsymbol{P}\mathbb{E}\big[\boldsymbol{r}_t\boldsymbol{1}_{\tau_{t_0} \geq t} \mid \mathcal{F}_t \big]\\ &+ \gamma_t^2 \mathbb{E}\big[ (\boldsymbol{G}_t \boldsymbol{e}_t + \boldsymbol{\xi}_{t} + \boldsymbol{r}_t)^\top \boldsymbol{P}(\boldsymbol{G}_t \boldsymbol{e}_t + \boldsymbol{\xi}_{t} + \boldsymbol{r}_t)\boldsymbol{1}_{\tau_{t_0} \geq t} \mid \mathcal{F}_t \big]. \end{align}\]

We now estimate each term separately.

Step 1: estimate of the drift term.

Since \(\boldsymbol{G}_t = \boldsymbol{J}+ \boldsymbol{H}_t\), we have \[\begin{align} 2 \boldsymbol{e}_t^\top \boldsymbol{P}\boldsymbol{G}_t \boldsymbol{e}_t &= 2 \boldsymbol{e}_t^\top \boldsymbol{P}\boldsymbol{J}\boldsymbol{e}_t + 2 \boldsymbol{e}_t^\top \boldsymbol{P}\boldsymbol{H}_t \boldsymbol{e}_t. \end{align}\] For the first term, \[\begin{align} 2 \boldsymbol{e}_t^\top \boldsymbol{P}\boldsymbol{J}\boldsymbol{e}_t = \boldsymbol{e}_t^\top (\boldsymbol{J}^\top \boldsymbol{P}+ \boldsymbol{P}\boldsymbol{J})\boldsymbol{e}_t = - \left\|\boldsymbol{e}_t\right\|^2. \end{align}\] For the perturbation term, \[\left|2 \boldsymbol{e}_t^\top \boldsymbol{P}\boldsymbol{H}_t \boldsymbol{e}_t\right| \le 2 \|\boldsymbol{P}\|\,\|\boldsymbol{H}_t\|\,\|\boldsymbol{e}_t\|^2.\] Hence \[2 \boldsymbol{e}_t^\top \boldsymbol{P}\boldsymbol{G}_t \boldsymbol{e}_t \leq - \|\boldsymbol{e}_t\|^2 + 2\|\boldsymbol{P}\|\,\|\boldsymbol{H}_t\|\,\|\boldsymbol{e}_t\|^2 \leq -\frac{1}{2}\|\boldsymbol{e}_t\|^2\] Using the norm equivalence \(V_t \le c_2 \|\boldsymbol{e}_t\|^2\), we further obtain \[\label{eq4} 2 \boldsymbol{e}_t^\top \boldsymbol{P}\boldsymbol{G}_t \boldsymbol{e}_t \boldsymbol{1}_{\tau_{t_0} \geq t} \le - \frac{1}{2c_2}V_t \boldsymbol{1}_{\tau_{t_0} \geq t} := -c_0V_t \boldsymbol{1}_{\tau_{t_0} \geq t},\tag{32}\] where we denote \(c_0 = \frac{1}{2c_2}\)

Step 2: estimate of the quadratic term.

Using \[(\boldsymbol{u}+\boldsymbol{v}+ \boldsymbol{w})^\top \boldsymbol{P}(\boldsymbol{u}+\boldsymbol{v}+ \boldsymbol{w}) \le 3 \left\| \boldsymbol{P}\right\|\left(\left\| \boldsymbol{u}\right\|^2 + \left\| \boldsymbol{v}\right\|^2 + \left\| \boldsymbol{w}\right\|^2\right),\] we get \[\begin{align} (\boldsymbol{G}_t \boldsymbol{e}_t + \boldsymbol{\xi}_{t} + \boldsymbol{r}_t)^\top \boldsymbol{P}(\boldsymbol{G}_t \boldsymbol{e}_t + \boldsymbol{\xi}_{t} + \boldsymbol{r}_t) \le 3 \|\boldsymbol{P}\|\,\left(\|\boldsymbol{G}_t \boldsymbol{e}_t\|^2 + \|\boldsymbol{\xi}_{t}\|^2+ \left\| \boldsymbol{r}_t\right\|^2\right). \end{align}\] Hence \[\begin{align} & \mathbb{E}\big[ (\boldsymbol{G}_t \boldsymbol{e}_t + \boldsymbol{\xi}_{t} + \boldsymbol{r}_t)^\top \boldsymbol{P}(\boldsymbol{G}_t \boldsymbol{e}_t + \boldsymbol{\xi}_{t} + \boldsymbol{r}_t) \mid \mathcal{F}_t \big]\\ & \leq 3 \|\boldsymbol{P}\| \left(\|\boldsymbol{G}_t\|^2 \|\boldsymbol{e}_t\|^2 + \mathbb{E}\big[\|\boldsymbol{\xi}_{t}\|^2 \mid \mathcal{F}_t\big] + \mathbb{E}\big[\|\boldsymbol{r}_t\|^2 \mid \mathcal{F}_t\big] \right). \end{align}\] By 1 (A3) and (A6), \[\mathbb{E}[\|\boldsymbol{\xi}_{t}\|^2\mid \mathcal{F}_t] \le C_\xi,\qquad \mathbb{E}[\|\boldsymbol{r}_{t}\|^2\mid \mathcal{F}_t] \le C_r \gamma_t^2,\] for some \(C_{\xi}, C_r > 0\). Therefore, for some constant \(C>0\), \[\mathbb{E}\big[ (\boldsymbol{G}_t \boldsymbol{e}_t + \boldsymbol{\xi}_{t} + \boldsymbol{r}_t)^\top \boldsymbol{P}(\boldsymbol{G}_t \boldsymbol{e}_t + \boldsymbol{\xi}_{t} + \boldsymbol{r}_t) \mid \mathcal{F}_t \big] \le C(\|\boldsymbol{e}_t\|^2+1).\] Using again the equivalence between \(V_t\) and \(\|\boldsymbol{e}_t\|^2\), we obtain \[\label{eq5} \mathbb{E}\big[ (\boldsymbol{G}_t \boldsymbol{e}_t + \boldsymbol{\xi}_{t} + \boldsymbol{r}_t)^\top \boldsymbol{P}(\boldsymbol{G}_t \boldsymbol{e}_t + \boldsymbol{\xi}_{t} + \boldsymbol{r}_t) \mid \mathcal{F}_t \big] \le C(V_t+1).\tag{33}\]

Step 3: one-step recursive inequality.

Combining the previous bounds in 32 and 33 yields \[\begin{align} \mathbb{E}[V_{t+1} \boldsymbol{1}_{\tau_{t_0} \geq t+1}\mid \mathcal{F}_t] &\le V_t \boldsymbol{1}_{\tau_{t_0} \geq t} - c_0 \gamma_t V_t\boldsymbol{1}_{\tau_{t_0} \geq t} + C \gamma_t^2 (V_t+1) \boldsymbol{1}_{\tau_{t_0} \geq t}. \end{align}\] Since \(\gamma_t \to 0\), for all sufficiently large \(t\), \[C\gamma_t^2 V_t \le \frac{c_0}{2}\gamma_t V_t.\] Thus, \[\mathbb{E}[V_{t+1}\boldsymbol{1}_{\tau_{t_0} \geq t+1}\mid \mathcal{F}_t] \le (1-c_1\gamma_t)V_t \boldsymbol{1}_{\tau_{t_0} \geq t} + C\gamma_t^2,\] for some constant \(c_1>0\).

Taking full expectation, we obtain \[\mathbb{E}[V_{t+1} \boldsymbol{1}_{\tau_{t_0} \geq t+1}] \le (1-c_1\gamma_t)\mathbb{E}[V_t \boldsymbol{1}_{\tau_{t_0} \geq t}] + C\gamma_t^2.\]

Step 4: the convergence rates for \(\boldsymbol{e}_t\) and \(\rho(\boldsymbol{e}_t)\).

The scalar recursion \[u_{t+1}\le (1-c_1\gamma_t)u_t + C\gamma_t^2\] implies \[u_t = \mathcal{O}(\gamma_t).\] Applying this with \(u_t=\mathbb{E}[V_t \boldsymbol{1}_{\tau_{t_0} \geq t}]\), we conclude that \[\mathbb{E}[V_t \boldsymbol{1}_{\tau_{t_0} \geq t}] = \mathcal{O}(\gamma_t).\] Finally, since \(c_1 \|\boldsymbol{e}_t\|^2 \le V_t\), we obtain \[\mathbb{E}[\|\boldsymbol{e}_t\|^2 \boldsymbol{1}_{\tau_{t_0} \geq t}] = \mathcal{O}(\gamma_t).\] Therefore, \(\|\boldsymbol{e}_t\| \boldsymbol{1}_{\tau_{t_0} \geq t} = \mathcal{O}_{p}\left(\gamma_t^{1/2}\right)\) and \(\left\|\rho(\boldsymbol{e}_t)\right\| \boldsymbol{1}_{\tau_{t_0} \geq t} = o_{p}\left(\gamma_t^{1/2}\right)\).

Step 5: the convergence rate for the weighted term.

Finally, we are ready to bound \[\begin{align} & \left\|\sum_{k=1}^{t} \Phi_{t,k+1} \gamma_k \rho(\boldsymbol{e}_k) \right\|\\ & \leq \sum_{k=0}^{t} \left\| \Phi_{t,k+1}\right\| \gamma_k \left\| \rho(\boldsymbol{e}_k) \right\| \\ & \leq C_0 \sum_{k=0}^{t} \prod_{j=k+1}^{t}\left(1-c\gamma_j\right) \gamma_k \left\| \rho(\boldsymbol{e}_k) \right\|. \end{align}\] Since \(\left\|\rho(\boldsymbol{e}_t)\right\| \boldsymbol{1}_{\tau_{t_0} \geq t} = o_{p}\left(\gamma_t^{1/2}\right)\), we have \(\left\|\sum_{k=1}^{t} \Phi_{t,k+1} \gamma_k \rho(\boldsymbol{e}_k) \right\| = o_{p}\left(\gamma_t^{1/2}\right)\), by Lemma C.3 in [69]. We also note that \(\left\|\boldsymbol{r}_t\right\| = \mathcal{O}_{p}(\gamma_t) = o_p(\gamma_t^{1/2})\) implies \(\left\|\sum_{k=1}^{t} \Phi_{t,k+1} \gamma_k \boldsymbol{r}_k \right\| = o_{p}\left(\gamma_t^{1/2}\right)\), which completes the proof.

10.3 Proof of 1↩︎

We prove the two statements separately.

10.3.0.1 Characterization of the limiting point.

Let \(\boldsymbol{z}^\star=(\boldsymbol{x}^\star,\boldsymbol{m}^\star,\boldsymbol{v}^\star)\) be a solution point of the AdamW mean-field dynamics. By definition, \[F(\boldsymbol{z}^\star)=\boldsymbol{0},\] that is, \[\begin{align} -\frac{\boldsymbol{m}^\star}{\sqrt{\boldsymbol{v}^\star}+\varepsilon}-\lambda \boldsymbol{x}^\star &= \boldsymbol{0}, \tag{34}\\ \alpha\bigl(\nabla f(\boldsymbol{x}^\star)-\boldsymbol{m}^\star\bigr) &= \boldsymbol{0}, \tag{35}\\ \beta\bigl(\boldsymbol{q}(\boldsymbol{x}^\star)-\boldsymbol{v}^\star\bigr) &= \boldsymbol{0}. \tag{36} \end{align}\] Since \(\alpha,\beta>0\), 35 and 36 imply \[\boldsymbol{m}^\star=\nabla f(\boldsymbol{x}^\star), \qquad \boldsymbol{v}^\star=\boldsymbol{q}(\boldsymbol{x}^\star)=\boldsymbol{q}^\star.\] Substituting these identities into 34 , we obtain \[-\frac{\nabla f(\boldsymbol{x}^\star)}{\sqrt{\boldsymbol{v}^\star}+\varepsilon}-\lambda \boldsymbol{x}^\star=\boldsymbol{0},\] or equivalently, \[\frac{\nabla f(\boldsymbol{x}^\star)}{\sqrt{\boldsymbol{v}^\star}+\varepsilon}+\lambda \boldsymbol{x}^\star=\boldsymbol{0}.\] This proves the first statement.

10.3.0.2 Computation of the Jacobian matrix.

Write \[F(\boldsymbol{z})= \begin{pmatrix} F_1(\boldsymbol{x},\boldsymbol{m},\boldsymbol{v})\\[1mm] F_2(\boldsymbol{x},\boldsymbol{m},\boldsymbol{v})\\[1mm] F_3(\boldsymbol{x},\boldsymbol{m},\boldsymbol{v}) \end{pmatrix},\] where \[\begin{align} F_1(\boldsymbol{x},\boldsymbol{m},\boldsymbol{v}) &= -\frac{\boldsymbol{m}}{\sqrt{\boldsymbol{v}}+\varepsilon}-\lambda \boldsymbol{x},\\ F_2(\boldsymbol{x},\boldsymbol{m},\boldsymbol{v}) &= \alpha\bigl(\nabla f(\boldsymbol{x})-\boldsymbol{m}\bigr),\\ F_3(\boldsymbol{x},\boldsymbol{m},\boldsymbol{v}) &= \beta\bigl(\boldsymbol{q}(\boldsymbol{x})-\boldsymbol{v}\bigr). \end{align}\]

We compute the block derivatives of \(F\) at \(\boldsymbol{z}^\star=(\boldsymbol{x}^\star,\boldsymbol{m}^\star,\boldsymbol{v}^\star)\).

Derivative of the first block. Since \[F_1(\boldsymbol{x},\boldsymbol{m},\boldsymbol{v}) = -\lambda \boldsymbol{x} - \operatorname{diag}\!\left(\frac{1}{\sqrt{\boldsymbol{v}}+\varepsilon}\right)\boldsymbol{m},\] we have \[\frac{\partial F_1}{\partial \boldsymbol{x}}(\boldsymbol{z}^\star)=-\lambda \boldsymbol{I}_d.\] Differentiating with respect to \(\boldsymbol{m}\) gives \[\frac{\partial F_1}{\partial \boldsymbol{m}}(\boldsymbol{z}^\star) = -\operatorname{diag}\!\left( \frac{1}{\sqrt{\boldsymbol{v}^\star}+\varepsilon} \right) =: -\boldsymbol{D}^\star.\] Next, for each component, \[[F_1]_i(\boldsymbol{x},\boldsymbol{m},\boldsymbol{v}) = -\frac{m_i}{\sqrt{v_i}+\varepsilon}-\lambda x_i.\] Hence, for \(v_i^\star>0\), \[\frac{\partial [F_1]_i}{\partial v_i}(\boldsymbol{z}^\star) = \frac{m_i^\star}{2\sqrt{v_i^\star}(\sqrt{v_i^\star}+\varepsilon)^2}.\] Therefore, \[\frac{\partial F_1}{\partial \boldsymbol{v}}(\boldsymbol{z}^\star) = \operatorname{diag}\!\left( \frac{\boldsymbol{m}^\star}{2\sqrt{\boldsymbol{v}^\star}(\sqrt{\boldsymbol{v}^\star}+\varepsilon)^{\odot 2}} \right) =: \boldsymbol{C}^\star,\] where the division and multiplication are conducted elementwise for vectors in \(\frac{\boldsymbol{m}^\star}{2\sqrt{\boldsymbol{v}^\star}(\sqrt{\boldsymbol{v}^\star}+\varepsilon)^{\odot 2}}\), i.e., \(\left[\frac{\boldsymbol{m}^\star}{2\sqrt{\boldsymbol{v}^\star}(\sqrt{\boldsymbol{v}^\star}+\varepsilon)^{\odot 2}}\right]_{i} = \frac{m_i^\star}{2 \sqrt{v_i^\star}(\sqrt{v_i^\star} + \varepsilon)^2}\).

Derivative of the second block. Since \[F_2(\boldsymbol{x},\boldsymbol{m},\boldsymbol{v})=\alpha\bigl(\nabla f(\boldsymbol{x})-\boldsymbol{m}\bigr),\] we immediately obtain \[\frac{\partial F_2}{\partial \boldsymbol{x}}(\boldsymbol{z}^\star)=\alpha \nabla^2 f(\boldsymbol{x}^\star), \qquad \frac{\partial F_2}{\partial \boldsymbol{m}}(\boldsymbol{z}^\star)=-\alpha \boldsymbol{I}_d, \qquad \frac{\partial F_2}{\partial \boldsymbol{v}}(\boldsymbol{z}^\star)=\boldsymbol{0}.\]

Derivative of the third block. For \[F_3(\boldsymbol{x},\boldsymbol{m},\boldsymbol{v})=\beta\bigl(\boldsymbol{q}(\boldsymbol{x})-\boldsymbol{v}\bigr),\] we have \[\frac{\partial F_3}{\partial \boldsymbol{m}}(\boldsymbol{z}^\star)=\boldsymbol{0}, \qquad \frac{\partial F_3}{\partial \boldsymbol{v}}(\boldsymbol{z}^\star)=-\beta \boldsymbol{I}_d,\] and \[\frac{\partial F_3}{\partial \boldsymbol{x}}(\boldsymbol{z}^\star) = \beta \nabla \boldsymbol{q}(\boldsymbol{x}^{\star}).\]

Collecting all block derivatives, we conclude that \[\boldsymbol{J}=DF(\boldsymbol{z}^\star)= \begin{pmatrix} -\lambda \boldsymbol{I}_d & -\boldsymbol{D}^\star & \boldsymbol{C}^\star\\[2mm] \alpha \nabla^2 f(\boldsymbol{x}^\star) & -\alpha \boldsymbol{I}_d & \boldsymbol{0}\\[2mm] \beta\nabla \boldsymbol{q}(\boldsymbol{x}^\star) & \boldsymbol{0}& -\beta \boldsymbol{I}_d \end{pmatrix}.\]

10.3.0.3 Hurwitz property.

Now we prove the Hurwitzness of the Jacobian matrix \(\boldsymbol{J}\), under certain conditions. We decompose \(\boldsymbol{J}\) as \[\boldsymbol{J}=\boldsymbol{J}_0+\boldsymbol{E},\] where \[\label{eq:J095adamw} \boldsymbol{J}_0= \begin{pmatrix} -\lambda \boldsymbol{I}_d & -\boldsymbol{D}^\star & \boldsymbol{0}\\[1mm] \alpha \nabla^2 f(\boldsymbol{x}^\star) & -\alpha \boldsymbol{I}_d & \boldsymbol{0}\\[1mm] \beta \nabla \boldsymbol{q}(\boldsymbol{x}^\star) & \boldsymbol{0}& -\beta \boldsymbol{I}_d \end{pmatrix},\tag{37}\] and \[\label{eq:E95adamw} \boldsymbol{E}= \begin{pmatrix} \boldsymbol{0}& \boldsymbol{0}& \boldsymbol{C}^\star\\[1mm] \boldsymbol{0}& \boldsymbol{0}& \boldsymbol{0}\\[1mm] \boldsymbol{0}& \boldsymbol{0}& \boldsymbol{0} \end{pmatrix}.\tag{38}\]

We first show that \(\boldsymbol{J}_0\) is Hurwitz. Observe that \(\boldsymbol{J}_0\) is block lower triangular, so its eigenvalues are the union of those of \[\boldsymbol{M}= \begin{pmatrix} -\lambda \boldsymbol{I}_d & -\boldsymbol{D}^\star\\ \alpha \nabla^2 f(\boldsymbol{x}^\star) & -\alpha \boldsymbol{I}_d \end{pmatrix}\] and those of \(-\beta \boldsymbol{I}_d\). Since \(\beta>0\), the block \(-\beta \boldsymbol{I}_d\) is clearly Hurwitz. It therefore suffices to show that \(\boldsymbol{M}\) is Hurwitz. The proof is similar to that for 3.

Let \(\mu\in\mathbb{C}\) be an eigenvalue of \(\boldsymbol{M}\) with corresponding eigenvector \((\boldsymbol{u}^\top,\boldsymbol{v}^\top)^\top\neq \boldsymbol{0}\). Then \[\begin{pmatrix} -\lambda \boldsymbol{I}_d & -\boldsymbol{D}^\star\\ \alpha \nabla^2 f(\boldsymbol{x}^\star) & -\alpha \boldsymbol{I}_d \end{pmatrix} \begin{pmatrix} \boldsymbol{u}\\ \boldsymbol{v} \end{pmatrix} = \mu \begin{pmatrix} \boldsymbol{u}\\ \boldsymbol{v} \end{pmatrix},\] which is equivalent to \[\begin{align} -(\lambda+\mu)\boldsymbol{u}-\boldsymbol{D}^\star\boldsymbol{v}&=\boldsymbol{0}, \tag{39}\\ \alpha \nabla^2 f(\boldsymbol{x}^\star)\boldsymbol{u}-(\alpha+\mu)\boldsymbol{v}&=\boldsymbol{0}. \tag{40} \end{align}\] From 40 , we obtain \[\boldsymbol{v}=\frac{\alpha}{\alpha+\mu}\nabla^2 f(\boldsymbol{x}^\star)\boldsymbol{u}.\] Substituting this into 39 , we get \[(\mu+\lambda)\boldsymbol{u}+\frac{\alpha}{\alpha+\mu}\boldsymbol{D}^\star \nabla^2 f(\boldsymbol{x}^\star)\boldsymbol{u}=\boldsymbol{0}.\] Equivalently, \[(\mu+\lambda)(\mu+\alpha)\boldsymbol{u}+\alpha \boldsymbol{D}^\star \nabla^2 f(\boldsymbol{x}^\star)\boldsymbol{u}=\boldsymbol{0}.\]

Now define \[\boldsymbol{H}^\star:=\nabla^2 f(\boldsymbol{x}^\star), \qquad \boldsymbol{B}:=(\boldsymbol{D}^\star)^{1/2}\boldsymbol{H}^\star(\boldsymbol{D}^\star)^{1/2}.\] Since \(\boldsymbol{D}^\star\succ \boldsymbol{0}\) and \(\boldsymbol{H}^\star\succ \boldsymbol{0}\), we have \(\boldsymbol{B}\succ \boldsymbol{0}\). Hence every eigenvalue \(\theta\) of \(\boldsymbol{B}\) satisfies \(\theta>0\). It follows that every eigenvalue \(\mu\) of \(\boldsymbol{M}\) must satisfy \[(\mu+\lambda)(\mu+\alpha)+\alpha\theta=0,\] that is, \[\mu^2+(\alpha+\lambda)\mu+\alpha(\lambda+\theta)=0, \qquad \theta>0.\] Since all coefficients are strictly positive, both roots have strictly negative real parts. Therefore, every eigenvalue of \(\boldsymbol{M}\) lies in the open left half-plane, and thus \(\boldsymbol{M}\) is Hurwitz. Consequently, \(\boldsymbol{J}_0\) is Hurwitz.

Since \(\boldsymbol{J}_0\) is Hurwitz, by the Lyapunov theorem there exists a symmetric positive definite matrix \(\boldsymbol{P}\succ \boldsymbol{0}\) such that \[\boldsymbol{J}_0^\top \boldsymbol{P}+\boldsymbol{P}\boldsymbol{J}_0=-\boldsymbol{I}.\] For the full matrix \(\boldsymbol{J}=\boldsymbol{J}_0+\boldsymbol{E}\), we have \[\boldsymbol{J}^\top \boldsymbol{P}+\boldsymbol{P}\boldsymbol{J} = -\boldsymbol{I}+\boldsymbol{E}^\top \boldsymbol{P}+\boldsymbol{P}\boldsymbol{E}.\] Thus, if \[\|\boldsymbol{E}^\top \boldsymbol{P}+\boldsymbol{P}\boldsymbol{E}\|<1,\] then \(\boldsymbol{J}^\top \boldsymbol{P}+\boldsymbol{P}\boldsymbol{J}\prec \boldsymbol{0}\), and therefore \(\boldsymbol{J}\) is Hurwitz.

Using the submultiplicative property of the operator norm, \[\|\boldsymbol{E}^\top \boldsymbol{P}+\boldsymbol{P}\boldsymbol{E}\| \le 2\|\boldsymbol{P}\|\,\|\boldsymbol{E}\| = 2\|\boldsymbol{P}\|\,\|\boldsymbol{C}^\star\|.\] Hence it suffices to require \[2\|\boldsymbol{P}\|\,\|\boldsymbol{C}^\star\|<1.\]

It remains to estimate \(\|\boldsymbol{C}^\star\|\). Since \[\boldsymbol{m}^\star=\nabla f(\boldsymbol{x}^\star)\] and \(v_i^\star\ge \underline v>0\), we have \[|C_{ii}^\star| = \left| \frac{m_i^\star}{2\sqrt{v_i^\star}(\sqrt{v_i^\star}+\varepsilon)^2} \right| \le \frac{\|\nabla f(\boldsymbol{x}^\star)\|_\infty}{2\sqrt{\underline v}(\sqrt{\underline v}+\varepsilon)^2}.\] Therefore, \[\|\boldsymbol{C}^\star\| \le \frac{\|\nabla f(\boldsymbol{x}^\star)\|_\infty}{2\sqrt{\underline v}(\sqrt{\underline v}+\varepsilon)^2}.\] It follows that if \[\|\nabla f(\boldsymbol{x}^\star)\|_\infty < \frac{\sqrt{\underline v}(\sqrt{\underline v}+\varepsilon)^2}{\|\boldsymbol{P}\|},\] then \[2\|\boldsymbol{P}\|\,\|\boldsymbol{C}^\star\|<1,\] and hence \(\boldsymbol{J}\) is Hurwitz.

10.4 Proof of 2↩︎

We provide a proof of the consistency of the batch-means estimator used in 2. The argument follows the standard batch-means variance-estimation principle for dependent stochastic processes. We state the regularity conditions explicitly for completeness.

Let \(s_t=s_{\ell,j}(t)\) be the empirical score process for a fixed LoRA component and let \(s^\star=s_{\ell,j}^{\star}\). Define \[X_t := s_t-s^\star .\] We assume that the centered score process satisfies a functional central limit theorem with long-run variance \(\bar\sigma_{\ell,j}^{2}\), that is, \[\frac{1}{\sqrt{t}}\sum_{\tau=1}^{\lfloor ct\rfloor} X_\tau \Rightarrow \bar\sigma_{\ell,j} B(c), \qquad c\in[0,1],\] where \(B(\cdot)\) is a standard Brownian motion. In addition, we assume the usual moment and weak-dependence conditions for non-overlapping batch-means consistency. Equivalently, for batch length \(b\to\infty\), number of batches \(M\to\infty\), and \(t=Mb\) with \(b/t\to0\), the block variables \[Y_{m,b} := \frac{1}{\sqrt b} \sum_{\tau=(m-1)b+1}^{mb} X_\tau, \qquad m=1,\ldots,M,\] satisfy \[\frac{1}{M}\sum_{m=1}^M Y_{m,b}^2 \xrightarrow{p} \bar\sigma_{\ell,j}^{2}, \qquad \frac{1}{M}\sum_{m=1}^M Y_{m,b} \xrightarrow{p} 0 .\] These conditions hold under standard mixing and moment assumptions used for batch-means estimators of long-run variances.

Recall that \[\bar s_{m,b} = \frac{1}{b}\sum_{\tau=(m-1)b+1}^{mb}s_\tau, \qquad \bar s_t = \frac{1}{t}\sum_{\tau=1}^{t}s_\tau .\] Since subtracting the constant \(s^\star\) does not change centered batch differences, we have \[\bar s_{m,b}-\bar s_t = \bar X_{m,b}-\bar X_t,\] where \[\bar X_{m,b} = \frac{1}{b}\sum_{\tau=(m-1)b+1}^{mb}X_\tau, \qquad \bar X_t = \frac{1}{t}\sum_{\tau=1}^{t}X_\tau .\] By the definition of \(Y_{m,b}\), \[\bar X_{m,b} = \frac{1}{\sqrt b}Y_{m,b}.\] Moreover, \[\bar X_t = \frac{1}{M}\sum_{m=1}^M \bar X_{m,b} = \frac{1}{\sqrt b} \left( \frac{1}{M}\sum_{m=1}^M Y_{m,b} \right).\] Therefore, \[\bar s_{m,b}-\bar s_t = \frac{1}{\sqrt b} \left( Y_{m,b} - \bar Y_{M,b} \right), \qquad \bar Y_{M,b} := \frac{1}{M}\sum_{m=1}^M Y_{m,b}.\] Substituting this expression into the batch-means estimator gives \[\begin{align} \widehat{\bar\sigma}_{\ell,j}^{2} &= \frac{b}{M-1} \sum_{m=1}^M \left( \bar s_{m,b}-\bar s_t \right)^2 \\ &= \frac{1}{M-1} \sum_{m=1}^M \left( Y_{m,b}-\bar Y_{M,b} \right)^2 . \end{align}\] Expanding the centered empirical second moment, we obtain \[\begin{align} \widehat{\bar\sigma}_{\ell,j}^{2} &= \frac{M}{M-1} \left[ \frac{1}{M}\sum_{m=1}^M Y_{m,b}^2 - \left( \frac{1}{M}\sum_{m=1}^M Y_{m,b} \right)^2 \right]. \end{align}\] By the batch-means regularity conditions, \[\frac{1}{M}\sum_{m=1}^M Y_{m,b}^2 \xrightarrow{p} \bar\sigma_{\ell,j}^{2}, \qquad \left( \frac{1}{M}\sum_{m=1}^M Y_{m,b} \right)^2 \xrightarrow{p} 0 .\] Since \(M/(M-1)\to1\), Slutsky’s theorem yields \[\widehat{\bar\sigma}_{\ell,j}^{2} \xrightarrow{p} \bar\sigma_{\ell,j}^{2}.\] This proves the consistency of the batch-means variance estimator.

10.5 Proof for 3↩︎

We prove the two statements separately.

10.5.0.1 Characterization of the limiting point.

Recall that the mean-field mapping for Adam is given by \[F(\boldsymbol{z})= \begin{pmatrix} -\dfrac{\boldsymbol{m}}{\sqrt{\boldsymbol{v}}+\varepsilon}\\[1mm] \alpha\bigl(\nabla f(\boldsymbol{x})-\boldsymbol{m}\bigr)\\[1mm] \beta\bigl(\boldsymbol{q}(\boldsymbol{x})-\boldsymbol{v}\bigr) \end{pmatrix}, \qquad \boldsymbol{z}=(\boldsymbol{x},\boldsymbol{m},\boldsymbol{v})\in\mathbb{R}^{3d},\] where the division is understood componentwise.

A solution \(\boldsymbol{z}^\star=(\boldsymbol{x}^\star,\boldsymbol{m}^\star,\boldsymbol{v}^\star)\) of the mean-field function satisfies \[F(\boldsymbol{z}^\star)=\boldsymbol{0}.\] Hence its three blocks must satisfy \[\begin{align} -\frac{\boldsymbol{m}^\star}{\sqrt{\boldsymbol{v}^\star}+\varepsilon} &= \boldsymbol{0}, \tag{41}\\ \alpha\bigl(\nabla f(\boldsymbol{x}^\star)-\boldsymbol{m}^\star\bigr) &= \boldsymbol{0}, \tag{42}\\ \beta\bigl(\boldsymbol{q}(\boldsymbol{x}^{\star})-\boldsymbol{v}^\star\bigr) &= \boldsymbol{0}. \tag{43} \end{align}\] Since \(\alpha,\beta>0\) and \(\sqrt{\boldsymbol{v}^\star}+\varepsilon\) is strictly positive componentwise, 41 implies \[\boldsymbol{m}^\star=\boldsymbol{0}.\] Substituting this into 42 yields \[\nabla f(\boldsymbol{x}^\star)=\boldsymbol{0},\] so that \(\boldsymbol{x}^\star\) is a stationary point of the objective function. Finally, 43 gives \[\boldsymbol{v}^\star=\boldsymbol{q}(\boldsymbol{x}^\star) = \boldsymbol{q}^{\star}.\] Therefore, every solution is of the form \[\boldsymbol{z}^\star=(\boldsymbol{x}^\star,\boldsymbol{0},\boldsymbol{q}^{\star}), \qquad \nabla f(\boldsymbol{x}^\star)=\boldsymbol{0}.\]

10.5.0.2 Jacobian matrix at the limiting point.

We now compute the Jacobian matrix \(\boldsymbol{J}=DF(\boldsymbol{z}^\star)\) at the limiting point \[\boldsymbol{z}^\star=(\boldsymbol{x}^\star,\boldsymbol{0},\boldsymbol{q}^{\star}), \qquad \nabla f(\boldsymbol{x}^\star)=\boldsymbol{0}.\] Write \[F(\boldsymbol{z})= \begin{pmatrix} F_1(\boldsymbol{x},\boldsymbol{m},\boldsymbol{v})\\[1mm] F_2(\boldsymbol{x},\boldsymbol{m},\boldsymbol{v})\\[1mm] F_3(\boldsymbol{x},\boldsymbol{m},\boldsymbol{v}) \end{pmatrix},\] where \[\begin{align} F_1(\boldsymbol{x},\boldsymbol{m},\boldsymbol{v}) &= -\frac{\boldsymbol{m}}{\sqrt{\boldsymbol{v}}+\varepsilon},\\ F_2(\boldsymbol{x},\boldsymbol{m},\boldsymbol{v}) &= \alpha\bigl(\nabla f(\boldsymbol{x})-\boldsymbol{m}\bigr),\\ F_3(\boldsymbol{x},\boldsymbol{m},\boldsymbol{v}) &= \beta\bigl(\boldsymbol{q}(\boldsymbol{x})-\boldsymbol{v}\bigr). \end{align}\]

We compute each block derivative.

Derivative of the first block. Since \(F_1\) does not depend explicitly on \(\boldsymbol{x}\), we have \[\frac{\partial F_1}{\partial \boldsymbol{x}}(\boldsymbol{z}^\star)=\boldsymbol{0}.\] Differentiating with respect to \(\boldsymbol{m}\) gives \[\frac{\partial F_1}{\partial \boldsymbol{m}}(\boldsymbol{z}^\star) = -\operatorname{diag}\!\left(\frac{1}{\sqrt{\boldsymbol{q}^\star}+\varepsilon}\right).\] Finally, differentiating \(F_1\) with respect to \(\boldsymbol{v}\) produces terms proportional to \(\boldsymbol{m}\), and therefore \[\frac{\partial F_1}{\partial \boldsymbol{v}}(\boldsymbol{z}^\star)=\boldsymbol{0}\] since \(\boldsymbol{m}^\star=\boldsymbol{0}\).

Derivative of the second block. For \[F_2(\boldsymbol{x},\boldsymbol{m},\boldsymbol{v})=\alpha\bigl(\nabla f(\boldsymbol{x})-\boldsymbol{m}\bigr),\] we immediately obtain \[\frac{\partial F_2}{\partial \boldsymbol{x}}(\boldsymbol{z}^\star)=\alpha \nabla^2 f(\boldsymbol{x}^\star), \qquad \frac{\partial F_2}{\partial \boldsymbol{m}}(\boldsymbol{z}^\star)=-\alpha \boldsymbol{I}_d, \qquad \frac{\partial F_2}{\partial \boldsymbol{v}}(\boldsymbol{z}^\star)=\boldsymbol{0}.\]

Derivative of the third block. For \[F_3(\boldsymbol{x},\boldsymbol{m},\boldsymbol{v})=\beta\bigl(\boldsymbol{q}(\boldsymbol{x})-\boldsymbol{v}\bigr),\] we have \[\frac{\partial F_3}{\partial \boldsymbol{m}}(\boldsymbol{z}^\star)=\boldsymbol{0}, \qquad \frac{\partial F_3}{\partial \boldsymbol{v}}(\boldsymbol{z}^\star)=-\beta \boldsymbol{I}_d.\] Moreover, the derivative of \(\boldsymbol{q}(\boldsymbol{x})\) with respect to \(\boldsymbol{x}\) is \(\beta\nabla \boldsymbol{q}(\boldsymbol{x}^{\star})\).

Collecting all the block derivatives, we conclude that \[\boldsymbol{J}= DF(\boldsymbol{z}^\star) = \begin{pmatrix} \boldsymbol{0}& - \operatorname{diag}\!\bigl(\frac{1}{\sqrt{\boldsymbol{q}^{\star}} + \varepsilon}\bigr) & \boldsymbol{0}\\ \alpha \nabla^2 f(\boldsymbol{x}^\star) & -\alpha \boldsymbol{I}_d & \boldsymbol{0}\\ \beta \nabla \boldsymbol{q}(\boldsymbol{x}^{\star}) & \boldsymbol{0}& -\beta \boldsymbol{I}_d \end{pmatrix},\]

10.5.0.3 Hurwitz property.

Write \(\boldsymbol{J}\) in the block form \[\boldsymbol{J}= \begin{pmatrix} \boldsymbol{M}& \boldsymbol{0}\\ \boldsymbol{N}& -\beta \boldsymbol{I}_d \end{pmatrix},\] where \[\boldsymbol{M}= \begin{pmatrix} \boldsymbol{0}& -\boldsymbol{D}^\star\\ \alpha \nabla^2 f(\boldsymbol{x}^\star) & -\alpha \boldsymbol{I}_d \end{pmatrix}, \qquad \boldsymbol{N}= \begin{pmatrix} \beta \nabla \boldsymbol{q}(\boldsymbol{x}^\star) & \boldsymbol{0} \end{pmatrix}.\] Since \(\boldsymbol{J}\) is block lower triangular, its eigenvalues are the union of those of \(\boldsymbol{M}\) and those of \(-\beta \boldsymbol{I}_d\). The latter are all equal to \(-\beta<0\). Therefore, it remains to show that \(\boldsymbol{M}\) is Hurwitz.

Let \(\mu\in\mathbb{C}\) be an eigenvalue of \(\boldsymbol{M}\) with corresponding eigenvector \((\boldsymbol{u}^\top,\boldsymbol{v}^\top)^\top\neq \boldsymbol{0}\). Then \[\begin{pmatrix} \boldsymbol{0}& -\boldsymbol{D}^\star\\ \alpha \nabla^2 f(\boldsymbol{x}^\star) & -\alpha \boldsymbol{I}_d \end{pmatrix} \begin{pmatrix} \boldsymbol{u}\\ \boldsymbol{v} \end{pmatrix} = \mu \begin{pmatrix} \boldsymbol{u}\\ \boldsymbol{v} \end{pmatrix},\] which yields \[\begin{align} -\boldsymbol{D}^\star \boldsymbol{v}&= \mu \boldsymbol{u}, \tag{44}\\ \alpha \nabla^2 f(\boldsymbol{x}^\star)\boldsymbol{u}- \alpha \boldsymbol{v}&= \mu \boldsymbol{v}. \tag{45} \end{align}\] From 44 , we obtain \[\boldsymbol{v}=-\mu (\boldsymbol{D}^\star)^{-1}\boldsymbol{u}.\] Substituting this into 45 gives \[\alpha \nabla^2 f(\boldsymbol{x}^\star)\boldsymbol{u}+ \alpha\mu (\boldsymbol{D}^\star)^{-1}\boldsymbol{u} = -\mu^2 (\boldsymbol{D}^\star)^{-1}\boldsymbol{u}.\] Equivalently, \[\alpha \nabla^2 f(\boldsymbol{x}^\star)\boldsymbol{u}+ (\alpha\mu+\mu^2)(\boldsymbol{D}^\star)^{-1}\boldsymbol{u}=\boldsymbol{0}.\]

Now define \[\boldsymbol{w}:=(\boldsymbol{D}^\star)^{-1/2}\boldsymbol{u}, \qquad \boldsymbol{B}:=(\boldsymbol{D}^\star)^{1/2}\nabla^2 f(\boldsymbol{x}^\star)(\boldsymbol{D}^\star)^{1/2}.\] Then the above equation becomes \[\alpha \boldsymbol{B}\boldsymbol{w}+ (\alpha\mu+\mu^2)\boldsymbol{w}=\boldsymbol{0}.\] Since \(\nabla^2 f(\boldsymbol{x}^\star)\succ \boldsymbol{0}\) and \(\boldsymbol{D}^\star\succ \boldsymbol{0}\), it follows that \(\boldsymbol{B}\succ \boldsymbol{0}\). Hence every eigenvalue \(\theta\) of \(\boldsymbol{B}\) is strictly positive, and \(\mu\) must satisfy \[\mu^2+\alpha\mu+\alpha\theta=0, \qquad \theta>0.\] Because both coefficients \(\alpha\) and \(\alpha\theta\) are strictly positive, both roots of this quadratic have strictly negative real parts. Therefore, every eigenvalue of \(\boldsymbol{M}\) lies in the open left half-plane, and hence \(\boldsymbol{M}\) is Hurwitz.

Combining this with the fact that \(-\beta \boldsymbol{I}_d\) is Hurwitz, we conclude that \(\boldsymbol{J}\) is Hurwitz.

10.6 Proof for 4↩︎

We prove the two statements separately.

10.6.0.1 Characterization of the limiting point.

Recall that the mean-field mapping for AdaFactor is \[F(\boldsymbol{z})= \begin{pmatrix} -\dfrac{\nabla f(\boldsymbol{x})}{\sqrt{\boldsymbol{v}}+\varepsilon}\\[2mm] \beta\bigl(\boldsymbol{q}(\boldsymbol{x})-\boldsymbol{v}\bigr) \end{pmatrix}, \qquad \boldsymbol{z}=(\boldsymbol{x},\boldsymbol{v})\in\mathbb{R}^{2d},\] where the division is understood componentwise.

Let \(\boldsymbol{z}^\star=(\boldsymbol{x}^\star,\boldsymbol{v}^\star)\) be a solution of the mean-field function. Then \[F(\boldsymbol{z}^\star)=\boldsymbol{0},\] which is equivalent to \[\begin{align} -\frac{\nabla f(\boldsymbol{x}^\star)}{\sqrt{\boldsymbol{v}^\star}+\varepsilon} &= \boldsymbol{0}, \tag{46}\\ \beta\bigl(\boldsymbol{q}(\boldsymbol{x}^\star)-\boldsymbol{v}^\star\bigr) &= \boldsymbol{0}. \tag{47} \end{align}\] Since \(\varepsilon>0\), the vector \(\sqrt{\boldsymbol{v}^\star}+\varepsilon\) is strictly positive componentwise. Therefore, 46 implies \[\nabla f(\boldsymbol{x}^\star)=\boldsymbol{0}.\] Substituting this identity into 47 , and using \(\beta>0\), we obtain \[\boldsymbol{v}^\star=\boldsymbol{q}(\boldsymbol{x}^\star) = \boldsymbol{q}^\star.\] Hence, every solution is of the form \[\boldsymbol{z}^\star=(\boldsymbol{x}^\star,\boldsymbol{v}^\star), \quad \nabla f(\boldsymbol{x}^\star)=\boldsymbol{0}, \quad \boldsymbol{v}^\star = \boldsymbol{q}^\star.\]

10.6.0.2 Computation of the Jacobian matrix.

We now compute the Jacobian matrix \(\boldsymbol{J}=DF(\boldsymbol{z}^\star)\) at the limiting point \[\boldsymbol{z}^\star=(\boldsymbol{x}^\star,\boldsymbol{v}^\star), \qquad \nabla f(\boldsymbol{x}^\star)=\boldsymbol{0}.\] Write \[F(\boldsymbol{z})= \begin{pmatrix} F_1(\boldsymbol{x},\boldsymbol{v})\\[1mm] F_2(\boldsymbol{x},\boldsymbol{v}) \end{pmatrix},\] where \[\begin{align} F_1(\boldsymbol{x},\boldsymbol{v}) &= -\frac{\nabla f(\boldsymbol{x})}{\sqrt{\boldsymbol{v}}+\varepsilon},\\ F_2(\boldsymbol{x},\boldsymbol{v}) &= \beta\bigl(\boldsymbol{q}(\boldsymbol{x})-\boldsymbol{v}\bigr). \end{align}\]

We compute each block derivative.

Derivative of the first block. Since \[F_1(\boldsymbol{x},\boldsymbol{v}) = -\operatorname{diag}\!\left(\frac{1}{\sqrt{\boldsymbol{v}}+\varepsilon}\right)\nabla f(\boldsymbol{x}),\] we have \[\frac{\partial F_1}{\partial \boldsymbol{x}}(\boldsymbol{z}^\star) = -\operatorname{diag}\!\left(\frac{1}{\sqrt{\boldsymbol{v}^\star}+\varepsilon}\right) \nabla^2 f(\boldsymbol{x}^\star).\]

Next, we compute the derivative with respect to \(\boldsymbol{v}\). Componentwise, \[[F_1]_i(\boldsymbol{x},\boldsymbol{v}) = -\frac{\partial_i f(\boldsymbol{x})}{\sqrt{v_i}+\varepsilon}.\] Hence \[\frac{\partial [F_1]_i}{\partial v_i}(\boldsymbol{z}^\star) = \frac{\partial_i f(\boldsymbol{x}^\star)}{2\sqrt{v_i^\star}(\sqrt{v_i^\star}+\varepsilon)^2}.\] Since \(\nabla f(\boldsymbol{x}^\star)=\boldsymbol{0}\), we conclude that \[\frac{\partial F_1}{\partial \boldsymbol{v}}(\boldsymbol{z}^\star)=\boldsymbol{0}.\]

Derivative of the second block. For \[F_2(\boldsymbol{x},\boldsymbol{v})=\beta\bigl(\boldsymbol{q}(\boldsymbol{x})-\boldsymbol{v}\bigr),\] we have \[\frac{\partial F_2}{\partial \boldsymbol{x}}(\boldsymbol{z}^\star) = \beta \nabla \boldsymbol{q}(\boldsymbol{x}^\star),\] and \[\frac{\partial F_2}{\partial \boldsymbol{v}}(\boldsymbol{z}^\star) = -\beta\boldsymbol{I}_d.\]

Collecting all block derivatives, we obtain \[\boldsymbol{J}=DF(\boldsymbol{z}^\star) = \begin{pmatrix} -\operatorname{diag}\!\left(\frac{1}{\sqrt{\boldsymbol{v}^\star}+\varepsilon}\right) \nabla^2 f(\boldsymbol{x}^\star) & \boldsymbol{0}\\[2mm] \beta \nabla \boldsymbol{q}(\boldsymbol{x}^\star) & -\beta \boldsymbol{I}_d \end{pmatrix}.\]

10.6.0.3 Hurwitz property.

The matrix \(\boldsymbol{J}\) is block lower triangular. Therefore, its eigenvalues are the union of those of \[-\operatorname{diag}\!\left(\frac{1}{\sqrt{\boldsymbol{v}^\star}+\varepsilon}\right)\nabla^2 f(\boldsymbol{x}^\star)\] and those of \[-\beta \boldsymbol{I}_d.\] Since \(\beta>0\), the block \(-\beta \boldsymbol{I}_d\) is Hurwitz.

Now assume that \(\nabla^2 f(\boldsymbol{x}^\star)\succ \boldsymbol{0}\). Since \(\varepsilon>0\), the diagonal matrix \[\boldsymbol{D}^\star:=\operatorname{diag}\!\left(\frac{1}{\sqrt{\boldsymbol{v}^\star}+\varepsilon}\right)\] is positive definite. Then \[\boldsymbol{D}^\star \nabla^2 f(\boldsymbol{x}^\star)\] has strictly positive eigenvalues, because it is similar to \[(\boldsymbol{D}^\star)^{1/2}\nabla^2 f(\boldsymbol{x}^\star)(\boldsymbol{D}^\star)^{1/2}\succ \boldsymbol{0}.\] Hence \[-\boldsymbol{D}^\star \nabla^2 f(\boldsymbol{x}^\star)\] has strictly negative eigenvalues. Therefore, the upper-left block is Hurwitz. Since both diagonal blocks are Hurwitz, \(\boldsymbol{J}\) is Hurwitz.

11 Proofs for Some Technical Lemmas↩︎

11.1 Proof for 2↩︎

11.2 Proof for 3↩︎

Since \(\boldsymbol{J}\) is Hurwitz, the Lyapunov equation \[\boldsymbol{J}\boldsymbol{\Sigma}+\boldsymbol{\Sigma}\boldsymbol{J}^\top + \boldsymbol{Q}= \boldsymbol{0}\] admits a unique positive semidefinite solution \(\boldsymbol{\Sigma}\). We show that \(\boldsymbol{U}_t\to \boldsymbol{\Sigma}\).

Step 1: a recursion for \(\boldsymbol{U}_t\).

By definition, \[\boldsymbol{U}_{t+1} = \sum_{k=1}^{t+1} \frac{\gamma_k^2}{\gamma_{t+1}}\, \Phi_{t+1,k+1}\boldsymbol{Q}\Phi_{t+1,k+1}^\top.\] Separating the last term \(k=t+1\), and noting that \(\Phi_{t+1,t+2}=\boldsymbol{I}\), we obtain \[\boldsymbol{U}_{t+1} = \frac{\gamma_t}{\gamma_{t+1}} (\boldsymbol{I}+\gamma_{t+1}\boldsymbol{J}) \left( \sum_{k=1}^{t} \frac{\gamma_k^2}{\gamma_t}\, \Phi_{t,k+1}\boldsymbol{Q}\Phi_{t,k+1}^\top \right) (\boldsymbol{I}+\gamma_{t+1}\boldsymbol{J})^\top + \gamma_{t+1}\boldsymbol{Q}.\] Hence \[\label{eq7} \boldsymbol{U}_{t+1} = \frac{\gamma_t}{\gamma_{t+1}} (\boldsymbol{I}+\gamma_{t+1}\boldsymbol{J})\boldsymbol{U}_t(\boldsymbol{I}+\gamma_{t+1}\boldsymbol{J})^\top + \gamma_{t+1}\boldsymbol{Q}.\tag{49}\]

Step 2: expand the ratio \(\gamma_t/\gamma_{t+1}\).

Since \[\gamma_t=\gamma_0 t^{-\kappa}, \qquad \kappa\in(1/2,1),\] we have \[\frac{\gamma_t}{\gamma_{t+1}} = \left(\frac{t+1}{t}\right)^\kappa = 1+\frac{\kappa}{t}+o(t^{-1}).\] Because \(\kappa<1\), we have \[\frac{1}{t}=o(\gamma_t),\] and thus \[\frac{\gamma_t}{\gamma_{t+1}} = 1+o(\gamma_t).\]

Therefore, 49 implies that \[\begin{align} \boldsymbol{U}_{t+1} &= (1+o(\gamma_t)) (\boldsymbol{I}+\gamma_{t+1}\boldsymbol{J})\boldsymbol{U}_t(\boldsymbol{I}+\gamma_{t+1}\boldsymbol{J})^\top +\gamma_{t+1}\boldsymbol{Q}. \end{align}\] and expanding the product yields \[\begin{align} \boldsymbol{U}_{t+1} &= \boldsymbol{U}_t +\gamma_t(\boldsymbol{J}\boldsymbol{U}_t+\boldsymbol{U}_t\boldsymbol{J}^\top + \boldsymbol{Q}) +\boldsymbol{R}_t, \end{align}\] where the residual term satisfies \[\|\boldsymbol{R}_t\| \le o(\gamma_t)(1+\|\boldsymbol{U}_t\|).\] Note that \[\left\|\boldsymbol{U}_t \right\| \leq C\gamma_t^{-1} \sum_{k=1}^{t} \prod_{j=k+1}^{t}\left(1 - c\gamma_j\right)^2 \gamma_k^2 = \mathcal{O}\left(1\right),\] we have \[\|\boldsymbol{R}_t\| \le o(\gamma_t).\]

Step 3: compare \(\boldsymbol{U}_t\) with the Lyapunov solution \(\boldsymbol{\Sigma}\).

Define \[\boldsymbol{E}_t:=\boldsymbol{U}_t-\boldsymbol{\Sigma}.\] Since \(\boldsymbol{\Sigma}\) solves \[\boldsymbol{J}\boldsymbol{\Sigma}+\boldsymbol{\Sigma}\boldsymbol{J}^\top+\boldsymbol{Q}=\boldsymbol{0},\] we get \[\label{eq8} \begin{align} \boldsymbol{E}_{t+1} &= \boldsymbol{E}_t+\gamma_t(\boldsymbol{J}\boldsymbol{E}_t+\boldsymbol{E}_t\boldsymbol{J}^\top)+\boldsymbol{R}_t\\ & = \left( \boldsymbol{I}+ \gamma_t\boldsymbol{J}\right)\boldsymbol{E}_t\left( \boldsymbol{I}+ \gamma_t\boldsymbol{J}^{\top}\right) + \widetilde{\boldsymbol{R}}_{t}, \end{align}\tag{50}\] where \(\widetilde{\boldsymbol{R}}_{t} = \boldsymbol{R}_t - \gamma_t^2 \boldsymbol{J}\boldsymbol{E}_t\boldsymbol{J}^{\top}\) and thus \(\left\|\widetilde{\boldsymbol{R}}_{t} \right\| = o(\gamma_t)\).

Step 4: stability of the homogeneous part.

Applying the recursion for \(\boldsymbol{E}_t\) in 50 multiple times, we obtain \[\boldsymbol{E}_t = \Phi_{t,1}\boldsymbol{E}_1\Phi_{t,1}^{\top} + \sum_{k=1}^{t-1}\Phi_{t,k+1}\widetilde{\boldsymbol{R}}_k\Phi_{t,k+1}^{\top}.\] The first term vanishes because of the exponential stability: \[\Phi_{t,1}\boldsymbol{E}_1\Phi_{t,1}^{\top} \to \boldsymbol{0}.\] For the second term, since \(\left\|\widetilde{\boldsymbol{R}}_k\right\|=o(\gamma_k)\), we have \[\begin{align} & \left\| \sum_{k=1}^{t-1}\Phi_{t,k+1}\widetilde{\boldsymbol{R}}_k\Phi_{t,k+1}^{\top}\right\|\\ & \leq \sum_{k=1}^{t-1} \left\|\Phi_{t,k+1} \right\|^2\left\| \widetilde{\boldsymbol{R}}_k \right\|\\ & \leq C \sum_{k=1}^{t-1} \prod_{j=k+1}^{t-1} \left(1-c\gamma_j\right)^2 \left\|\widetilde{\boldsymbol{R}}_k\right\| = o(1). \end{align}\]

Hence \[\boldsymbol{E}_t\to \boldsymbol{0},\] that is, \[\boldsymbol{U}_t\to \boldsymbol{\Sigma}.\] This completes the proof.

References↩︎

---

1. Department of Mathematics, National University of Singapore. Email: gaoyh@connect.hku.hk↩︎

2. Department of Mathematics and Department of Electrical and Computer Engineering, National University of Singapore. Email: vtan@nus.edu.sg↩︎