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) [1,2,3]. Transformers are one of the key players in this development [3]. Initially designed in the context of NLP, they were adapted for Computer Vision tasks through the introduction of Vision Transformers (ViTs) [4]. The working principle of ViTs is similar to those of Transformers, but instead of dividing a sentence into words, they split an image into non-overlapping rectangular patches and look for semantic correlations between them. ViTs have proven to be very competitive, and in some contexts, their performance has been superior to that of Convolutional Neural Networks (CNNs) [4]. The main problem with ViTs is their computational cost, since for each layer it is necessary to compute the attention of each token with respect to all others. To overcome this problem, several variants of ViTs have been proposed to reduce the cost of the attention layers [5,6,7,8,9].
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 [10,11,12,13,14,15].
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 [16] that returns a list of selected patches. The agent is composed of three dense layers. For each training batch, it observes the attention values produced by the first ViT layer and returns a subset of the original patches to use for the training of the ViT. After a certain number of training epochs, the agent receives a reward that takes into account training loss and training time. The user can decide how much weight to give to each of these two parameters, thus favoring a set of patches that guarantees a low training time or one that guarantees a low training loss.
Several approaches have been proposed in the literature to reduce the computational load of attention layers. They are based on different techniques such as quantization [17,18,19,20,21,22,23], pruning [24,25,26,27], low-rank factorization [28,29] and knowledge distillation [30,31,32,33,34,35]. Other approaches perturbate the input of a ViT to optimize the resources it uses [36,8,37]. Others compute the importance of each token and remove less important tokens as inference proceeds [9,6,7,38,39,40]. AgentViT shares with some of the above approaches the policy of setting a variable number of tokens based on the input images. This allows it to fit the images in the best possible way. However, it has a completely different fitting mechanism than the other approaches. In fact, the latter requires the user to specify the maximum number of tokens to be used. If, after resampling, the number of tokens is greater than the number specified by the user, the excess tokens are removed. AgentViT also allows the user to specify the maximum number of tokens desired, and it trains its RL agent to select a number of tokens as close as possible to the number specified by the user. However, if it obtains a particularly low training loss during training, its reward mechanism will prompt it to select a smaller number of tokens for that particular batch of images. Conversely, if it obtains a high training loss for a particular batch, its reward mechanism will prompt it to increase the number of tokens to be used, giving less weight to the number specified by the user.
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, 𝒫 : 𝒮 × 𝒜 → 𝒮 is a 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 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 (𝑠, 𝑎). Q-Learning 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 (𝑠, 𝑎).
Here:
• 𝜂 is the learning rate; it is a number in the real range [0, 1] and specifies the rate at which the agent learns; • 𝛾 is the discount factor; it represents the importance of the immediacy of the reward. If 𝛾 is closer to 0, actions with immediate rewards are favored; if 𝛾 is closer to 1, all rewards are given equal weight, regardless of their immediacy, which favors a long-term view; • 𝑠 * is the next state, i.e., the state in which the agent arrives when it starts from 𝑠 and executes the action 𝑎.
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.
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 different 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]:
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 difference). 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:
Here:
• ℛ 𝑙𝑜𝑠𝑠 𝑡 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 same function at iteration 𝑡.
• ℛ 𝑝𝑎𝑡𝑐ℎ 𝑡 is the reward related to the number of patches; it is defined as the ratio of the difference 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 [0, 1], that determines the weight to assign to ℛ 𝑙𝑜𝑠𝑠 𝑡 with respect to ℛ 𝑝𝑎𝑡𝑐ℎ 𝑡 .
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 ℛ 𝑙𝑜𝑠𝑠 𝑡 and ℛ 𝑝𝑎𝑡𝑐ℎ 𝑡 in Equation 2.3. The weight 𝛼 allows the user to specify how 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 different 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 [6]; the latter, to the best of our knowledge, is the approach most similar to AgentViT. For each of these approaches, we computed their Cumulative Training Time (measured in seconds), Accuracy, Precision, Recall, and F1-Score. Table 2 shows the corresponding values.
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 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-off between accuracy and training time. Finally, AgentViT is the approach capable of providing the best trade-off 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 difficulty 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-off 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 different ViT architectures, possibly including multiple attention layers, which would make our framework more robust and versatile.