In recent years, Transformers have led a revolution in Natural Language Processing, and Vision Transformers (ViTs) promise to do the same in Computer Vision. The main obstacle to the widespread use of ViTs is their computational cost. Indeed, given an image divided into a list of patches, ViTs compute, for each layer, the attention of each patch with respect to all others. In the literature, many solutions try to reduce the computational cost of attention layers using quantization, knowledge distillation, and input perturbation. In this paper, we aim to make a contribution in this setting. In particular, we propose AgentViT, a framework that uses Reinforcement Learning to train an agent whose task is to identify the least important patches during the training of a ViT. Once such patches are identified, AgentViT removes them, thus reducing the number of patches processed by the ViT. Our goal is to reduce the training time of the ViT while maintaining competitive performance.
In recent years, thanks also to the massive development of deep learning systems, Artificial
Intelligence is experiencing a golden age in many sectors, including Natural Language Processing
(NLP) and Computer Vision (CV) [
Another area of Artificial Intelligence that has shown great promise in recent years is
Reinforcement Learning (RL). In fact, it is being applied in a wide range of contexts, from robotics
to intelligent transportation systems [
In this paper, we propose AgentViT, a framework for ViT optimization. To achieve its goal,
AgentViT uses RL to reduce the computational complexity of the attention layer and thus the
training time of ViTs. In AgentViT, an RL agent selects a subset of the image patches so that the
ViT has to process only them for its training, thus reducing the training time while maintaining
competitive performances. The RL agent is a Deep Q-Learning Network [
Several approaches have been proposed in the literature to reduce the computational load
of attention layers. They are based on diferent techniques such as quantization [
This paper is organized as follows: in Section 2, we describe AgentViT. In Section 3, we illustrate our experimental campaign aimed at determining the values of its hyperparameters, comparing it with related approaches, and deriving interesting insights. Finally, in Section 4, we draw our conclusions and look at some possible future developments.
AgentViT uses a Markov Decision Process (MDP) ℳ = ⟨, , ℛ, , ⟩ [41]. Here, is a state space, is a discrete set of actions, ℛ : × → R is a reward function, : × → transition kernel, and ∈ [0, 1) is a discount factor. In an MDP, a stationary policy : → is a mapping from states to actions; it specifies the action an agent takes when it is in a given is a state. It is used to describe how an agent interacts with the environment.
AgentViT uses an Action Value Function (, ), introduced in Q-Learning [42], to estimate the expected cumulative reward an RL agent can obtain from a given state-action (, ). QLearning uses a table with a row for each observable state and a column for each possible action. As the algorithm runs, the values in the table are updated using the formula expressed in Equation 2.1. This allows us to recursively obtain the cumulative reward (, ) associated with each action-state pair (, ).
︂[ (, ) + · ℛ, + · m′∈ax((* , ′)) − (, ) ︂] (2.1) Here:
the agent learns;
• is the learning rate; it is a number in the real range [
The Q-Learning algorithm struggles in the presence of a large number of states or when the states involved are continuous. In these cases, the table is replaced by a neural network called Deep Q-Network. It receives the vector representing the state and computes the Q-values corresponding to each pair (, * ), where * represents any action the agent can take in . The agent chooses the action with the highest Q-value. Deep Q-Network uses this reward to tune its behavior for the next step.
Attention Scores
Reward
Selected
Patches Deep Q-Network
Vision
Transformer 2.2. State In the context of the MDP, a state ∈ observed by the agent is represented by a vector of real numbers and models the current conditions of the environment. In AgentViT, the state of an environment is represented by the attention score obtained from a batch of images processed by the first attention layer of the ViT.
More specifically, given an image composed of patches, the output of the attention layer is represented by an × matrix, where is the embedding size of the image. This matrix is given in input to a ViT module downstream of the attention layer, which transforms the matrix itself into a vector of elements obtained by averaging the values along the dimensions. This vector represents the average values of attention (and thus importance) that the attention layer assigns to the diferent patches. It is also the state given as input to AgentViT’s Deep Q-Network.
Given a vector of elements representing a state ∈ , the agent returns a list of elements. Each element is associated with a patch and is a real value that, according to the Deep Q-Learning algorithm, represents the cumulative reward associated with the corresponding patches as estimated by the agent [42]. The list is sorted in descending order so that the first elements represent the most important patches. AgentViT selects all patches in the list whose associated value is greater than the average value of the elements in the list. Therefore, the action ∈ associated with the state ∈ corresponds to the selection of the most promising patches.
The transition kernel function (see Section 2.1) is the one provided by Deep Q-Learning. Based on it, after the agent chooses an action ∈ and receives a reward ℛ,, the update of the Q-value associated with the pair (, ) is done through the following formula [42]: ︂[
(, ) + · ℒ (ℛ, + · m′∈ax((* , ′)), (, ))
Similarly to Equation 2.1, this formula describes the update of (, ) by taking into account the previous value and the distance between the maximum cumulative reward associated with the next state * and the Q-value associated with the current state. In AgentViT, we adopted a Huber function ℒ [43] to compute this distance (unlike [42] that used an algebraic diference). The reasoning behind this choice is that ℒ is not sensitive to outliers and, in some cases, prevents the gradient explosion problem. The cumulative reward for the next state is predicted by the Target Network, which consists of an exact copy of the agent network except that its weights are not updated by backpropagation, but are periodically copied from the agent network by a soft-copy mechanism. As shown in [44], this way of proceeding allows us to stabilize the learning process.
AgentViT also has a mechanism to avoid falling into a local minimum. Indeed, the agent chooses a random action with a probability equal to instead of the action that maximizes the value of . The value of decays exponentially as training progresses to avoid instability. In this way, AgentViT is able to ensure good exploratory analysis in the early stages of ViT training and good stability of results as training progresses.
Finally, AgentViT uses a replay memory [45, 46] to improve the stability and generalizability of the agent. It can store observed data for later use during training in a way that breaks unwanted temporal correlations.
In AgentViT, the reward ℛ obtained at iteration by starting from a state and executing an action plays a key role since it serves to define the quality of training. As mentioned above, this quality must take into account the time required for training and the accuracy. Consequently, ℛ must consider both the training loss and the training time. For this purpose, it is defined as a weighted mean of the training loss and the number of patches selected by the agent, which is proportional to the training time.
Based on this reasoning, ℛ can be formulated as: ℛ = · ℛ + (1 − ) · ℛ ℎ (2.3) Here: • ℛ • ℛ same function at iteration .
is the reward related to the training loss; it is equal to the ratio between the value
ℒ(0) of the loss function of the ViT at the starting iteration and the value ℒ() of the
ℎ is the reward related to the number of patches; it is defined as the ratio of the
diference between the actual number of patches selected by the agent and the user’s
desired number of selected patches, to the user’s desired number of selected patches.
• is a value belonging to the real interval [
ℎ.
In this way, the agent is incentivized to select a number of patches close to the number desired by the user (or a very small number if the user does not specify a value). But, it is also incentivized to select a subset of patches that can minimize training loss. These two goals are represented by ℛ ℎ in Equation 2.3. The weight allows the user to specify how and ℛ much importance to place on each of these goals. If tends to 1, the agent has an incentive to choose a large number of patches. On the other hand, if tends to 0, it is incentivized to minimize the number of patches selected, subject to the accuracy constraints to be achieved.
AgentViT can be applied to any ViT, since it is based on observing the attention scores it returns. Consequently, in our experiments, we could have applied AgentViT to any ViT proposed in the literature. To conduct the experiments in a reasonable time, we decided to employ SimpleViT [47] that splits images into 64 patches (SimpleViT64), since it can be trained faster than a classical ViT. We performed our experiments on the CIFAR10 dataset, which is a collection of 60,000 color images 32x32 divided into 10 diferent classes, designed for training and testing machine learning models in computer vision tasks. CIFAR10 is widely used for benchmarking classification algorithms in the deep learning field. For the training and testing phases of our experiments, we used Google Colab, which provides an Intel Xeon CPU with 2 vCPUs, 13 GB of RAM, and an NVIDIA Tesla K80 GPU with 12 GB of VRAM. We refer the reader to the link https://github.com/DavideTraini/RL-for-ViT for the code used to implement AgentViT.
As a first step in our experiments, we had to define the values of the hyperparameters of AgentViT. Due to space limitations, we cannot report in detail the tasks we performed to determine these values. At the end of these tasks, we obtained the values reported in Table 1.
As a next step, we decided to compare the performance of AgentViT with that of related
approaches already proposed in the literature. Specifically, the approaches we considered
for comparison are the original ViT, SimpleViT64, and ATSViT [
This table shows that there are approaches able to guarantee low values of Cumulative Training Time, but at the expense of Accuracy, Precision, Recall, and F1-Score. Conversely, other approaches can obtain high values of these measures, but at the expense of Cumulative Training Time. AgentViT is able to achieve a suboptimal value for all five metrics. In other words, it does not achieve the maximum value for any metric but is able to ensure the best compromise among all metrics.
As seen above, the goal of AgentViT is to use RL to select patches for optimal filtering. Comparison with other related approaches has shown us that it has a satisfactory performance. In fact, it is able to train a ViT in less time than the ViT itself, if it is trained without removing Hyperparameter
Value patches. This saving in training time does not come at the expense of accuracy, which remains comparable to that of the SimpleViT trained without patch removal. Moreover, AgentViT allows the user to specify the desired trade-of between accuracy and training time. Finally, AgentViT is the approach capable of providing the best trade-of between Cumulative Training Time on the one hand, and Accuracy, Precision, Recall, and F1-Score on the other hand.
In addition, AgentViT has other interesting implications. One of them is the possibility of using larger Vision Transformers. In fact, AgentViT’s ability to reduce training time makes it possible to adopt architectures that would not normally have been adopted due to their excessive computational load. A second implication concerns the use of AgentViT to build smaller synthetic datasets from the original ones, which can be used to train deep neural networks. A further implication concerns the possibility of extending the use of AgentViT to contexts other than Vision Transformers. In fact, the idea behind AgentViT is general and independent of the type of transformers to which it is applied; therefore, it could work with any transformers, such as those used in the context of NLP. The only condition is that the input of the RL agent within AgentViT can receive an attention matrix as input.
Finally, we highlight some limitations of AgentViT. The first concerns the fact that the decision to use Deep Q-Learning within AgentViT involves the need to set various hyperparameters, which makes the setup phase rather complex. A second limitation is related to the number of patches required for AgentViT to work properly. In fact, if the Vision Transformer underlying AgentViT works with only a few patches, the agent has dificulty selecting the most important ones.
In this paper, we proposed AgentViT, a framework that uses RL to reduce the training time of a ViT without significantly reducing its performance. The RL agent present in AgentViT uses a classical MDP-based mechanism to represent an environment for the image classification task. As for this process, we redefined the state, action, and reward needed to train our RL agent. We tested AgentViT using SimpleViT64 as the internal ViT and Deep Q-Network as the internal RL agent. Our experiments showed that AgentViT can achieve the best trade-of between Cumulative Training Time on the one hand and Accuracy, Precision, Recall, and F1-Score on the other hand. The experiments conducted allowed us to draw several implications regarding the strengths and limitations of our framework.
We can think of several possible future developments of our approach. For example, we could improve the reward function to consider validation loss instead of training loss. We could also define a new metric, similar to the Akaike information criterion [ 48], which takes into account both model performance and the number of tokens. Moreover, we could test other Reinforcement Learning algorithms, such as Multi-Agent RL and Contextual Multi-Armed Bandit, instead of the Deep Q-Network, and see if they can further improve the performance of ViTs. These algorithms could assist in selecting the best actions and the corresponding patches to speed up ViT training. Finally, we could evaluate the impact of our approach on diferent ViT architectures, possibly including multiple attention layers, which would make our framework more robust and versatile. Recognition (CVPR’22), New Orleans, LA, USA, 2022, pp. 12165–12174. [39] Y. Rao, Z. Liu, W. Zhao, J. Zhou, J. Lu, Dynamic spatial sparsification for eficient vision transformers and convolutional neural networks, IEEE Transactions on Pattern Analysis and Machine Intelligence (2023). IEEE. [40] L. Meng, H. Li, B. Chen, S. Lan, Z. Wu, Y. J. e S.N. Lim, Adavit: Adaptive vision transformers for eficient image recognition, in: Proc. of the International Conference on Computer Vision and Pattern Recognition (CVPR’22), New Orleans, LA, USA, 2022, pp. 12309–12318. [41] C. L. Lan, S. Tu, A. Oberman, R. Agarwal, M. Bellemare, On the generalization of representations in reinforcement learning, arXiv preprint arXiv:2203.00543 (2022). [42] K. Arulkumaran, M. Deisenroth, M. Brundage, A. Bharath, Deep Reinforcement Learning:
A brief survey, IEEE Signal Processing Magazine 34 (2017) 26–38. IEEE. [43] P. J. Huber, Robust estimation of a location parameter, in: Breakthroughs in statistics:
Methodology and distribution, Springer, 1992, pp. 492–518. [44] J. Fan, Z. Wang, Y. Xie, Z. Yang, A theoretical analysis of deep Q-learning, in: Proc. of the International Conference on Learning for Dynamics and Control (L4DC’20), Berkeley, CA, USA, 2020, pp. 486–489. PMLR. [45] R. Liu, J. Zou, The efects of memory replay in reinforcement learning, in: Proc. of the Annual Allerton Conference on Communication, Control, and Computing (Allerton’18), Monticello, IL, USA, 2018, pp. 478–485. IEEE. [46] L. Lin, Self-improving reactive agents based on reinforcement learning, planning and teaching, Machine Learning 8 (1992) 293–321. Springer. [47] L. Beyer, X. Zhai, A. Kolesnikov, Better plain ViT baselines for ImageNet-1k, arXiv preprint arXiv:2205.01580 (2022). [48] H. Akaike, A new look at the statistical model identification, IEEE Transactions on Automatic Control 19 (1974) 716–723. IEEE.