SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Baohao Liao Shaomu Tan Christof Monz Language Technology Lab, University of Amsterdam {b.liao, s.tan, c.monz}@uva.nl
Parameter-efficient fine-tuning (PEFT) of pre-trained language models (PLMs) has emerged as a highly successful approach, with training only a small number of parameters without sacrificing performance and becoming the de-facto learning paradigm with the increasing size of PLMs. However, existing PEFT methods are not memory-efficient, because they still require caching most of the intermediate activations for the gradient calculation, akin to fine-tuning. One effective way to reduce the activation memory is to apply a reversible model, so the intermediate activations are not necessary to be cached and can be recomputed. Nevertheless, modifying a PLM to its reversible variant is not straightforward, since the reversible model has a distinct architecture from the currently released PLMs. In this paper, we first investigate what is a key factor for the success of existing PEFT methods, and realize that it's essential to preserve the PLM's starting point when initializing a PEFT method. With this finding, we propose memory-efficient fine-tuning (MEFT) that inserts adapters into a PLM, preserving the PLM's starting point and making it reversible without additional pre-training. We evaluate MEFT on the GLUE benchmark and five question-answering tasks with various backbones, BERT, RoBERTa, BART and OPT. MEFT significantly reduces the activation memory up to  $84\%$  of full fine-tuning with a negligible amount of trainable parameters. Moreover, MEFT achieves the same score on GLUE and a comparable score on the question-answering tasks as full fine-tuning. A similar finding is also observed for the image classification task. $^{1}$
Large-scale pre-trained models have achieved great success across various domains and applications [1, 2, 3, 4, 5, 6, 7, 8]. As their capabilities continue to evolve, the released pre-trained language models (PLMs) have grown exponentially in size, even reaching a scale of 100 billion parameters [3, 9, 10, 11, 12]. Consequently, it presents unprecedented challenges in effectively leveraging these models for downstream tasks due to limited computing resources.
A historically common approach to adapting PLMs to downstream tasks is updating all pre-trained parameters, full fine-tuning. Although full fine-tuning has yielded numerous state-of-the-art results, its applicability is limited in storage-constrained environments. This constraint arises from maintaining a complete copy of the fine-tuned model for each task. An alternative
Figure 1: Average performance of different tasks vs. activation memory. The memory usage for full finetuning is denoted as  $100\%$ .
$^{1}$ Code at https://github.com/baohaoliao/mefts. Up-to-date version at https://arxiv.org/abs/2306.00477.
adaptation approach is parameter-efficient fine-tuning (PEFT) [13, 14, 15, 16, 17, 18, 19] which involves selectively updating a small number of task-specific parameters while keeping the majority of the PLM's parameters frozen. PEFT offers significant advantages in reducing storage requirements by only saving one general PLM alongside the modified parameters for each task. In addition to storage savings, PEFT achieves comparable performance to full fine-tuning, sparking considerable interest in the adoption of PEFT.
Despite their advantages in parameter efficiency, existing PEFT methods still face challenges in terms of memory efficiency [20, 21]. PEFTs necessitate the caching of intermediate activations, similar to the requirements of full fine-tuning, to calculate the gradients of the trainable parameters. Typically, they consume more than  $70\%$  activation memory of full fine-tuning (see Figure 1). Since activations significantly contribute to the memory requirements during training, there are instances where fine-tuning a large-scale PLM with PEFT is not feasible due to memory constraints. To address this issue, a commonly employed approach is to treat the PLM as a feature extractor, such as knowledge distillation to a smaller model [22, 23], adding additional trainable layers on top [20] or aligned [21, 24] with it, and so on. These approaches circumvent the need to store the PLM's activations since the gradient computation graph does not traverse through the PLM. However, these methods often require additional pre-training or exhibit a substantial performance gap compared to full fine-tuning when using the same underlying model [20, 21].
In this paper, we propose a novel method called memory-efficient fine-tuning (MEFT) to modify PLMs in a parameter- and memory-efficient manner, without requiring additional pre-training. Initially, we investigate a crucial factor for the success of existing PEFT methods and determine that the proper initialization of newly added parameters is essential to maintain the continuity of information from the PLM ( $\S 2$ ). Leveraging this insight, we design three MEFT methods that enable the modification of a PLM to its reversible variant, so it only necessitates caching the final output and allows for the recomputation of intermediate activations during back-propagation ( $\S 3$ ). Consequently, MEFT significantly reduces the memory required for caching activations (see Figure 1).
To validate the effectiveness of our MEFT methods, we conduct extensive evaluations on the GLUE benchmark [25] with BERT [1], RoBERTa [2] and BART [26] ( $\S 4$ ). The experimental results consistently demonstrate that our MEFT methods outperform both full fine-tuning and strong PEFT baselines in terms of parameter and memory efficiency. Remarkably, our methods achieve the same score as full fine-tuning while updating only  $0.2\%$  of the parameters and saving up to  $84\%$  of the activation memory. Furthermore, we evaluate MEFT on five question-answering tasks with a larger model, OPT [9]. The results show that our approach achieves a comparable score as full fine-tuning while saving  $50\%$  of the activation memory and updating only  $0.64\%$  of the parameters. A similar finding is also observed on the image classification task, SVHN [27]. Collectively, these experiments establish the effectiveness of MEFT as a powerful parameter- and memory-efficient approach that does not compromise performance.
In this section, we aim to provide essential background knowledge by addressing the following questions: (1) Why are existing PEFTs not sufficiently memory-efficient ( $\S 2.1$ )? (2) What is a key factor for the success of PEFT ( $\S 2.2$ )? (3) What challenges does a reversible model have ( $\S 2.3$ )?
2.1 Parameter-efficient fine-tuning is not sufficiently memory-efficient
Given a  $N$  multilayer perception:  $\pmb{h}_N = f_N(f_{N - 1}(\dots (f_2(f_1(\pmb{h}_0))))\dots))$  with  $\pmb{h}_0$  as the initial input, the  $n^{th}$  layer  $\pmb{h}_n = f_n(\pmb{h}_{n - 1}) = \sigma_n(\pmb{W}_nh_{n - 1})$  consists of a nonlinear function  $\sigma_{n}$  and a weight matrix  $\pmb{W}_{n}$ , where the bias term is ignored for simplicity. Denoting  $\pmb{x}_n = \pmb{W}_nh_{n - 1}$ , in backpropagation with a loss  $\mathcal{L}$ , the gradient of  $\pmb{W}_{n}$  is calculated with the chain rule as:
$$
\frac {\partial \mathcal {L}}{\partial \boldsymbol {W} _ {n}} = \frac {\partial \mathcal {L}}{\partial \boldsymbol {h} _ {N}} \left(\prod_ {i = n + 1} ^ {N} \frac {\partial \boldsymbol {h} _ {i}}{\partial \boldsymbol {x} _ {i}} \frac {\partial \boldsymbol {x} _ {i}}{\partial \boldsymbol {h} _ {i - 1}}\right) \frac {\partial \boldsymbol {h} _ {n}}{\partial \boldsymbol {x} _ {n}} \frac {\partial \boldsymbol {x} _ {n}}{\partial \boldsymbol {W} _ {n}} = \frac {\partial \mathcal {L}}{\partial \boldsymbol {h} _ {N}} \left(\prod_ {i = n + 1} ^ {N} \boldsymbol {\sigma} _ {i} ^ {\prime} \boldsymbol {W} _ {i}\right) \boldsymbol {\sigma} _ {n} ^ {\prime} \boldsymbol {h} _ {n - 1} \tag {1}
$$
where  $\sigma^{\prime}$  is the derivative of  $\sigma$  and the calculation of  $\sigma_{n}^{\prime}$  requires  $\pmb{x}_{n}$ . Therefore,  $\{x_i\}_{i = n}^N$  are cached during the forward pass to obtain the gradient of  $W_{n}$ , even though  $\{W_{i}\}_{i > n}$  are frozen.
(a) Memory trade-off.
(b) Initialization effect, Left: LoRA, Right: (IA) $^3$ .
Figure 2: Exploration of existing PEFTs: (a) The trade-off between memory and the number of trainable parameters. The dashed and solid lines denote the peak and activation memory, respectively. The model size for  $\mathrm{BERT}_{\mathrm{base}}$  is  $0.4\mathrm{GB}^2$ . (b) The initialization effect of PEFT on  $\mathrm{RoBERTa}_{\mathrm{base}}$ . Random PLM denotes that we initialize the backbone randomly instead of using a pre-trained model.
During training, the peak memory footprint is mainly occupied by three components: model's parameters  $\{W_{n}\}_{n = 1}^{N}$ , optimizer state whose size is three times as large as the size of trainable parameters for Adam [28] (one for gradient and two for moments), and activations. The memory footprint for all three components is related to the model's depth and width. In addition, the memory footprint for activations is also related to some training settings, like batch size and sequence length.
Compared to full fine-tuning, existing PEFT methods, such as (Houlsby and Pfeiffer) Adapters [14, 16], LoRA [17],  $(\mathrm{IA})^3$  [29], Prompt-Tuning [19] and Prefix-Tuning [15], tune a small number of parameters, making the size of the optimizer state negligible. However, the memory footprint required for activations is not significantly reduced. As shown in Figure 2a, where we set the batch size as 64 and the sequence length as 512 on RTE [30, 31, 32, 33] with  $\mathrm{BERT}_{\mathrm{base}}$  [1], the activation memory of all PEFT methods is  $>75\%$  of full fine-tuning, even with  $<1\%$  trainable parameters.
2.2 Initialization is significant for parameter-efficient fine-Tuning
Pre-trained models learn generic and distributed enough representations to facilitate downstream learning of highly pressed task representation [36], i.e. offering a robust starting point for the training of downstream tasks. When modifying a PLM with PEFT, we hypothesize that one needs to preserve this starting point at the beginning of training for better performance.
The Starting Point Hypothesis. When modifying a pre-trained model by adding new parameters, one needs to initialize the new parameters in a way to preserve the starting point from the pretrained model at the beginning of training, such that fine-tuning the modified model can match the performance of full fine-tuning.
More formally, supposed  $f_{n}$  is a PLM layer and  $h_n = f_n(h_{n - 1})$ , the output from a modified layer  $f_{n}^{\prime},h_{n}^{\prime} = f_{n}^{\prime}(h_{n - 1})$ , should be close to  $h_n$  at the beginning of training. I.e.  $h_n^\prime = h_n + \delta$ , where  $\| \delta \| \to 0$ . Intuitively, we want  $h_n^\prime \approx h_n$ , because  $h_n^\prime$  is the input to the next (modified) PLM layer. If they are dissimilar, the representation continuity will be broken down. Though most PEFT methods [14, 16, 17, 29] initialize their added modules in this way, we couldn't find a thorough investigation of the significance of this initialization in the existing literature. In this section, we explore the significance of PEFT's initialization for two methods, LoRA and (IA)<sup>3</sup> [29].
LoRA and (IA) $^3$  represent two common methods for introducing new parameters, involving addition and scaling operations, respectively. Given a pre-trained weight matrix  $\pmb{W} \in \mathbb{R}^{d \times d}$ , LoRA modifies it as  $\pmb{h}' = (\pmb{W} + \frac{\alpha}{r} \pmb{W}_{down} \pmb{W}_{up})\pmb{h}$ , where  $\pmb{W}_{down} \in \mathbb{R}^{d \times r}$  and  $\pmb{W}_{up} \in \mathbb{R}^{r \times d}$  are the added trainable parameters,  $\alpha$  is a constant scale factor and normally  $r \ll d$ . LoRA's default initialization is  $\pmb{W}_{down} \sim \mathcal{N}(0, \sigma^2)$  and  $\pmb{W}_{up} = \mathbf{0}$ . In this way,  $\pmb{W}_{down} \pmb{W}_{up} = \mathbf{0}$  and the starting point from the
Though we train in FP16, the PLM is first loaded in FP32, then auto-casted to FP16 for the forward pass in Transformers [34]. Since the memory required for the model in FP32 is always there during training, we report the memory for models in FP32 in this paper (see Table 9). More discussions about this are here. We believe it's a bug in the framework and can be resolved with further investigation. Especially Huggingface's new PEFT framework [35] allows loading INT8 model for fine-tuning.
(a) Reversible architecture.
Figure 3: (a)  $\mathcal{F}$  and  $\mathcal{G}$  are two arbitrary functions (sub-networks), taking two inputs,  $h_n^1$  and  $h_n^2$  (b) Reconstruction error between the vanilla and reversible gradients. The default setting is RevViT [40] with 8 layers,  $\lambda = 1$ ,  $\beta = 1$ ,  $\mu = 0$  and  $\sigma = 0.02$ . Left: Different number of layers. Middle: Different scaling values. Right: Initialization with different means and standard deviations.
(b) Instability of reversible model.
PLM is preserved perfectly. (IA) $^3$  modifies  $\pmb{W}$  by multiplying it to a trainable vector  $l\in \mathbb{R}^d$  as  $\pmb {h}' = (\pmb {l}\odot \pmb {W})\pmb {h}$ , where  $\odot$  represents element-wise multiplication. The default initialization of (IA) $^3$  is  $l = 1$ , also making the starting point untouched.
To facilitate the initialization process of LoRA, we opt for the following initial values:  $W_{down} = 1$ ,  $W_{up} = c$  and  $\alpha = 1$ , where  $c$  is a matrix with all elements equal to an initialized value  $c$ , resulting in  $\frac{\alpha}{r} W_{down} W_{up} = c$ . When  $c = 0$ , the starting point from a PLM is preserved. By adjusting  $c$ , we exert control over the degree of departure from the starting point. Similarly, we replace  $l$  with  $l' = \alpha l = \alpha c$  for (IA) $^3$ .
In Figure 2b, we train the newly added parameters on RoBERTa $_{\text{base}}$  [2] for four tasks (CoLA [37], STS-B [38], MRPC [39] and RTE [30, 31, 32, 33]). For LoRA  $(r = 8)$ , though we modify the initialization method, our result  $(c = 0)$  is very close to the default initialization. When the starting point is broken by  $c \neq 0$  ( $\alpha = 1$ ), all results are even worse than a randomly initialized model. However, when we set  $\alpha = 0^3$  to preserve the starting point, all results become much better than the ones with  $\alpha = 1$ . For (IA) $^3$ , when we decrease  $c$  from 1 (default initialization) to 0, the results  $(\alpha = 1)$  become worse and worse. However, when we set  $\alpha = 1/c$  to preserve the starting point, all results become better. Some of them are even better than the default initialization. All of the above-mentioned results show that it's significant to preserve the starting point from a PLM at the beginning of training when applying or designing a PEFT method. A different initialization scheme is in Figure 10 which leads to a similar finding.
2.3 Challenges of reversible neural network
Recapping a reversible model [41] in Figure 3a, one can reconstruct inputs from outputs as:
$$
\boldsymbol {h} _ {n + 1} ^ {1} = \lambda \boldsymbol {h} _ {n} ^ {1} + \mathcal {F} _ {n} \left(\boldsymbol {h} _ {n} ^ {2}\right)
$$
$$
\boldsymbol {h} _ {n} ^ {2} = \left(\boldsymbol {h} _ {n + 1} ^ {2} - \mathcal {G} _ {n} \left(\boldsymbol {h} _ {n + 1} ^ {1}\right)\right) / \beta
$$
$$
\boldsymbol {h} _ {n + 1} ^ {2} = \beta \boldsymbol {h} _ {n} ^ {2} + \mathcal {G} _ {n} (\boldsymbol {h} _ {n + 1} ^ {1})
$$
$$
\boldsymbol {h} _ {n} ^ {1} = \left(\boldsymbol {h} _ {n + 1} ^ {1} - \mathcal {F} _ {n} \left(\boldsymbol {h} _ {n} ^ {2}\right)\right) / \lambda
$$
where  $\lambda$  and  $\beta$  are scaling factors. Theoretically,  $\mathcal{F}_n$  and  $\mathcal{G}_n$  could be two arbitrary functions (subnetworks). Given a multilayer reversible network, intermediate activations for each layer during the forward pass are not necessary to be cached. One only needs to store the final outputs, then reconstruct the intermediate activations and calculate the gradient layer-by-layer in a backward manner (See Listing 1 in §Appendix). In this way, the memory footprint required for activations can be reduced significantly and has no relationship with the model's depth, i.e.  $\mathcal{O}(1)$  instead of  $\mathcal{O}(N)$ .
To investigate the training stability of a reversible model, we run experiments on RevViT [40].<sup>4</sup> RevViT shares the same architecture as Reformer [42], except applying a convolutional layer at the beginning to project an image into a sequence of vectors. When running RevViT, one could still cache the intermediate activations and treat it as an irreversible model. We term the gradient calculated in this way as vanilla gradient. One could also train RevViT in a reversible way, and the corresponding
$^3\alpha$  has to be trainable when  $\alpha = 0$ . Otherwise, the newly added parameters are useless.
Our experiments are based on this file, https://github.com/karttikeya/minREV/blob/main/rev.py.
(c)
Figure 4: MEFT architectures. (a) Unfold reversible model. (b) A PLM layer. (c) Two MEFT architectures: (1)  $\mathcal{F}$  is the PLM layer with an adapter (up) and  $\mathcal{G}$  is an adapter (down); (2)  $\mathcal{G}$  is the PLM layer with an adapter (up) and  $\mathcal{F}$  is an adapter (down). (d) The third MEFT architecture:  $\mathcal{G}$  is the MLP block with an adapter (up) and  $\mathcal{F}$  is the attention block with an adapter (down). For initialization,  $W_{down}, W_{up} \sim \mathcal{N}(0, \sigma^2)$  and  $\sigma = 0.02$ . Only the adapter is trainable.
gradient is called reversible gradient. We input the same random noises into the same RevViT twice to obtain the parameter gradients from the convolutional layer, in a vanilla and reversible way. Then we calculate the absolute difference between these two gradients and report the maximum and mean values. In this way, we want to check whether the vanilla gradient can be reconstructed in a reversible way. If the reconstruction error is large, it means that the vanilla gradient could not be recovered in a reversible way due to numerical stability, which might cause unstable training or bad performance.
As shown in Figure 3b, with an increasing number of layers in RevViT, the reconstruction error becomes larger, but still around  $10^{-8}$  which is negligible. However, RevViT is sensitive to the scaling factors,  $\lambda$  and  $\beta$ . When both scaling factors or one of them are less than 1, the reconstruction error increases dramatically. We also explore the initialization of the linear layers in RevViT and find that a larger standard deviation or mean can cause a bigger reconstruction error. In sum, a larger number of layers, smaller scaling factors  $(< 1)$  and a larger standard deviation or mean for initialization tend to cause a bigger reconstruction error, which might result in the unstable training or low performance of a reversible model. Last but not least, RevViT [40] finds that residual connections inside  $\mathcal{F}$  and  $\mathcal{G}$  deteriorate the performance of a reversible Transformer [43].
3 Memory-efficient fine-tuning
This paper aims to modify a PLM to its reversible variant without additional pre-training, so the PLM can still be fine-tuned with a limited memory footprint. The fundamental guiding principle behind our design is: preserving the starting point from a PLM to the greatest extent possible (discus
Table 1: A summarization of three MEFT methods.
MEFT?	F	G	λ	β	Switch	hn1	hn2
1	layer	adapter	→0	any	✓	βhn-1	hn
2	adapter	layer	→1	→0	✓	hn	hn-1
3	attention	MLP	→0	→0	X	-	hn
sion in  $\S 2.2)$  . In this section, we propose three methods to modify a PLM to a reversible one.
3.1  $\mathbf{MEFT}_1$ : PLM layer as  $\mathcal{F}$ , adapter as  $\mathcal{G}$
As shown in Figure 4c, we design  $\mathcal{F}$  as a pre-trained layer with an adapter, where the insertion position for the adapter is borrowed from He et al. [18].  $\mathcal{G}$  is simply an adapter. We initialize the adapters as  $W_{down}, W_{up} \sim \mathcal{N}(0, \sigma^2)$ , same for the following methods. In this way, the output from the adapter is close to  $\mathbf{0}$  at the beginning of the training, so  $h_n \approx \mathcal{F}_n(h_{n-1})$ . For the following discussion, we only focus on the beginning of the training, making sure our design preserves the starting point from a PLM.
$h_0$  and  $h_1$  are the input to and output from the  $1^{st}$  layer of a PLM without any modification, respectively. I.e.  $h_0$  is the representation after the position and word embedding layers of a PLM. We
<sup>5</sup>In RevViT,  $\mathcal{F}$  and  $\mathcal{G}$  are the attention and MLP block (Figure 4b) without residual connection, respectively.
assign  $\pmb{h}_0^1 = \pmb{h}_0^2 = \pmb{h}_0$ , same for the following methods. At the beginning of the training (see Figure 4a),  $\pmb{h}_1^1 = \lambda \pmb{h}_0^1 + \mathcal{F}_1(\pmb{h}_0^2) = \lambda \pmb{h}_0 + \mathcal{F}_1(\pmb{h}_0) \approx \lambda \pmb{h}_0 + \pmb{h}_1$ ,  $\pmb{h}_1^2 = \beta \pmb{h}_0^2 + \mathcal{G}_1(\pmb{h}_1^1) = \beta \pmb{h}_0 + \mathcal{G}_1(\pmb{h}_1^1) \approx \beta \pmb{h}_0$ , where the approximation holds because of our initialization of the adapters.
For now,  $h_1^1$  and  $h_1^2$  are not desired. When we input  $h_1^1$  and  $h_1^2$  to the  $2^{nd}$  reversible layer, especially when we input  $h_1^2$  to  $\mathcal{F}_2$ , the representation continuity is broken, because  $h_1^2 \neq h_1$ . We introduce two modifications to address this issue: (1) We set  $\lambda \to 0$ , so  $h_1^1 \approx h_1$ . (2) Then we switch the order of  $h_1^1$  and  $h_1^2$  before feeding to the next reversible layer, i.e. making  $h_1^1 \approx \beta h_0$  and  $h_1^2 \approx h_1$ . In this way,  $h_1^2$  preserves the starting point. We don't require  $h_1^1$  to preserve any starting point, because it is entered to  $\mathcal{G}_2$  which is not a pre-trained layer.
With the same above-mentioned design for the  $2^{nd}$  reversible layer, we obtain  $\pmb{h}_2^1\approx \beta \pmb{h}_1$  and  $\pmb{h}_2^2\approx \pmb{h}_2$ . By analogy,  $\pmb{h}_n^1\approx \beta \pmb{h}_{n - 1}$  and  $\pmb{h}_n^2\approx \pmb{h}_n$ , which means  $\pmb{h}_n^2$  always preserves the starting point from the PLM. Feeding  $\pmb{h}_n^2$  to the next reversible layer,  $\mathcal{F}_{n + 1}$ , doesn't break the representation continuity. After all layers, we input  $\pmb{h}_N^\prime = (\pmb{h}_N^1 +\pmb{h}_N^2) / 2$  to a task-specific head that is a brand new layer, same for the following methods.
3.2  $\mathbf{MEFT}_2$ : Adapter as  $\mathcal{F}$ , PLM layer as  $\mathcal{G}$
Opposite to  $\mathrm{MEFT}_1$ , we design  $\mathcal{F}$  as an adapter and  $\mathcal{G}$  as the PLM layer with an adapter for  $\mathrm{MEFT}_2$  (see Figure 4c). In this case, we need to make sure that the input to  $\mathcal{G}$  preserves the starting point. Let's also start with the first layer,  $h_1^1 = \lambda h_0^1 + \mathcal{F}_1(h_0^2) = \lambda h_0 + \mathcal{F}_1(h_0) \approx \lambda h_0$ ,  $h_1^2 = \beta h_0^2 + \mathcal{G}_1(h_1^1) = \beta h_0 + \mathcal{G}_1(h_1^1) \approx \beta h_0 + \mathcal{G}_1(\lambda h_0)$ , where the approximation holds because of our initialization of the adapters.
To preserve the starting point from the PLM, we set  $\lambda \rightarrow 1$ ,  $\beta \rightarrow 0$  and switch the order of  $h_1^1$  and  $h_1^2$  before feeding to the next reversible layer. When setting  $\lambda \rightarrow 1$ , we make sure the representation continuity is preserved for  $\mathcal{G}_1$ , resulting in  $h_1^2 \approx \beta h_0 + h_1$ . When  $\beta \rightarrow 0$  and the order of  $h_1^1$  and  $h_1^2$  is switched,  $h_1^1 \approx h_1$  and  $h_1^2 \approx h_0$ . In this way,  $h_1^1$  preserves the initialization point, and we won't break the representation continuity when feeding it to  $\mathcal{G}_2$  in the next reversible layer. With the same setting for each layer,  $h_n^1 \approx h_n$  and  $h_n^2 \approx h_{n-1}$ , so  $h_n^1$  always preserves the starting point.
3.3  $\mathbf{MEFT}_3$ : Attention block as  $\mathcal{F}$ , MLP block as  $\mathcal{G}$