FLOPS: Forward Learning with OPtimal Sampling

## Abstract

Abstract Given the limitations of backpropagation, perturbation-based gradient computation methods have recently gained focus for learning with only forward passes, also referred to as queries. Conventional forward learning consumes enormous queries on each data point for accurate gradient estimation through Monte Carlo sampling, which hinders the scalability of those algorithms. However, not all data points deserve equal queries for gradient estimation. In this paper, we study the problem of improving the forward learning efficiency from a novel perspective: how to reduce the gradient estimation variance with minimum cost? For this, we propose to allocate the optimal number of queries over each data in one batch during training to achieve a good balance between estimation accuracy and computational efficiency. Specifically, with a simplified proxy objective and a reparameterization technique, we derive a novel plug-and-play query allocator with minimal parameters. Theoretical results are carried out to verify its optimality. We conduct extensive experiments for fine-tuning Vision Transformers on various datasets and further deploy the allocator to two black-box applications: prompt tuning and multimodal alignment for foundation models. All findings demonstrate that our proposed allocator significantly enhances the scalability of forward-learning algorithms, paving the way for real-world applications.

## 1 Introduction

Ever since the success of backpropagation (BP) (Rumelhart et al., ), researchers have sought alternate biologically plausible methods as the brain doesn’t perform backward computation to learn high-level representation (Ma et al., ; Nøkland, ; Jacot et al., ). Moreover, existing practical scenarios urge for scalable forward learning methods, when the black-box nature renders taking the derivative infeasible or difficult. For example, as the model parameters increase, machine learning (ML) systems usually integrate with third-party black-box APIs (Achiam et al., ).

There have been continuous efforts to develop learning algorithms that rely solely on the forward pass (Lillicrap et al., ). The forward learning uses multiple queries on one data point to estimate the gradient (Spall, ; Zhang et al., ). The high dimensionality and the rugged loss landscape of ML problems pose arduous challenges for deriving a low-variance gradient estimator. Current Literature has demonstrated that a sufficient number of queries on each data point is necessary for successful training (Chen et al., ; Zhang et al., ). To achieve good gradient estimation and task performance, all existing methods equally assign a large number of queries to different data points, neglecting the varying difficulty of gradient estimation at these data.

However, it is not always helpful for the forward learning algorithms to increase the number of queries. Especially for the large model training on a large dataset, it usually requires an extensive number of queries, which poses great challenges to the memory and computation cost. For example, The forward-forward (FF) (Hinton, ) algorithm and the feedback alignment (FA) (Nøkland, ) are only capable of training multilayer perceptrons on MNIST (LeCun, ). To scale up the forward algorithm to large model training on large datasets, DeepZero (Chen et al., ) and Mezo (Malladi et al., ) explore the use of SPSA (Spall, ) in training the convolution network from scratch and fine-tune the language model, respectively. But with such “large” scale optimization problem, DeepZero requires $192$ queries to train a pruned ResNet-20 with only $12K$ parameters. Mezo needs meticulously selected prompts to ensure performance, and its gradient estimator is too noisy due to insufficient query numbers, thus suffering from highly unstable performance.

There is a great challenge required to be considered carefully in forward learning. On the one hand, consuming large amounts of queries on every data point results in high computational costs and long training time. On the other hand, using insufficient queries can save memory and computational costs but lead to unstable performance. Efficiently utilizing the queries with limited memory and computation budget is at the core of the scalability for forward learning to reduce the gradient estimation variance and improve the model training efficiency.

To achieve a good balance between the query cost and gradient estimation variance, we propose to study the query allocation problem under the forward learning framework. In our paper, we first unify the different estimators under the perturbation-based framework and formulate the allocation problems. Then, we propose the optimal allocation by constructing a query budget allocator to assign queries adaptively following a data-aware paradigm.
An overview of our method and the comparison with previous work can be found in Figure 1. Specifically, our contribution can be summarized as follows,

- 1.
We find a simplified objective for the allocation problem and solve it via the likelihood ratio method in a lightweight style. With the optimal allocation, we can compress the queries to the minimum number of $20$.
- 2.
We theoretically deduce that the enhancement in allocation is assured by a defined lower bound. The allocator would assign the queries for different data in a mini-batch according to the variance of the gradient estimation. Intuitively, the estimator with high variance would be assigned with more queries, and vice versa.
- 3.
Extensive experiments are conducted on fine-tuning Vision Transformer (Dosovitskiy, ) on various challenging datasets including the ImageNet-1K Deng et al. (), achieving the SOTA performance in the realm of forward learning.
- 4.
We leverage the proposed method to solve two real-world black-box problems, including black-box prompt tuning and alignment between foundation models, showing its scalability and efficiency improvement versus a series of optimization methods, including the BP.

Figure: Figure 1: Illustration of allocating the query budget of the forward learning paradigm. As shown in (a), previous methods equally allocate the query across different data. Our method, as shown in (b), adaptively allocate the queries under theoretically guaranteed optimality.
Refer to caption: x1.png

## 2 Related Work

Backpropagation-free learning. Without BP, people have designed algorithms drawing inspiration from biological social behavior and evolution theory Wierstra et al. (). Swarm intelligent algorithms, for example, particle swarm optimization (Vaz & Vicente, ), optimize a black-box problem by iteratively tracking the best candidate in the population. Evolutionary algorithms, for example, natural evolutionary strategies, inject random noise into the model and optimize along the optimal direction along the noise (Salimans et al., ). Both swarm intelligence algorithms and evolution algorithms lack sufficient theoretical analysis and the performance would drop as the parameters increase. Forward gradient (Baydin et al., ; Silver et al., ), applying the chain rule from a different direction than BP, utilizes forward-mode automatic differentiation (AD) to train the ML model. Forward-mode AD needs specific AD software and full access to the model structure. Therefore, training the deep model under black-box scenarios is not feasible by the forward gradient (Ren et al., ).

Perturbation-based forward learning. Besides the heuristic algorithms, researchers have devised surrogate methods to approximate the first-order gradients, such as SPSA. The SPSA injects noise with opposite signs into the model parameters to form a finite-difference style estimator. Recently, Mezo (Malladi et al., ) applied SPSA to fine-tune the language model. DeepZero (Chen et al., ) integrates SPSA with sparsity and trains the network from scratch with many queries. However, their methods still have some limitations and their query budgets are equally allocated, neglecting the data difference. Jiang et al. () extends the likelihood ratio estimator to train a wide range of neural networks, i.e. CNN (He et al., ), RNN (Dey & Salem, ), GNN (Scarselli et al., ), SNN (Ghosh-Dastidar & Adeli, ), whereas they don’t verify their method on modern architecture, i.e., Transformers (Vaswani, ). There is still a huge research gap for improving the scalability of forward learning.

Machine learning for black box scenarios.
The applications to real-world black-box scenarios call for scalable BP-free algorithms. When deploying ML algorithms on specific hardware systems, computation resources are limited and BP becomes infeasible. In the fields of physics and chemistry, ML models have to interact with environments that include non-differentiable operations (Momeni et al., ; Gu et al., ). Moreover, extremely large models are often only accessible through third-party API (Sun et al., ). Evolutionary algorithms and reinforcement learning have been applied to tune the prompt for the black-box language model (Diao et al., ).

## 3 Preliminary on Forward Learning

Given a generic neural network with no assumption on its architecture, the input is denoted as $x\in\mathbb{R}^{d_{x}}$, and the output is given by
$y=\phi_{2}(\varphi(\phi_{1}(x);\theta))\in\mathbb{R}^{d_{y}}$, where $\phi_{1}(\cdot)$ and $\phi_{2}(\cdot)$ are black-box parts, and $\varphi(\cdot;\theta)$ is the trainable structure parameterized by $\theta\in\Theta\subset\mathbb{R}^{d_{\theta}}$. Training the neural network is to solve such an optimization:
$\min_{\theta\in\Theta}\frac{1}{|\mathcal{X}|}\sum_{x\in\mathcal{X}}\mathcal{L}
(y)$, where $\mathcal{X}$ is the dataset for training, and $\mathcal{L}(\cdot)$ is an evaluation procedure yielding the loss feedback. Due to the growing scale of parameters in modern neural networks and the dataset, stochastic gradient descent has essentially become the only feasible solution method.
The primary challenge remains in computing the gradients for each batch efficiently.

In addition to approaches relying on computation graphs like BP, forward-learning algorithms estimate gradients by perturbing the neural activity and observing the correlation between the injected perturbation and the final loss value, thereby solving the stochastic version of the problem:
$\min_{\theta\in\Theta}\frac{1}{|\mathcal{X}|}\sum_{x\in\mathcal{X}}\mathbb{E}[
\mathcal{L}(y)]$.
For instance, Jiang et al. () perturb the intermediate result to perform LR estimation, i.e.,
$y=\phi_{2}(\varphi(\phi_{1}(x);\theta)+z)$, where $z$ is the injected noise with a density $f(\cdot)$. Denote the Jacobian matrix of $a\in\mathbb{R}^{d_{a}}$ with respect to $x\in\mathbb{R}^{d_{b}}$ as $D_{b}a\in\mathbb{R}^{d_{a}\times d_{b}}$. The loss gradient can be estimated by

$$ $\displaystyle\nabla_{\theta}\mathbb{E}[\mathcal{L}(y)]=\mathbb{E}[-\mathcal{L} (y)D_{\theta}^{\top}\varphi(\phi_{1}(x);\theta)\nabla_{\theta}\ln f(z)].$ (1) $$

The SPSA family injects noise to the parameters (Spall, ; Salimans et al., ), i.e., $y=\phi_{2}(\varphi(\phi_{1}(x);\theta+z))$.
For continuous noise distributions, the gradient estimation can be derived as a special case of equality (1) as follows:

$$ $\displaystyle\nabla_{\theta}\mathbb{E}[\mathcal{L}(y)]=\mathbb{E}[-\mathcal{L} (y)\nabla_{\theta}\ln f(z)].$ (2) $$

It can also be interpreted as a method for randomized finite differences in high-dimensional space if we use Gaussian noise $z\sim\mathcal{N}(0,\sigma^{2}I_{d})$ and apply an antithetic variable trick for variance reduction to equality (2), i.e.,

$$ $\displaystyle\nabla_{\theta}\mathbb{E}[\mathcal{L}(y)]=\mathbb{E}\big{[}\frac{ z}{2\sigma}(\mathcal{L}(\phi_{2}(\varphi(\phi_{1}(x);\theta+z))-\mathcal{L}( \phi_{2}(\varphi(\phi_{1}(x);\theta-z)))\big{]}.$ $$

Both equations (1) and (2) use Monte-Carlo simulation with $A>0$ queries to compute the loss gradient, and the corresponding estimates take the common form $\nabla_{\theta}\hat{\mathcal{L}}(y)=\frac{1}{A}\sum_{i=1}^{A}G(x,\theta,z_{i})$, where $z_{i}$ is sampled independently from $f(\cdot)$.

Given the limited computational resources, the parameter is updated iteratively with randomly sampled mini-batches of data denoted as $\{x_{1},\cdots,x_{B}\}$, where $B$ is the batch size.
For readability, we abbreviate $\nabla_{\theta}\hat{\mathcal{L}}(y_{j})=\frac{1}{A_{j}}\sum_{i=1}^{A_{j}}G(x_{
j},\theta,z_{i,j})$ as $\nabla_{\theta}\hat{\mathcal{L}}_{j}(\theta)$ and denote the batch-wise gradient estimates as $\nabla_{\theta}\hat{\mathcal{L}}(\theta_{t})$.
The updated recursion can be written as

$$ $\displaystyle\theta_{t+1}=\theta_{t}-\eta_{t}\nabla_{\theta}\hat{\mathcal{L}}( \theta_{t})=\theta_{t}-\frac{\eta_{t}}{B}\sum_{j=1}^{B}\nabla_{\theta}\hat{ \mathcal{L}}_{j}(\theta_{t}),$ (3) $$

where $\theta_{t}$ is the network parameter at the $t$-th step, and $\eta_{t}$ is the learning rate.

By default, the query amount $A_{j}$ allocated to estimate the gradient associated with each data point is equally allocated. However, the difficulty of gradient estimation and the model performance differ over data points. To conserve computational resources while maintaining the quality of gradient estimation, the number of queries should be properly allocated across data within the mini-batch.
Utilizing the likelihood ratio technique, we propose a universal accelerator for all perturbation-based training methods from a data-aware perspective.

## 4 Method

### 4.1 Optimal allocation via likelihood ratio method

Definition.
The query allocator is defined as a random vector $\bm{A}=(A_{1},\cdots,A_{B})$, where each component $A_{j}$ represents the (relative) number of queries allocated to the $j$-th data point in the batch. The allocator is assumed to follow a probability measure $\bm{A}\sim P(\bm{A}|\bm{\lambda};\Phi)$, which is parameterized by $\bm{\lambda}\in\Lambda$ and conditioned on features $\Phi$, such as the loss. At each step of neural network training, a sample of $\bm{A}$ is drawn from this distribution, and queries are allocated to data points based on the sample.
In this work, we select a Gaussian Allocator (GA)

$$ $\displaystyle\bm{A}\sim N(\bm{\mu},\Sigma|\bm{\lambda};\Phi),$ (4) $$

where $\bm{\lambda}=(\bm{\mu},\Sigma)\in\mathbb{R}^{B}\times\mathbb{R}^{B\times B}$. The matrix $\Sigma$ captures the correlation structure between data points and facilitates the allocation of queries. Intuitively, structurally similar data points possess similar levels of importance, leading to a similar trend in query allocation.
It is important to note that GA has a positive probability of generating negative samples. In practice, any data point receiving a negative allocation is assigned zero queries.

Bernoulli Allocator (BA) is another lightweight alternative, similar to the idea in Qin et al. (). In an equal allocation scenario, each data in the batch receives $\overline{A}$ queries. The BA introduces a probabilistic mechanism where, with probability $p$, the number of queries assigned to a data point is reduced to $\overline{A}/2$; otherwise, it remains at $\overline{A}$. The probability $p$ is the only parameter of the allocator.

It is important to note that $\bm{A}$ represents the relative number of queries, and thus the actual allocation proportion to the
$j$-th data point is given by $A_{j}/{\sum_{i=1}^{B}A_{i}}$. This definition eliminates the need for imposing a fixed total budget constraint in the following allocator optimization problem. However, as a consequence, we must specify that the total query budget per step is fixed at $A_{0}=\overline{A}\cdot B$ during the training process.
Note that, since the data points selected in the mini-batch differ at each step, $A_{j}^{t}$ is actually a function of $t$. For simplicity of notation, we omit this dependency in the remainder of the paper.

Optimization of the Query Allocator.
We employ a gradient-based method to optimize the allocator parameters $\bm{\lambda}$, introducing an additional optimization step before each iteration of neural network training. However, by formulating an equivalent optimization objective and utilizing reparameterization techniques, we ensure that this auxiliary optimization remains computationally lightweight and incurs minimal overhead.

Our objective is to maximize the performance improvement at this step, i.e.,

$$ $\displaystyle\arg\max_{\bm{\lambda}\in\Lambda}PI_{t}(\bm{\lambda})\triangleq \mathbb{E}_{\bm{A}\sim P(\cdot|\bm{\lambda})}[\sum_{j=1}^{B}\big{|}\mathcal{L} _{j}(\theta_{t+1})-\mathcal{L}_{j}(\theta_{t})\big{|}\big{]},$ (5) $$

where $\mathcal{L}_{j}$ is the loss corresponding to the $j$-th data point. Assuming that $\mathcal{L}_{j}$ is $L$-smooth, $\forall j=1,\cdots,B$, and $\Theta$ is a convex set, we can derive a lower bound of the optimization objective $PI_{t}(\bm{\lambda})$.

$$ $\displaystyle PI_{t}(\bm{\lambda})\geq LB_{t}(\bm{\lambda})\triangleq\mathbb{E }_{\bm{A}\sim P(\cdot|\bm{\lambda})}\sum_{j=1}^{B}\bigg{[}\big{(}\eta_{t} \nabla_{\theta}\hat{\mathcal{L}}(\theta_{t}))^{\top}\nabla_{\theta}\mathcal{L} _{j}(\theta_{t})-\frac{1}{2}\eta_{t}^{2}L\|\nabla_{\theta}\hat{\mathcal{L}}( \theta_{t})\|^{2}\bigg{]}.$ (6) $$

Inequality (6) (the proof is provided in Appendix A.1) closely resembles the Descent Lemma commonly used in the analysis of gradient descent. The lower bound $LB_{t}$ is often utilized as a surrogate optimization objective in the literature on adaptive step size (Pirotta et al., ) and adaptive batch size (Papini et al., ).
Since our decision variable is $\bm{\lambda}$ now, we can further simplify the form of $LB_{t}(\bm{\lambda})$. Before proceeding, according to the Central Limit Theorem, it is canonical to assume that the perturbation-based gradient estimator $\nabla_{\theta}\hat{\mathcal{L}}_{j}(\theta_{t})=\frac{1}{A_{j}}\sum_{i=1}^{A_
{j}}G(x_{j},\theta_{t},z_{i,j})$, $z_{i,j}\sim f(\cdot)$, $\forall j=1,\cdots,B,$ approximately follows a $d$-dimensional Gaussian distribution.

###### Assumption 1 .

$\forall j=1,\cdots,B$,
$\nabla_{\theta}\hat{\mathcal{L}}_{j}(\theta_{t})\sim N(\mu^{j,t}_{G},\frac{
\Sigma^{j,t}_{G}}{A_{j}})$,
where $\mu^{j,t}_{G}$ and $\Sigma^{j,t}_{G}$ are the expectation and the covariance matrix of $G(x_{j},\theta_{t},z)$, $z\sim f(\cdot)$, respectively.

###### Theorem 1 ( Equivalent Objective ) .

Suppose that the Assumption 1 holds, maximizing the lower bound $LB_{t}(\bm{\lambda})$ over $\bm{\lambda}\in\Lambda$ is equivalent to minimizing

$$ $J_{t}(\bm{\lambda})\triangleq\mathbb{E}_{\bm{A}\sim P(\cdot|\bm{\lambda})} \bigg{[}\sum_{j=1}^{B}\mathrm{Tr}(\frac{\Sigma_{G}^{j,t}}{A_{j}})\bigg{]}$ (7) $$

over $\bm{\lambda}\in\Lambda$, i.e.,

$$ $\max_{\bm{\lambda}\in\Lambda}LB_{t}(\bm{\lambda})\Longleftrightarrow\min_{\bm{ \lambda}\in\Lambda}J_{t}(\bm{\lambda}).$ $$

The equivalent expression (7) of the optimization objective makes the optimization of the allocator both computationally and statistically tractable (see the proof in Appendix A.2). First, it demonstrates that the optimization of the allocator is independent of both the smoothness parameter $L$ and the expectation $\mu^{j,t}_{G}$ of the estimator sample. Consequently, we can avoid estimating these unknown and bias-inducing high-dimensional parameters.
More importantly, the optimization of the allocator does not require estimating the off-diagonal elements of the covariance matrix $\Sigma_{G}^{j,t}$, which can be statistically and computationally challenging (Fan et al., ), especially since $\Sigma_{G}^{j,t}$ is of dimension $d^{2}$. It only requires estimating the trace of $\Sigma_{G}^{j,t}$, i.e., the sum of the diagonal variances, which significantly reduces the computational complexity.

Moreover, we adopt a subsampling technique in which the trace of the covariance matrix corresponding to the parameters of the deep layers of the neural network serves as an approximation for the overall trace. This approximation is justified as these layers typically account for the majority of the variance in the model. In practice, prior to the optimization process of the allocator, we sample a small number of initial queries for each data point to estimate $\mathrm{Tr}(\Sigma^{j,t}_{G})$. In the subsequent computations,
$\mathrm{Tr}(\Sigma^{j,t}_{G})$ is replaced by its estimate in a plug-in manner.
In addition, from Theorem 1, we can see that generally data points with higher variance should be allocated more query resources.

Then, the gradient of $J_{t}(\bm{\lambda})$ with respect to allocator parameter $\bm{\lambda}$ is given by

$$ $\displaystyle\nabla_{\bm{\lambda}}J_{t}(\bm{\lambda})=\nabla_{\bm{\lambda}} \mathbb{E}\bigg{[}\sum_{j=1}^{B}\mathrm{Tr}(\frac{\Sigma^{j,t}_{G}}{A_{j}}) \bigg{]}=\mathbb{E}\bigg{[}\sum_{j=1}^{B}\mathrm{Tr}(\frac{\Sigma^{j,t}_{G}}{A _{j}})\nabla_{\bm{\lambda}}\ln(P(\bm{A}))\bigg{]}.$ (8) $$

The LR gradient estimator

$$ $\displaystyle\nabla_{\bm{\lambda}}\hat{J}(\bm{\lambda})=\frac{1}{K}\sum_{k=1}^ {K}\bigg{[}\sum_{j=1}^{B}\mathrm{Tr}(\frac{\Sigma^{j,t}_{G}}{A^{(k)}_{j}}) \nabla_{\bm{\lambda}}\ln(P(\bm{A}^{(k)}))\bigg{]},$ (9) $$

where $\bm{A}^{(k)}=(A_{1}^{(k)},\cdots,A_{B}^{(k)})$ are i.i.d samples from $P(\cdot|\bm{\lambda})$. At each step $t$ of neural network training, we employ the above gradient estimator to search for the optimal allocator parameter $\bm{\lambda}_{t}^{\ast}\triangleq\arg\max_{\bm{\lambda}\in\Lambda}J(\bm{
\lambda})$.
Let $LB^{\ast}_{1:T}$ and $LB_{1:T}^{\mathrm{equal}}$ represent the lower bound of the cumulative performance improvement of the neural network after $T$ steps of training when using the optimal allocator $\bm{A}^{\ast}\sim P(\cdot|\lambda^{\ast}_{t})$ and using an equal allocator $\bm{A}^{\mathrm{equal}}$, respectively. Then the difference between these two allocation strategies is lower bound by the following expression.

###### Theorem 2 ( Theoretical Improvement ) .

Suppose that the Assumption 1 holds, then we have

$$ $\mathbb{E}\big{(}LB^{\ast}_{1:T}-LB_{1:T}^{\mathrm{equal}}\big{)}\geq\sum_{t=1 }^{T}\frac{\eta_{t}^{2}L}{2BA_{0}}\sum_{j<k}\left(\sqrt{\mathrm{Tr}(\Sigma^{j, t}_{G})}-\sqrt{\mathrm{Tr}(\Sigma^{k,t}_{G})}\right)^{2}\geq 0.$ $$

Theorem 2 (see the proof in Appendix A.3) quantifies the superiority of our proposed optimal allocator over the vanilla equal allocation strategy. The greater the differences in $\mathrm{Tr}(\Sigma^{j,t}_{G})$ across different data points (i.e., the stronger the heterogeneity), the more pronounced the superiority of the optimal allocator, and this difference equals zero if and only if all $\mathrm{Tr}(\Sigma^{j,t}_{G})$ are equal.

Reparameterization of the Gaussian allocator.
The original parameter dimension of the Gaussian allocator is $B+\frac{B(B+1)}{2}$, which might be problematic due to its high dimensionality. To reduce the dimensionality of the allocator’s optimization problem and to incorporate more relevant features, we adopt a reparameterization approach, i.e., imposing a specific structure to the parameter $\bm{\lambda}$.

Drawing inspiration from the literature on data pruning (Qin et al., ), the loss value serves as a strong indicator of the data point’s characteristics. Therefore, we assume that the mean of the allocator follows a linear structure:

$$ $\bm{\mu}=\beta_{0}+\beta_{1}\Phi,$ $$

where $\Phi=\mathrm{Tanh}(\mathcal{L}(x;\bm{\theta};\cdot))$ is the “clean” loss without injected noise. The original loss value is processed through a $\mathrm{Tanh}$ function to constraint $\Phi$ within the range $[-1,1]$. The initial value for $\beta_{0}$ and $\beta_{1}$ are set to $\overline{A}$ and $\overline{A}/2$, respectively.

For the covariance structure $\Sigma$, based on the observation that similar data may have similar influences on the training of neural networks, we reparameterize the covariance matrix as follows: the covariance between the allocation of the $i$-th data point and the $j$-th data point is

$$ $\Sigma(i,j)=\sigma^{2}\mathrm{exp}\bigg{(}-\frac{d_{i,j}}{2\gamma^{2}}\bigg{)},$ $$

where $d_{i,j}$ is the cosine similarity of the embedding between the $i$-th and the $j$-th data points. $\sigma^{2}$ and $\gamma^{2}$ are two parameters that control the scale and decay of the covariance, respectively. The $\sigma$ is initialized as $\overline{A}/5$ and $\gamma$ is initialized as $1$. This kernel-based reparameterization form, also known as the Radial Basis Function (RBF), is widely used in Gaussian process regression and Bayesian Optimization literature (Shahriari et al., ). It enables us to model dependencies between data based on their pairwise distances.

By applying the reparameterization techniques to the covariance matrix $\Sigma$ and mean vector
$\bm{\mu}$, the original high-dimensional optimization problem for these parameters has been effectively transformed into a more manageable optimization problem for 4 variables $\bm{\lambda}=(\beta_{0},\beta_{1},\sigma,\gamma)$. The pseudocode for the forward learning under optimal allocation is in Appendix B.

###### Assumption 1 .

###### Theorem 1 ( Equivalent Objective ) .

###### Theorem 2 ( Theoretical Improvement ) .

## 5 Experiments

### 5.1 Finetuning Vision Transformer

Experimental setting: We evaluate our model’s performance using a diverse set of widely used benchmark datasets, each chosen for its unique characteristics and specific challenges contributing to a comprehensive analysis of our method’s generalization across different domains: ImageNet (Deng et al., ), Caltech101 (Fei-Fei et al., ), Food101 (Bossard et al., ), Flowers102 (Nilsback & Zisserman, ), CIFAR10/100 (Krizhevsky et al., ), and EuroSAT (Helber et al., ).

Evaluation metrics: We focus on image classification and evaluate the performance by accuracy in different training contexts: BP, LR, and SPSA. We compare the performance with and without the query allocation. The GA is employed for all experiments.

Table: Table 1: Classification accuracy ($\%$) for finetuning vit on downstream dataset.

Using different gradient estimators, we conduct full-parameters fine-tuning for the Vision Transformer (both base and large model). We include DeepZero (SPSA with $n$ queries, denoted as $n$-SPSA), Mezo (SPSA with only $2$ queries), and LR with equally allocated query budgets as baselines in our experiments. For the SPSA-based methods, we include a version with the optimal allocator (OPS-SPSA) to compare with their corresponding baselines. We apply the same settings to the LR, comparing the optimally allocated version (OPS-LR) with the vanilla version.
We train the network for 100 epochs with the learning rate of $1\times 10^{-4}$ and the Adam optimizer. The batch size is 64. All the methods in our experiments use the same query budgets, except for Mezo, which uses only 2 queries per data point in accordance with its original memory-efficient settings. As shown in Table 1, OPS-LR outperforms all other methods. In contrast, Mezo with only 1 query performs poorly in the experiment. On the contrary, DeepZero achieves acceptable accuracy compared with Mezo. The phenomenon suggests that an adequate number of queries is crucial for the success of perturbation-based methods since the gradient’s variance must be properly controlled. The OPS-SPSA achieves better accuracy than the DeepZero with an equally allocated budget. Our proposed optimal allocator provides universal acceleration across all perturbation-based methods. Moreover, LR achieves higher accuracy than DeepZero, suggesting that LR is a better gradient estimator than the SPSA in this experiment context. Please refer to the Appendix C.1 for the details and further analysis.

### 5.2 Black-box tuning for Model as a Service

Figure: Figure 2: illustration paradigm of fine-tuning the prompt for black-box vision language model
Refer to caption: x2.png

We can tune the prompts for vision-language models (VLMs) in a Model-as-a-Service (MaaS) environment using the OPS-LR framework, as shown in Figure 2. We address the challenge of fine-tuning VLMs without access to model internals or gradient information through a scalable black-box optimization approach. We use CLIP (Radford et al., ), with embedding dimensions of $D_{T}=512$ for the text encoder and $D_{I}=756$ for the image encoder, as the backbone foundation model. Both the text encoder and the image encoder are kept frozen as black-box modules.
Tunable prefix or suffix prompt tokens are injected into the text and image token sequences, respectively: $S_{T}=[cls,{\color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}p
_{T}},e_{T}]\in\mathbb{R}^{(1+n_{T}+m_{T})\times D_{T}}$ is the input sequence for the text encoder, $S_{I}=[cls,e_{I},{\color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{
0,0,1}p_{I}}]\in\mathbb{R}^{(1+m_{I}+n_{I})\times D_{I}}$ for the image counterpart. Previous researches have suggested that large foundation models have low intrinsic dimensions. Therefore, it is plausible to project the original embedding space onto a low-dimension subspace: $p_{T}=z_{T}\cdot W_{T};W_{T}\in\mathbb{R}^{d_{T}\times D_{T}},z_{T}\in\mathbb{
R}^{m_{T}\times d_{T}}$ and $p_{I}=z_{I}\cdot W_{I};W_{I}\in\mathbb{R}^{d_{I}\times D_{I}},z_{I}\in\mathbb{
R}^{m_{I}\times d_{I}}$. $W_{T}$ and $W_{I}$ are the projection matrix. The projection matrix is kept frozen after being initialized with a Gaussian distribution. With the projection, our black-box tuning problem is reduced to dimensions of $d_{T}+d_{I}\ll D_{T}+D_{I}$. The optimization problem, using the cross-entropy loss, takes the form:

$$ $\displaystyle Z^{*}=\mathop{\arg\max}\limits_{z_{T},z_{I}}\mathcal{L}([e_{T},e _{I}];\bm{\theta}=[z_{T}\cdot D_{T},z_{I}\cdot D_{I}];\cdot),$ (10) $$

Table: Table 2: Classification accuracy ($\%$) for black-box prompt tuning.

where $Z$ is the unified formulation of $[z_{T},z_{I}]$. By applying the OPS-LR framework, we can optimize the black-box prompt tuning problems without access to the model structure. We can achieve an accuracy of $69.7\%$ on ImageNet by tuning only a few thousand parameters, surpassing the performance of other black-box methods and Manual Prompting (Radford et al., ). Please refer to the Appendix C.2 for more details.

### 5.3 Alignment between black-box foundation model

Figure: (a)
Refer to caption: x3.png

Foundation models like GPT (Brown, ), LLaMa Touvron et al. (), Vicuna (Chiang et al., ), and CLIP (Radford et al., ) contain rich domain knowledge from the pretraining. Multimodal tasks, such as video captioning (Krishna et al., ) and video grounding(Gao et al., ), require models to have sufficient cross-domain understanding, and proper alignment between different modalities is essential.
Aligning foundation models from a single modality is a more economical approach than pretraining from scratch using multimodal data. It is a common practice to add an adapter between the visual encoder and the Large Language model. The adapter is a linear projector, $g(\cdot)$, that projects the visual embedding to the textual embedding space. Since foundation models contain billions of parameters, the backward pass through the model requires extra memory for storing the computation graph, and recursive computations through such a large pre-trained model are time-consuming. Our OPS-LR framework efficiently bypasses the need for recursive computation to estimate the gradient with a large foundation model in between, as shown in Figure 3(a).

The training paradigm for the multimodal alignment follows an autoregressive style using image-text pairs $<I,T>$. We use a frozen CLIP ViT-L/14 as the visual encoder, denoted as $ViT$. The CLS tokens of the output sequence of $ViT$ are utilized as the global feature for the frame. We denote the visual token after the projection as $Z_{v}=g(ViT(I)^{CLS})$. The visual token is then concatenated with the text tokens sequence: $[Z_{v},t_{1},t_{2},...,t_{m}]$, where $[t_{1},t_{2},...,t_{m}]$ is the text tokens sequence with a length of $m$. Consequently, we can train the adapter on the concatenated sequence with the autoregressive objective of the LLM. We report the wall time for the feature alignment procedure on the LCS-558K dataset (Liu et al., ). Our OPS-LR framework can reduce almost half of the training time since the recursive backward computation is skipped. From Figure 3(b), the wall time with or without the allocator is almost the same, which further illustrates that our allocator is lightweight and has minimal impact on runtime. Please refer to the Appendix C.3 for more details.

### 5.4 Abation study

Figure: (a)
Refer to caption: x5.png

We conduct an ablation study using ViT-base on the CIFAR100 dataset. We investigate the necessity of the allocation policy for the forward-learning training and determine whether GA or BA is the better allocator. Furthermore, we discern the varying degrees of difficulty in estimation across different network components. The results are shown in Figure 4.

Impact of different query budgets. To further investigate the necessity of the procedure of query budget allocation, we compare the performance increment under various query budgets for different allocation policies. Traditionally, the queries are equally allocated (EA) across different data. The performance improvements offered by GA and BA are considerably more significant than those of EA. When the number of queries per data is small, the benefit from the allocation is still non-ignorable. In other words, the more query budget you have, the greater the improvement gains from allocation. Meanwhile, the ablation on different transformer layers indicates that the variance of the gradient increases as the layer goes deeper, which validates our subsampling technique.

Impact of different allocators, Gaussian or Bernoulli. We compare different allocators in Figure 4(a), 4(b),4(d), and 4(e). GA controls the allocation with the mean vectors and the covariance matrix. The covariance matrix captures the similarity between different data points.
Structurally similar data points possess
similar levels of importance. Therefore, the allocation policy for those data points should be correlated. If the covariance matrix is diagonal, the similarity between data points is ignored, which could undermine the performance of the allocation strategy. Compared with the GA, the BA is a simpler strategy. Intuitively, it prunes the queries on unimportant data points to eliminate unnecessary computation. Our ablation study suggests that GA is superior to BA.

Estimation on different components of Transformer. We present the cosine similarity results on the multi-head attention (MHA) and the feed-forward network (FFN) in Figure 4(a), 4(b), 4(c), and 4(f). Estimation on the FFN is easier than on the MHA. In the Transformer block, the MHA precedes the FFN, which may explain this phenomenon. The low cosine similarity on MHA attention of deep layers may lead to the performance gap on large-scale dataset like ImageNet.
It is also worth noting that a large number of queries when training Transformer can improve the estimation quality. However, the memory cost of increasing queries is extremely high since the complexity of the attention module is $\mathcal{O}(n^{2})$. Integrating efficient MHA techniques, such as Flash Attention (Dao et al., ), is left for future work.

## 6 Conclusions

In this paper, we unify different forward learning under the perturbation-based framework. We propose an optimal allocation method to accelerate forward learning by controlling the gradient variance in the mini-batch. Theoretical results show that the allocation can be conducted in a lightweight style and the optimality of the allocation is proved. Compared with other forward learning methods, we achieve the SOTA performance in fine-tuning ViT. With the improved scalability, we apply our method to two black-box scenarios for the foundation model. However, there is still a non-neglectable gap compared with BP. Other techniques of accelerating forward learning are left to future work. Moreover, forward learning can optimize black-box objectives, such as human preference, where taking derivative is infeasible. It is also an interesting future direction.

## Appendix A Theoretical Details

### A.1 Proof of inequality ( 6 )

$$ $\displaystyle PI_{t}(\bm{\lambda})\geq\mathbb{E}_{\bm{A}\sim P(\cdot|\bm{ \lambda})}\sum_{j=1}^{B}\bigg{[}\big{(}\eta_{t}\nabla_{\theta}\hat{\mathcal{L} }(\theta_{t}))^{\top}\nabla_{\theta}\mathcal{L}_{j}(\theta_{t})-\frac{1}{2} \eta_{t}^{2}L\|\nabla_{\theta}\hat{\mathcal{L}}(\theta_{t})\|^{2}\bigg{]}.$ $$

###### Proof.

For each loss function $\mathcal{L}_{j}$, by exploiting Taylor’s theorem, we have

$$ $\displaystyle\mathcal{L}_{j}(\theta_{t+1})-\mathcal{L}_{j}(\theta_{t})$ $\displaystyle=\int_{0}^{1}\frac{d}{d\tau}\mathcal{L}_{j}(\theta_{t}+\tau( \theta_{t+1}-\theta_{t}))\,d\tau$ (11) $\displaystyle=\int_{0}^{1}\langle\nabla_{\theta}\mathcal{L}_{j}(\theta_{t}+ \tau(\theta_{t+1}-\theta_{t})),\theta_{t+1}-\theta_{t}\rangle\,d\tau.$ $$

By adding and subtracting $\nabla_{\theta}\mathcal{L}_{j}(\theta_{t})$,

$$ $\displaystyle(\ref{appendix_1})$ $\displaystyle=\int_{0}^{1}\langle\nabla_{\theta}\mathcal{L}_{j}(\theta_{t}), \theta_{t+1}-\theta_{t}\rangle\,d\tau+\int_{0}^{1}\langle\nabla_{\theta} \mathcal{L}_{j}(\theta_{t}+\tau(\theta_{t+1}-\theta_{t}))-\nabla_{\theta} \mathcal{L}_{j}(\theta_{t}),\theta_{t+1}-\theta_{t}\rangle\,d\tau$ (12) $\displaystyle=\langle\nabla_{\theta}\mathcal{L}_{j}(\theta_{t}),\theta_{t+1}- \theta_{t}\rangle+\int_{0}^{1}\langle\nabla_{\theta}\mathcal{L}_{j}(\theta_{t} +\tau(\theta_{t+1}-\theta_{t}))-\nabla_{\theta}\mathcal{L}_{j}(\theta_{t}), \theta_{t+1}-\theta_{t}\rangle\,d\tau.$ $$

Now we focus on the second term in (12). By the Cauchy-Schwarz inequality, we have

$$ $\left|\langle\nabla_{\theta}\mathcal{L}_{j}(\theta_{t}+\tau(\theta_{t+1}- \theta_{t}))-\nabla_{\theta}\mathcal{L}_{j}(\theta_{t}),\theta_{t+1}-\theta_{t }\rangle\right|\leq\|\nabla_{\theta}\mathcal{L}_{j}(\theta_{t}+\tau(\theta_{t+ 1}-\theta_{t}))-\nabla_{\theta}\mathcal{L}_{j}(\theta_{t})\|\cdot\|\theta_{t+1 }-\theta_{t}\|.$ $$

And according to the $L$-smoothness assumption, we have

$$ $\|\nabla_{\theta}\mathcal{L}_{j}(\theta_{t}+\tau(\theta_{t+1}-\theta_{t}))- \nabla_{\theta}\mathcal{L}_{j}(\theta_{t})\|\leq L\|\tau(\theta_{t+1}-\theta_{ t})\|=L\tau\|\theta_{t+1}-\theta_{t}\|.$ $$

Therefore, the second term in (12) becomes

$$ $\displaystyle\int_{0}^{1}\langle\nabla_{\theta}\mathcal{L}_{j}(\theta_{t}$ $\displaystyle+\tau(\theta_{t+1}-\theta_{t}))-\nabla_{\theta}\mathcal{L}_{j}( \theta_{t}),\theta_{t+1}-\theta_{t}\rangle\,d\tau$ (13) $\displaystyle\geq-\int_{0}^{1}L\tau\|\theta_{t+1}-\theta_{t}\|^{2}\,d\tau=- \frac{L}{2}\|\theta_{t+1}-\theta_{t}\|^{2}.$ $$

Combining (12) and (13), and substitute $\theta_{t+1}=\theta_{t}+\eta_{t}\nabla_{\theta}\hat{\mathcal{L}}(\theta_{t})$, we have

$$ $\mathcal{L}_{j}(\theta_{t}+\eta_{t}\nabla_{\theta}\hat{\mathcal{L}}(\theta_{t} ))-\mathcal{L}_{j}(\theta_{t})\geq\eta_{t}\langle\nabla_{\theta}\mathcal{L}_{j }(\theta_{t}),\nabla_{\theta}\hat{\mathcal{L}}(\theta_{t})\rangle-\frac{L}{2} \eta_{t}^{2}\|\nabla_{\theta}\hat{\mathcal{L}}(\theta_{t})\|^{2}.$ $$

Then, taking expectation over $\bm{A}\sim P(\cdot|\bm{\lambda})$ and summing over $j$,

$$ $\displaystyle PI_{t}(\bm{\lambda})\geq\mathbb{E}_{\bm{A}\sim P(\cdot|\bm{ \lambda})}\sum_{j=1}^{B}\bigg{[}\big{(}\eta_{t}\nabla_{\theta}\hat{\mathcal{L} }(\theta_{t}))^{\top}\nabla_{\theta}\mathcal{L}_{j}(\theta_{t})-\frac{1}{2} \eta_{t}^{2}L\|\nabla_{\theta}\hat{\mathcal{L}}(\theta_{t})\|^{2}\bigg{]},$ $$

which completes the proof.
∎

###### Proof.

### A.2 Proof of Theorem 1

###### Proof.

$$ $\displaystyle LB_{t}$ $\displaystyle=\mathbb{E}_{\bm{A}}\sum_{j=1}^{B}\bigg{[}\big{(}\eta_{t}\nabla_{ \theta}\hat{\mathcal{L}}(\theta_{t}))^{\top}\nabla_{\theta}\mathcal{L}_{j}( \theta_{t})-\frac{1}{2}\eta_{t}^{2}L\|\nabla_{\theta}\hat{\mathcal{L}}(\theta_ {t})\|^{2}\bigg{]}$ (14) $\displaystyle=\mathbb{E}_{\bm{A}}\sum_{j=1}^{B}\bigg{[}\big{(}\frac{\eta_{t}}{ B}\sum_{j=1}^{B}\nabla_{\theta}\hat{\mathcal{L}}_{j}(\theta_{t})\big{)}^{\top} \nabla_{\theta}\mathcal{L}_{j}(\theta_{t})-\frac{1}{2}\eta_{t}^{2}L\|\frac{1}{ B}\sum_{j=1}^{B}\nabla_{\theta}\hat{\mathcal{L}}_{j}(\theta_{t})\|^{2}\bigg{]},$ $$

where $\nabla_{\theta}\hat{\mathcal{L}}_{j}(\theta_{t})\sim N(\mu^{j,t}_{G},\frac{
\Sigma^{j,t}_{G}}{A_{j}})$, and $j=1,\cdots,B$.
Then, by taking the conditional expectation.

$$ $\displaystyle LB_{t}$ $\displaystyle=\mathbb{E}_{\bm{A}}\sum_{j=1}^{B}\bigg{[}\mathbb{E}_{\nabla_{ \theta}\hat{\mathcal{L}}(\theta_{t})}\bigg{[}\big{(}\frac{\eta_{t}}{B}\sum_{j= 1}^{B}\nabla_{\theta}\hat{\mathcal{L}}_{j}(\theta_{t})\big{)}^{\top}\nabla_{ \theta}\mathcal{L}_{j}(\theta_{t})-\frac{1}{2}\eta_{t}^{2}L\bigg{\|}\frac{1}{B }\sum_{j=1}^{B}\nabla_{\theta}\hat{\mathcal{L}}_{j}(\theta_{t})\bigg{\|}^{2} \bigg{|}\bm{A}\bigg{]}\bigg{]}$ (15) $$

$\forall j=1,\cdots,B$, the expected inner product on the RHS of (15)

$$ $\displaystyle\mathbb{E}_{\nabla_{\theta}\hat{\mathcal{L}}(\theta_{t})}\bigg{[} \big{(}\frac{\eta_{t}}{B}\sum_{j=1}^{B}\nabla_{\theta}\hat{\mathcal{L}}_{j}( \theta_{t})\big{)}^{\top}\nabla_{\theta}\mathcal{L}_{j}(\theta_{t})\bigg{|}\bm {A}\bigg{]}=\frac{\eta_{t}}{B}\big{(}\sum_{j=1}^{B}\mu^{j,t}_{G}\big{)}^{\top} \nabla_{\theta}\mathcal{L}_{j}(\theta_{t}).$ (16) $$

Next, we focus on the expected squared norm on the RHS of (15).

$$ $\displaystyle\mathbb{E}_{\nabla_{\theta}\hat{\mathcal{L}}(\theta_{t})}\bigg{[} \bigg{\|}\frac{1}{B}\sum_{j=1}^{B}\nabla_{\theta}\hat{\mathcal{L}}_{j}(\theta_ {t})\bigg{\|}^{2}\bigg{|}\bm{A}\bigg{]}=\frac{1}{B^{2}}\bigg{(}\sum_{j=1}^{B} \sum_{k=1}^{B}\mathbb{E}\bigg{[}\nabla_{\theta}\hat{\mathcal{L}}_{j}(\theta_{t })^{\top}\nabla_{\theta}\hat{\mathcal{L}}_{k}(\theta_{t})\bigg{|}\bm{A}\bigg{] }\bigg{)}.$ (17) $$

For $j\neq k$, we have

$$ $\mathbb{E}\left[\nabla_{\theta}\hat{\mathcal{L}}_{j}(\theta_{t})^{\top}\nabla_ {\theta}\hat{\mathcal{L}}_{k}(\theta_{t})\bigg{|}\bm{A}\right]=\left(\mu_{G}^{ j,t}\right)^{\top}\mu_{G}^{k,t}.$ (18) $$

For $j=k$, we express $\nabla_{\theta}\hat{\mathcal{L}}_{j}(\theta_{t})$ as the sum of the mean and a zero-mean normal random vector:

$$ $\nabla_{\theta}\hat{\mathcal{L}}_{j}(\theta_{t})=\mu_{G}^{j,t}+Z,$ $$

where $Z\sim N(0,\frac{\Sigma_{G}^{j,t}}{A_{j}})$, meaning $Z$ is a zero-mean random vector with covariance $\frac{\Sigma_{G}^{j,t}}{A_{j}}$. Then we have

$$ $\displaystyle\mathbb{E}\left[\nabla_{\theta}\hat{\mathcal{L}}_{j}(\theta_{t})^ {\top}\nabla_{\theta}\hat{\mathcal{L}}_{k}(\theta_{t})\bigg{|}\bm{A}\right]$ $\displaystyle=\mathbb{E}\left[\left(\mu_{G}^{j,t}+Z\right)^{\top}\left(\mu_{G} ^{j,t}+Z\right)\bigg{|}\bm{A}\right]$ (19) $\displaystyle=\mathbb{E}\left[\left(\mu_{G}^{j,t}\right)^{\top}\mu_{G}^{j,t}+2 \left(\mu_{G}^{j,t}\right)^{\top}Z+Z^{\top}Z\bigg{|}\bm{A}\right]$ $\displaystyle=\left\|\mu_{G}^{j,t}\right\|^{2}+0+\mathbb{E}\left[Z^{\top}Z|A \right].$ $$

The expected value of the quadratic form $Z^{\top}Z$ in (19) for $Z\sim N(0,\frac{\Sigma_{G}^{j,t}}{A_{j}})$ is given by the trace of the covariance matrix, i.e.,

$$ $\mathbb{E}\left[Z^{\top}Z\right]=\operatorname{Tr}\left(\frac{\Sigma_{G}^{j,t} }{A_{j}}\right).$ (20) $$

Therefore, with (18), (19) and (20), we have

$$ $(\ref{appendix_5})=\frac{1}{B^{2}}\bigg{(}\bigg{\|}\sum_{j=1}^{B}\mu_{G}^{j,t} \bigg{\|}^{2}+\sum_{j=1}^{B}\operatorname{Tr}\bigg{(}\frac{\Sigma_{G}^{j,t}}{A _{j}}\bigg{)}\bigg{)}.$ (21) $$

Then, with (15), (16) and (21), we have

$$ $\displaystyle LB_{t}(\bm{\lambda})=\mathbb{E}_{\bm{A}\sim p(\cdot|\bm{\lambda} )}\bigg{[}\frac{\eta_{t}}{B}\big{(}\sum_{j=1}^{B}\mu^{j,t}_{G}\big{)}^{\top}$ $\displaystyle(\sum_{j=1}^{B}\nabla_{\theta}\mathcal{L}_{j}(\theta_{t}))$ (22) $\displaystyle-\frac{1}{2}\sum_{j=1}^{B}\frac{\eta_{t}^{2}L}{B^{2}}\bigg{(} \bigg{\|}\sum_{j=1}^{B}\mu_{G}^{j,t}\bigg{\|}^{2}+\sum_{j=1}^{B}\operatorname{ Tr}\bigg{(}\frac{\Sigma_{G}^{j,t}}{A_{j}}\bigg{)}\bigg{)}\bigg{]}.$ $$

Notice that, $\mu^{j,t}_{G}$ is independent of the allocation $\bm{A}$, which only depends on the estimator type and the parameters of the injected noise in the perturbation. Therefore, it is easy to see that maximizing the lower bound $LB_{t}(\bm{\lambda})$ over $\bm{\lambda}\in\Lambda$ is equivalent to minimizing

$$ $J_{t}(\bm{\lambda})\triangleq\mathbb{E}_{\bm{A}\sim P(\cdot|\bm{\lambda})} \bigg{[}\sum_{j=1}^{B}\mathrm{Tr}(\frac{\Sigma_{G}^{j,t}}{A_{j}})\bigg{]}$ (23) $$

over $\bm{\lambda}\in\Lambda$.
∎

###### Proof.

### A.3 Proof of Theorem 2

###### Proof.

Based on the optimality of $\lambda_{t}^{\ast}$ at step $t$, we know that for any GA following $N(\mu,\Sigma)$, the corresponding $\mathbb{E}(LB_{t}^{GA})\leq\mathbb{E}(LB_{t}^{\ast})$. Now, if we set $\Sigma=\bm{0}_{B\times B}$, the Gaussian distribution becomes deterministic, and the query allocation becomes a fixed, non-random strategy. To establish a lower bound for $\mathbb{E}(LB_{t}^{\ast})-\mathbb{E}(LB_{t}^{\mathrm{equal}})$, we construct an optimal deterministic allocator $\bm{A}^{\mathrm{det}}$ as the intermediary strategy between $\bm{A}^{\ast}$ and $\bm{A}^{\mathrm{equal}}$, which satisfies

$$ $\mathbb{E}(LB_{t}^{\mathrm{equal}})\leq\mathbb{E}(LB_{t}^{\mathrm{det}})\leq \mathbb{E}(LB_{t}^{\ast}),$ $$

where $LB_{t}^{\mathrm{det}}$ is lower bound of the PI corresponding to $\bm{A}^{\mathrm{det}}$.

For the optimal deterministic allocator $\bm{A}^{\mathrm{det}}$, we aim to maximize the objective function

$$ $J_{t}(\bm{A})=\sum_{j=1}^{B}\frac{\mathrm{Tr}(\Sigma^{j,t}_{G})}{A_{j}}$ $$

subject to the constraint $\sum_{j=1}^{B}A_{j}=A_{0}$. The Lagrangian function is

$$ $L=\sum_{j=1}^{B}\frac{\mathrm{Tr}(\Sigma^{j,t}_{G})}{A_{j}}-\lambda\bigg{(} \sum_{j=1}^{B}A_{j}-A_{0}\bigg{)},$ $$

where $\lambda$ is the Lagrange multiplier. Take the partial derivative of $L$ with respect to each $A_{j}$ and set it to zero.

$$ $\frac{\partial L}{\partial A_{j}}=-\frac{\mathrm{Tr}(\Sigma^{j,t}_{G})}{A_{j}^ {2}}-\lambda=0.$ $$

Then we have

$$ $A^{\mathrm{det}}_{j}=A_{0}\cdot\frac{\sqrt{\mathrm{Tr}(\Sigma^{j,t}_{G})}}{ \sum_{k=1}^{B}\sqrt{\mathrm{Tr}(\Sigma^{k,t}_{G})}}$ $$

Then corresponding objective is

$$ $J_{t}(\bm{A}^{\mathrm{det}})=\sum_{j=1}^{B}\frac{\mathrm{Tr}(\Sigma^{j,t}_{G}) }{A_{j}}=\sum_{j=1}^{B}\frac{\mathrm{Tr}(\Sigma^{j,t}_{G})}{A_{0}\cdot\frac{ \sqrt{\mathrm{Tr}(\Sigma^{j,t}_{G})}}{\sum_{k=1}^{B}\sqrt{\mathrm{Tr}(\Sigma^{ k,t}_{G})}}}=\frac{\left(\sum_{j=1}^{B}\sqrt{\mathrm{Tr}(\Sigma^{j,t}_{G})} \right)^{2}}{A_{0}}.$ $$

As for the equal allocator $A^{\mathrm{equal}}_{j}=\frac{A_{0}}{B}$,
the corresponding objective function becomes

$$ $J_{t}(\bm{A}^{\mathrm{equal}})=\frac{A_{0}}{B}=\sum_{j=1}^{B}\frac{\mathrm{Tr} (\Sigma^{j,t}_{G})}{\frac{A_{0}}{B}}=\frac{B}{A_{0}}\sum_{j=1}^{B}\mathrm{Tr}( \Sigma^{j,t}_{G}).$ $$

It follows that

$$ $J_{t}(\bm{A}^{\mathrm{equal}})-J_{t}(\bm{A}^{\mathrm{det}})=\frac{\left(\sum_{ j=1}^{B}\sqrt{\mathrm{Tr}(\Sigma^{j,t}_{G})}\right)^{2}}{A_{0}}-\frac{B}{A_{0} }\sum_{j=1}^{B}\mathrm{Tr}(\Sigma^{j,t}_{G}).$ $$

Notice that

$$ $B\sum_{j=1}^{B}\mathrm{Tr}(\Sigma^{j,t}_{G})-\bigg{(}\sum_{j=1}^{B}\sqrt{ \mathrm{Tr}(\Sigma^{j,t}_{G})}\bigg{)}^{2}=\sum_{j<k}\left(\sqrt{\mathrm{Tr}( \Sigma^{j,t}_{G})}-\sqrt{\mathrm{Tr}(\Sigma^{k,t}_{G})}\right)^{2}.$ $$

Then

$$ $J_{t}(\bm{A}^{\mathrm{equal}})-J_{t}(\bm{A}^{\mathrm{det}})=\frac{1}{A_{0}} \sum_{j<k}\left(\sqrt{\mathrm{Tr}(\Sigma^{j,t}_{G})}-\sqrt{\mathrm{Tr}(\Sigma^ {k,t}_{G})}\right)^{2}.$ $$

This difference is always non-negative and equals zero if and only if all $\mathrm{Tr}(\Sigma^{j,t}_{G})$ are equal.
Then we have

$$ $\displaystyle\mathbb{E}\big{(}LB^{\ast}_{1:T}-LB_{1:T}^{\mathrm{equal}}\big{)}$ $\displaystyle\geq\mathbb{E}\big{(}LB^{\mathrm{det}}_{1:T}-LB_{1:T}^{\mathrm{ equal}}\big{)}=\sum_{t=1}^{\top}\frac{\eta_{t}^{2}L}{2B}(J_{t}(\bm{A}^{\mathrm {equal}})-J_{t}(\bm{A}^{\mathrm{det}}))$ (24) $\displaystyle\geq\sum_{t=1}^{\top}\frac{\eta_{t}^{2}L}{2BA_{0}}\sum_{j<k}\left (\sqrt{\mathrm{Tr}(\Sigma^{j,t}_{G})}-\sqrt{\mathrm{Tr}(\Sigma^{k,t}_{G})} \right)^{2},$ $$

which completes the proof.
∎

###### Proof.

## Appendix B Pseudocode

Figure: Algorithm 1 Perturbation-based training via optimal allocation

## Appendix C Experiment Details

Platform: All the experiments are conducted on a machine with 8 NVIDIA A800 GPUs. Each A800 GPU has 80GB of memory.

### C.1 Experiments for ViT

Datasets:
ImageNet: Over 1.2 million images in 1,000 classes, a standard large-scale challenge for object classification.
Caltech101: 9,144 images across 101 categories, focusing on object recognition with varying perspectives and poses.
Food101: 101,000 images in 101 food categories, presenting challenges in texture, lighting, and fine-grained classification.
Flower102: 8,189 images across 102 flower species, testing subtle feature recognition.
CIFAR-10/100: Consisting of 60,000 32x32 images, CIFAR-10 has 10 classes, while CIFAR-100 has 100 classes, offering challenges in small-scale image classification with increasing complexity in category distinctions.
EuroSAT: 27,000 satellite images in 10 land-use categories, challenging models to perform in remote sensing tasks.
All the input images are resized to 224 before entering the ViT backbone.

ViT-base: The base model has 12 transformer layers. The hidden dimension is 786, with 12 heads for the multi-head attention. The intermediate dimension for feed-forward mlp is 3072. 12 layers are equally divided into three groups from bottom to top. The standard deviation of the injected Gaussian noise for each group is initialized as $1\times 10^{-5},1\times 10^{-4},1\times 10^{-3}$. The query budget for each group is $20$, $60$, and $120$ queries per data. The query budget for the classification head is $60$ queries per data.

ViT-large: The large model has 24 transformer layers. The hidden dimension is 1024, with 16 heads for the multi-head attention. The intermediate dimension for feed-forward mlp is 4096. 24 layers are equally divided into four groups from bottom to top. The standard deviation of the injected Gaussian noise for each group is initialized as $5\times 10^{-6},1\times 10^{-5},1\times 10^{-4},1\times 10^{-3}$. The query budget for each group is $20$, $20$, $60$, and $120$ queries per data. The query budget for the classification head is $60$ queries per data.

We use the CLS token on top of the Transformer backbone to calculate the cosine similarity between different data points, $d_{i,j}$, in the kernel for reparameterization. The pruned queries for the BA will be equally allocator to the other unpruned data points to ensure the total number of queries is the same.

Table: Table 3: Results of ViT-base under natural noise and adversarial noise.

Evaluation of robustness: Furthermore, we assess the robustness of BP, OPS-LR, and OPS-SPSA across three data sets—ImageNet, Food101, and CIFAR100—through three distinct criteria: 1) Primary task efficacy: accuracy of classification on unaltered samples (Original); 2) Endurance against natural corruption: accuracy of classification on datasets tainted with natural noise, incorporating Gaussian, uniform, and Poisson disturbances; 3) Robustness to adversarial attacks: accuracy of classification on samples altered by adversarial tactics, specifically the fast gradient sign method (FGSM), iterative fast gradient sign method (I-FGSM), and momentum-based iterative fast gradient sign method (MI-FGSM). For adversarial assaults, we cap the allowable perturbation per pixel at $8/255$, and for I-FGSM and MI-FGSM, we limit the iteration count to a maximum of 5. In our evaluations, we encompass the entire test set for corruption assessment. As shown in Table 3, BP achieves the best accuracy on the original data. However, the robustness of OPS-LR and OPS-SPSA is superior to BP. Especially under adversarial attacks, the perturbation methods have a significant advantage over BP.

### C.2 Experiments for Black-box tuning

We use the open-source CLIP with ViT-B/32 as the visual encoder. The intrinsic dimension, $d_{I}+d_{T}$, is 1000. The visual prompt length is $10$ and the text prompt length is $12$. The batch size is 64. We use 60 queries per data to tune the intrinsic dimension. We use Adam optimizer with a learning rate of $1\times 10^{-4}$ and train for 50 epochs with early stopping. All methods use the same 16-shot split for training and are evaluated on the full test sets for evaluation.

### C.3 Experiments for multimodal alignment

we use Vicuna v1.5 7B as the Large Language Model and train the 7B version. We train one epoch for the feature alignment by an Adam optimizer with the learning rate of $5\times 10^{-4}$ and cosine learning rate decay. The batch size is 64 and the query budget is 60 queries per data.

### C.4 Extended ablation

Figure: (a)
Refer to caption: x5.png

We provide a more extensive ablation study on ViT-Large following the same setting as the main text. Figure 5(a)-5(c) and Figure 5(d)-5(f) show the result of LR and SPSA estimator. Both the estimators can benefit from the allocation and the Gaussain allocator has better performance than the Bernoulli allocator. The gradient of deep layer like layer 1 is more difficult to estimate than the shallow layer and the corresponding variance accounts for most of the variance of the whole model. Comparing Figure 5(a)-5(c) to Figure 5(g)-5(i), estimation on the feed-forward network is easier than on the multi-head attention. In practice, it is recommended to add gradient clip to the multi-head attention layer.

## Peer Reviews

confidence: 3
contribution: 2
initial_rating: 3
presentation: 1
questions: For enhanced clarity, please focus on the weaknesses section, which includes my raised questions.

Concerning reproducibility, I reviewed the supplementary material, presumably the code. However, the absence of a README file and the presence of numerous extraneous files make it challenging to determine which files are essential. Additionally, the provided code lacks meaningful comments and contains a lot of debug information, which further complicates understanding its logic. If the authors aim to demonstrate reproducibility through the attached code, I recommend including a clean version of the code with detailed instructions in a README file to at least guide reviewers through the main logic.
soundness: 2
strengths: 1. The empirical accuracy results for ViT and CLIP appear to surpass those of the baselines. The authors have also conducted essential ablation studies to further validate their findings.

2. While the text in Figures 2 and 3 is smaller than the standard text size, making it challenging to read, the color combinations used in these figures are visually appealing.
summary: This paper examines perturbation-based gradient computation methods tailored for forward-only learning. The authors introduce an optimal sampling strategy based on a Gaussian Allocator designed to maximize performance improvements incrementally. They evaluate this approach using pretrained transformers and demonstrate that it outperforms selected baseline methods.
weaknesses: 1. Sections 3 and 4 introduce several undefined annotations, leading to ambiguity in the exposition. For instance, Line 148 mentions a term "a" whose role and relation to Equation (1) are unclear—is it a distribution, a hidden representation, or something else? Furthermore, "G(·)" on Line 160 and "y_j" on Line 164 are undefined, with no clarification of the indexing or distribution from which j is sampled. Additionally, Equation (9)'s term "K" lacks a defined scope. The abbreviation "LR" is also used without definition. The paper introduces many terms that, while relevant, obscure the main contributions. I recommend a thorough review of these sections to clarify the foundational concepts and distinctly outline the problems the paper aims to solve.

2. The rationale behind formulating the optimization of the query allocator as maximizing Equation (5) is not sufficiently clear. A deeper analysis of this formulation is necessary. Moreover, the paper mentions only the initialization of a Gaussian Allocator, which might suggest a broader applicability than is actually the case. Either a comparative analysis of different allocators should be included, or the focus on the Gaussian allocator should be explicitly stated in the abstract and introduction to manage expectations.

3. I have several concerns regarding the experimental results presented in Table 1, which utilizes a pretrained ViT network. Firstly, the rationale for the notably small number of queries for baseline methods, such as only 2 for Mezo, is not explained. What is the expected number of queries for your OPS-LR model in comparison? In Line 344, you mention alignment with Mezo's original memory-efficient setting, but detailed statistical data on memory consumption and training time for both baseline methods and your approach are lacking. Additionally, the experiments are conducted on pretrained models like ViT and CLIP, but the performance of models trained from scratch is not shown. It is crucial to demonstrate whether OPS-LR offers any advantages over these baselines when applied to models not pretrained. Furthermore, the validation datasets used are generally small, with all except ImageNet containing under 150,000 images. There is also concern that using pretrained transformers on the well-known ImageNet dataset could lead to performance biases due to data leakage.

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 2
questions: 1. Please refer to Weakness 1. Is it possible to integrate the proposed method into a callable class or function without rewriting model architectures like Linear and Conv2d to implement the specific add_noise operation?

2. Please refer to Weakness 2. Could you provide detailed memory usage for different methods during training for a more thorough comparison? I noticed that the code uses a repeat operation to expand the batch for varying numbers of queries on different data. Does this operation increase memory usage?

3. The authors mention in the main text that 'All the methods in the experiments use the same query budgets, except for Mezo, which uses only 2 queries per data point in accordance with its original memory-efficient setting.' However, could you provide a more detailed comparison of runtime (e.g., clock time) compared to other ZO methods and the BP baseline?
soundness: 3
strengths: 1. The idea of dynamically allocating different numbers of queries to each data point within a batch during training is novel, which is indeed a point that previous zeorth-order optimization (forward learning) methods have not considered.

2. The proposed method is intuitive. The approach of leveraging a Gaussian Allocator (GA) combined with a likelihood ratio method introduces a creative solution to minimize gradient estimation variance. Through appropriate approximations, the computational cost is effectively reduced. Theoretical analysis is also provided.

3. The experimental setup is extensive and reasonable, and the results are convincing. Both prompt tuning for large models and multimodal alignment for foundation models are promising application scenarios for zeroth-order (ZO) methods, and the proposed approach demonstrates good performance on these tasks.
summary: The paper introduces FLOPS: Forward Learning with Optimal Sampling, which aims to improve the efficiency of gradient estimation in forward-only learning methods by optimally allocating computational resources (queries) across data points within each mini-batch. The approach is motivated by the limitations of backpropagation, particularly in settings where only forward passes are feasible or desirable, such as in black-box optimization scenarios.  With a simplified proxy objective and a reparameterization technique, the authors derive a novel plug-and-play query allocator with minimal parameters. Extensive experiments show the superior performance of this method. Theoretical analysis is also provided.
weaknesses: 1. Although the authors provide part of the source code, I believe the coding is not advisable. Specifically, the authors override nn.Linear to create a custom Linear class and similarly override nn.Conv2d to create a custom Conv2d class. This approach results in the proposed method being tied to a specific model architecture, making it difficult to adapt to other architectures. In fact, existing zeroth-order optimization methods, such as ZO-SGD [1], ZO-AdaMM [2], and DeepZero [3], all have core optimization logic that can be implemented by inheriting from torch.optim.Optimizer, thereby aligning with gradient-based methods like SGD. Alternatively, they can be integrated into a specific function for easier migration.

2. One important reason why zeroth-order optimization is suitable for large model prompt fine-tuning is that these methods do not require backpropagation, which significantly saves memory compared to gradient-based methods like SGD. However, in the experimental results presented in the main paper, only the fine-tuning results are provided, without comparing their memory usage with backpropagation and other zeroth-order optimization methods.

[1] Saeed Ghadimi and Guanghui Lan. Stochastic first-and zeroth-order methods for nonconvex stochastic programming. SIAM Journal on Optimization, 23(4):2341–2368, 2013.

[2] Xiangyi Chen, Sijia Liu, Kaidi Xu, Xingguo Li, Xue Lin, Mingyi Hong, and David Cox. ZO-AdaMM: Zeroth-order adaptive momentum method for black-box optimization. NeurIPS, 32, 2019.

[3] Chen A, Zhang Y, Jia J, et al. Deepzero: Scaling up zeroth-order optimization for deep model training[J]. arXiv preprint arXiv:2310.02025, 2023.

---

confidence: 1
contribution: 3
initial_rating: 6
presentation: 3
questions: 1) I am curious about why other methods that utilize all queries would perform worse than this method, that utilizes limited quries for each data.  Since MEZO is tailored to another baseline, using the same hyperparameters might not be entirely fair. Have you considered tuning MEZO with more available queries per data point for comparison?
2) What is the exact computational cost when using all queries for each data point compared to your allocation method under different budget constraints?

If reasonable answers are provided, I will consider rasing scores accordingly.
soundness: 3
strengths: 1) The motivation is clear: optimizing the allocation of queries to effectively reduce computational overhead.
2) The experimental results show strong performance relative to the baselines.
3) The study provides both experimental and theoretical results, offering a well-rounded evaluation.
summary: This paper proposes an efficient query allocation strategy for forward-learning algorithms in gradient computation, reducing query usage by focusing on data points that need it most. Using a simplified objective and reparameterization, the authors introduce a lightweight query allocator that minimizes gradient estimation variance with low computational cost. Both experiments and theoretical analysis are provided.
weaknesses: 1) I am curious about why other methods that utilize all queries would perform worse than this method, that utilizes limited quries for each data. 
2) The comparison of exact computational cost between equally using all queries for each data point and your allocation method is unknown. However, it is one of the main motivation.  

Minor:   
3) In the abstract, the phrase “propose to allocate the optimal number of queries over each data” isn’t entirely accurate, as a total query budget must be pre-defined rather than learning an optimal number. Perhaps rephrasing to “allocate the optimal number of queries within a set budget” would be clearer.   
4) Table 2 is not well-formatted and appears misaligned.

---

confidence: 4
contribution: 3
initial_rating: 8
presentation: 1
questions: 1. **Clarification on LR and OPS-LR**: The distinction between LR and OPS-LR is not clear. More explanation would be appreciated.

2. **Clarification on experiments**: Do the experiments include cross-validation with multiple random seeds? If not, please show the experiment results with multiple seeds with dataset cross-validation. If yes, please provide more details.

3. **Ablation studies**: The proposed algorithms update four parameters. What if only three or two of them are updated? Which parameters are dispensable for this process? Conducting these experiments would provide more insights into the paper's contributions.
soundness: 2
strengths: 1. The empirical evaluation is comprehensive, demonstrating significant performance improvements through query sampler optimization.
summary: This paper presents an approach for optimizing a differentiable sampler in forward learning in foundation models, supported by theoretical proofs and extensive empirical evaluations. The work lies at the intersection of zeroth-order optimization and sampling optimization.
weaknesses: 1. **Overall outline and structure**: The paper builds upon DeepZero and Mezo by devising an optimization allocation for forward learning. However, the idea of optimal allocation is not new in the ML context. Papers such as "Stochastic Optimization with Importance Sampling" or "A General Analysis of Example-Selection for Stochastic Gradient Descent" (and several derived works) have explored similar concepts. For me, the main difference here is the focus on forward learning (or zeroth-order optimization) rather than backpropagation. The authors are encouraged to: (1) Review and mention this existing research (strongly encouraged) and (2) Compare these methods with their proposed algorithm (encouraged). Addressing these points would strengthen the paper's contribution and contextualize it within the field.

2. **Introduction and framing**:
The decision to begin with biologically plausible algorithms (BioPA) seems unexpected and may not effectively frame the paper's contributions. Nevertheless, the general scheme could be unfolded more clearly:
     1. The citation of Jacot et al. for "learning high-level representation" appears unrelated in the BioPA context.
     2. Consider including more relevant BioPA works such as "Direct Random Target Projection" [1], "SoftHebb" [2], and "Counter-Current Learning" [3]

3. **Writing and proofreading**:
   - Correct typographical errors (e.g., "Current Literature" should be "Current literature")
   - Address factual inaccuracies (e.g., L48 states that the FF algorithm is only capable of training MLPs on MNIST, but results for CIFAR are also presented)
   - Provide explanations for abbreviations (e.g., SPSA, LR)

4. **Related Work section**: The first subsection could be restructured. When discussing backpropagation-free learning, it's typically in the context of multi-layered neural networks. Also, including evolution theory and particle swarm optimization seem tangential. I suggest reorganizing this section and incorporating the suggested papers for a more focused discussion.

References

[1] "Learning Without Feedback: Fixed Random Learning Signals Allow for Feedforward Training of Deep Neural Networks" (Frontiers in Neuroscience, 2021)

[2] "Hebbian Deep Learning Without Feedback" (ICLR 2023)

[3] "Counter-Current Learning: A Biologically Plausible Dual Network Approach for Deep Learning" (NeurIPS 2024)

Meta-review: **Summary:** The paper proposes a method for perturbation-based forward gradient learning. The authors claim that the technique has lower variance and is more scalable than the previous approaches. Their method relies on importance estimation for data points and unifies the previous approaches. They provide theoretical results for optimality and show results on foundation models.  

**Strengths:** The reviewers believe that the idea is novel and the approach is intuitive. The method combines Gaussian Allocator (GA) with a likelihood ratio for importance estimation. The experimental results are thorough and promising and might be the strongest aspect of the paper. Some results such as ViT and CLIP appear to surpass those of the baselines, and the additional results on foundational models show good performance on these tasks. The reviewers do not raise any concerns about the theoretical results.  

**Weaknesses:** Some of the notations and definitions are ambiguous. The motivations are unclear at parts, for instance, the query allocator as maximizing Equation (5). Memory usage compared to backpropagation and other zeroth-order optimization methods is not provided. Also, Reviewer HWGc claims that ``the ​​exact computational cost between equally using all queries for each data point and your allocation method is unknown''.  

**Decision:** The paper addresses an important and practical problem in ML, which has immediate applications for more efficiently training large-scale models. The authors have done an excellent job of providing experimental evidence for the applicability of their method. The main weaknesses are related to 1) writing, proofreading, and presentation, which the authors have improved greatly during the rebuttal period, 2) some missing experiments and justifications, which the authors provide further additional results to address. Thus, based on the vote by the majority of the reviewers, and the authors’ response during the rebuttal period, I recommend acceptance.

Decision: Accept (Poster)