Efficient transformer with reinforced position embedding for language models

## Abstract

Abstract In this paper, we propose an efficient transformer architecture that uses reinforced positional embedding to obtain superior performance with half the number of encoder decoder layers. We demonstrate that concatenating positional encoding with trainable token embeddings, normalizing columns in the token embedding matrix, and using the normalized token embedding matrix as the value of the attention layer improve the training and validation loss and the training time in an encoder-decoder Transformer model for a Portuguese-English translation task with 10 epochs or 12 hours of training across 10 trials. Our method, with roughly a threefold parameter reduction compared to the baseline model, yields a mean training loss of 1.21, a mean validation loss of 1.51, and an average training time of 1352.27 seconds per epoch, surpassing the baseline model with the same embedding dimension that employs addition of positional encoding and token embeddings, which achieves a mean training loss of 1.96, a validation loss of 2.18, and an average training time of 4297.79 seconds per epoch. Additionally, we evaluated our proposed architecture and the baseline across 14 diverse translation datasets from TensorFlow. The results indicate that our method consistently achieves lower or comparable training and validation losses, suggesting enhanced learning efficiency.

## 1 Introduction

Transformers have become the state-of-the-art in natural language processing (NLP) due to their ability to capture long-range dependencies between input and output sequences using self-attention mechanisms, effectively overcoming the limitations of Recurrent Neural Networks (RNNs) and Neural Machine Translation (NMT) models Fournier et al. (). While RNNs rely on sequential processing, leading to inefficiencies and limited scalability, Transformers eliminate sequential dependencies and instead use self-attention to process sequences in parallel, significantly improving training times and accuracy Chaudhari et al. (). Though models like Long Short-Term Memory (LSTM) and Gated Recurrent Units (GRUs) were developed to overcome some of the limitations of standard RNNs, they still face challenges in processing long sequences effectively, making Transformers a more viable option for tasks requiring longer-term dependencies Khurana et al. (). Despite their advantages, the self-attention mechanism in Transformer leads to quadratic computational complexity with respect to sequence length, limiting their scalability for very long sequences Vyas et al. (). Nevertheless, the parallel processing capabilities and the ability to handle multiple modalities have positioned Transformers as the dominant architecture in NLP tasks like machine translation, text summarization, and automatic speech recognition Zhang & Shafiq ().

Previous work by Ke et al. Ke et al. () demonstrated that employing separate query and key projection matrices for token embeddings and absolute positional embeddings in the self-attention layer accelerates convergence with respect to validation loss in BERT models Lee & Toutanova (). They also indicated that the standard attention mechanism in Transformers might be inefficient due to weak correlations between words and their absolute positions. Huang et al. () introduced a Gated Cumulative Attention Network, which replaces the self-attention mechanism in the Transformer decoder to enhance inference speed. Liu & Abbeel () proposed a Blockwise Parallel Transformer that reduces memory usage. DuSell & Chiang () presented Stack Attention, enabling the Transformer to effectively learn context-free languages.

In this study, we propose modifications to the encoder-decoder Transformer architecture from Vaswani () to enhance training performance on machine translation tasks using the dataset from Ye et al. (). We present three modifications to enhance the efficiency and performance of the encoder-decoder Transformer architecture. Firstly, we reason that the addition of the token and positional embedding matrix may cause loss of information. To address this, we concatenate the token and positional embedding matrices before the initial encoder and decoder blocks, as shown in Figure 1 (b). Second, we normalize each column in the token embedding matrix, as shown in Figure 1 (c). Third, we use the normalized token embedding matrix as the value in the attention layer, as shown in Figure 1 (d).

Figure: Figure 1: (a) The matrix on the left-hand side is a token embedding matrix with the number of rows the same as the number of tokens and with the number of columns the same as the number of token features for the tokens. Each row in the token embedding matrix corresponds to the token feature vector for the token at that row. The matrix at the right-hand side is a positional embedding matrix with the number of rows the same as the number of tokens and with the number of columns the same as the number of features at that position (row). Each row in the positional embedding matrix is the feature vector correspond to that position (row). (b) We normalize each column of the token embedding matrix to make each column having elements with zero mean and unit variance. (c) After the token feature normalization, we concatenate the positional embedding matrix to the right of the normalized token embedding matrix. (d) For the scaled dot-product attention in each attention layer, the value is the normalized token embedding matrix from the input. Created with BioRender.com.
Refer to caption: extracted/5906057/Corrected_Simplified_of_Transformer_Inovations.png

Specifically, let $m$ denote the token embedding dimension and $n$ the sequence length. The input to the first encoder or decoder block consists of the concatenated token and positional embedding matrices, resulting in a matrix of dimension $n\times 2m$. Each column of the token embedding matrix is normalized to have a mean of zero and a standard deviation of one. In the attention layer, the key and query matrices are of dimension $n\times 2m$, while the value matrix is the normalized token embedding matrix, yielding an $n\times m$ dimension. The output of each attention layer is an $n\times 2m$ matrix, which allows for residual connections that maintain positional information across layers. The feedforward layer processes an input matrix of dimension $n\times 2m$ and outputs a matrix of the same dimension. The modified architecture is shown in Figure 2.

Figure: Figure 2: The proposed modified transformer architecture. We made three modification from the transformer architecture in Vaswani (). Firstly, each column in the token embedding matrix is normalized to have zero mean and unit variance for both the encoder and decoder. Secondly, The token embedding matrix and the positional embedding matrix is concatenated before the first encoder block and the first decoder block. Lastly, each attention layer has the value without concatenation. Created with BioRender.com.
Refer to caption: extracted/5906057/Corrected2_Clear_Modified_Transformer_Architecture.png

We compare our proposed model against a baseline architecture similar to that of Vaswani () but with double the number of layers and the same embedding dimension. Namely, in the baseline, the token and positional embedding matrices, both of dimension $n\times 2m$, are summed to form the input matrix for the initial encoder or decoder block.

Using the Portuguese-English translation dataset from Ye et al. () and training both models for 10 epochs, our proposed model achieves a training loss of 1.22 and a validation loss of 1.53, outperforming the baseline model, which recorded a training loss of 1.84 and a validation loss of 2.13. Additionally, our proposed model demonstrates an average training time of 1352 seconds per epoch with 2,809,634 parameters, both of which are roughly three times lower than the baseline model’s averages of 4298 seconds per epoch and 10,184,162 parameters.

These results indicate that our model not only achieves improved training and validation loss but also demonstrates greater training efficiency compared to the baseline Transformer on the Portuguese-English dataset Ye et al. (). To further validate the robustness of our findings across different translation datasets, we applied both models to thirteen additional translation datasets from Ye et al. () training for 10 epochs or 12 hours. Both models successfully completed 10 epochs of training on seven datasets, where our proposed method showed lower training and validation loss in six datasets. However, our method did not yield improved results for the Belarusian-to-English translation dataset, with our model showing a training loss of 4.60 and a validation loss of 4.35, compared to the baseline’s training loss of 4.56 and validation loss of 4.30. Notably, our proposed method achieved nearly a threefold reduction in training time compared to the baseline model.

## 2 Modified Transformer Architecture

Given a sequence $\bm{x}=\begin{bmatrix}x_{1}&x_{2}&\cdots&x_{n}\end{bmatrix}\in\mathbb{R}$ with $n$ elements, the output of a token embedding is a $n\times m$ matrix and it can be written as

$$ $\bm{X}_{e}=\begin{bmatrix}\bm{e}^{(1)}&\bm{e}^{(2)}&\cdots&\bm{e}^{(n)}\end{ bmatrix}^{T}\bm{W}_{1},$ (1) $$

where $\bm{W}_{1}\in\mathbb{R}^{|Vocab|\times m}$ is an embedding matrix, $|Vocab|$ is the number of tokens in the vocabulary set, $m$ is the dimension of an embedding vector, and $\bm{e}^{(i)}\in\mathbb{R}^{m\times 1}$ is a unit column vector with a $1$ at position $i$ and 0 0 elsewhere.

Given a sequence $\bm{y}=\begin{bmatrix}y_{1}&y_{2}&\cdots&y_{n}\end{bmatrix}\in\mathbb{R}^{n}$ with $n$ elements, the output of a token embedding is a $n\times m$ matrix and it can be written as

$$ $\bm{Y}_{e}=\begin{bmatrix}\bm{e}^{(1)}&\bm{e}^{(2)}&\cdots&\bm{e}^{(n)}\end{ bmatrix}^{T}\bm{W}_{2},$ (2) $$

where $\bm{W}_{2}\in\mathbb{R}^{|Vocab|\times m}$ is an embedding matrix, $|Vocab|$ is the number of tokens in the vocabulary set, $m$ is the dimension of an embedding vector, and $\bm{e}^{(i)}\in\mathbb{R}^{m\times 1}$ is a unit column vector with a $1$ at position $i$ and 0 0 elsewhere.

For the positional encoding, given the sequence length $n$ and the embedding dimension $m$, the output is $\bm{P}\in\mathbb{R}^{n\times m}$, where

$$ $\bm{P}_{i,j}=\begin{cases}sin(i/(10000^{2j/m}))&\text{if $1\leq j\leq m/2$}\\ cos(i/(10000^{2j/m}))&\text{if $m/2+1\leq j\leq m$}\\ \end{cases}$ (3) $$

for $1\leq i\leq n$ and $1\leq j\leq m$.

### 2.1 Normalization in token embedding

Given a matrix $\bm{X}\in\mathbb{R}^{n\times m}$, the output of a function for token normalization can be written as

$$ $TN(\bm{X})=\begin{bmatrix}\bm{x}_{1}&\bm{x}_{2}&\cdots&\bm{x}_{m}\end{bmatrix},$ (4) $$

where

$$ $\bm{x}_{j}=\frac{\bm{X}_{:,j}-E[\bm{X}_{:,j}]}{\sqrt{Var[\bm{X}_{:,j}]}},$ (5) $$

$$ $E[\bm{X}_{:,j}]=\frac{1}{n}\sum_{i=1}^{n}\bm{X}_{i,j},$ (6) $$

$$ $Var[\bm{X}_{:,j}]=E[(\bm{X}_{:,j}-E[\bm{X}_{:,j}])^{2}],$ (7) $$

and $\bm{X}_{:,j}$ is the $j$-th column of matrix $\bm{X}$ for $1\leq i\leq n$ and $1\leq j\leq m$.

First, we propose to normalize each column in the token embedding matrix $\bm{X}_{e}$ and $\bm{Y}_{e}$ such that

$$ $\bar{\bm{X}}_{e}=TN(\bm{X}_{e})$ (8) $$

and

$$ $\bar{\bm{Y}}_{e}=TN(\bm{Y}_{e}).$ (9) $$

### 2.2 Concatenation between token embedding and positional embedding

Second, we propose to concatenate the output of the token embedding and the output of the positional encoding such that

$$ $\bm{X}_{I}^{concat}=\begin{bmatrix}\bar{\bm{X}}_{e}&\bm{P}\end{bmatrix}$ (10) $$

and

$$ $\bm{Y}_{I}^{concat}=\begin{bmatrix}\bar{\bm{Y}}_{e}&\bm{P}\end{bmatrix}.$ (11) $$

### 2.3 Avoid concatenation in value in all the attention layers

Third, we propose to use the normalized token embedding matrix in value in each multi-head attention layer.

Given three matrices $\bm{Q},\bm{K}\in\mathbb{R}^{n\times 2m}$ and $\bm{V}\in\mathbb{R}^{n\times m}$, the scaled dot-product attention is defined as Vaswani ()

$$ $Attention(\bm{Q},\bm{K},\bm{V})=softmax(\frac{\bm{Q}\bm{K}^{T}}{\sqrt{m}})\bm{ V},$ (12) $$

where $softmax:\mathbb{R}^{n\times s}\rightarrow[0,1]^{n\times s}$ for $s\in\mathbb{N^{+}}$ is a Softmax function with $softmax(\bm{A})_{i,j}=\frac{\bm{A}_{i,j}}{\sum_{j=1}^{s}e^{\bm{A}_{i,j}}}$ for $\bm{A}\in\mathbb{R}^{n\times s}$.

For the proposed multi-head attention layer without concatenation in value in the $u$-th encoder block, given three matrices $\bm{Q}_{1},\bm{K}_{1}\in\mathbb{R}^{n\times 2m}$ and $\bm{V}_{1}\in\mathbb{R}^{n\times m}$, the output is computed by

$$ $PMHA(\bm{Q}_{1},\bm{K}_{1},\bm{V}_{1},u)=\begin{bmatrix}\bm{H}_{1}&\bm{H}_{2}& \cdots&\bm{H}_{p}\end{bmatrix}\bm{W}^{O_{u}},$ (13) $$

where

$$ $\bm{H}_{k}=Attention(\bm{Q}_{1}\bm{W}^{Q_{1}}_{k},\bm{K}_{1}\bm{W}^{K_{1}}_{k} ,\bm{V}_{1}\bm{W}^{V_{1}}_{k}),$ (14) $$

$u=1,2,\cdots,N$, $k=1,2,\cdots,p$, $N$ is the total number of encoders, $p$ is the number of heads, $\bm{W}^{O_{u}}\in\mathbb{R}^{qp\times 2m}$, $\bm{W}^{Q_{1}}_{k}\in\mathbb{R}^{2m\times r}$, $\bm{W}^{K_{1}}_{k}\in\mathbb{R}^{2m\times r}$, $\bm{W}^{V_{1}}_{k}\in\mathbb{R}^{m\times q}$, $r$ is the dimension of the query and key projection, $q$ is the dimension of the value projection. For the multi-head attention layer in each encoder block and decoder block, $\bm{V}_{1}=\bar{\bm{X}}_{e}$. For the masked multi-head attention layer in each decoder block, $\bm{V}_{1}=\bar{\bm{Y}}_{e}$.

## 3 Experiments and results

### 3.1 Experimental setup

We first evaluated our proposed model and the baseline model using the Portuguese-English translation dataset from TensorFlow Abadi et al. () for 10 epochs or 12 hours of training across 10 trials. The dataset contains 51,785 training, 1,193 validation, and 1,803 test pairs of Portuguese and English text. The encoder’s input is the Portuguese text, and the decoder’s input is the English text, with both sequences capped at 128 tokens. The target labels are the decoder inputs shifted right by one token. We applied wordpiece tokenization for the Portuguese-English translation dataset from TensorFlow Abadi et al. (). For the second experiment, we evaluated our proposed model and the baseline model using the Portuguese-English translation dataset from TensorFlow Abadi et al. () for 10 epochs of training or 12 hours in one trial. We applied wordpiece tokenization for each dataset from TensorFlow Abadi et al. ().

The model was trained using the Adam optimizer from Vaswani ():

$$ $lr(N_{s})=m^{-0.5}min(N_{s}^{-0.5},N_{s}N_{w}^{-1.5}),$ (15) $$

where $m$ is the embedding dimension, $N_{s}\in\mathbb{Z}$ is the current training step, and $N_{w}=4000$ is called the warm-up step. The loss function is a masked cross-entropy loss. The training epoch is $10$. The drop-out rate is 0.1. The batch size is $64$.

For our proposed model (Section 2.2), we used 2 encoder and decoder blocks, the token embedding dimension of $m=64$, $256$ hidden neurons in the feed-forward network, 4 attention heads, and the projection dimension of $64$ for query, key, and value. The proposed model contains 2,809,634 parameters.

For the baseline model, the number of encoder and decoder blocks is $4$; the token embedding dimension is $m=128$; the number of hidden neurons in the feed-forward network is $512$; the number of attention heads is $p=8$; the projection dimension is $128$ for query, key, and value. The baseline model contains 10,184,162 parameters, which is roughly three times more than the proposed model.

### 3.2 Experimental results and discussion

We trained both the baseline and proposed models for 10 epochs or 12 hours across 10 different trials on two A100 40GB GPUs on the Portuguese to English translation dataset from Ye et al. (). For the baseline model, the 10 epochs of training were completed for 4 trails; 9 epochs of training is completed for 4 trails; and 8 epochs of training were completed for 2 trails. For the proposed model, 10 epochs of training were completed for all the trials. The comparison of the learning curves between the baseline and the proposed method are shown in Figure 3 (a). The mean training losses of the baseline are roughly 6.60, 4.55, 3.82, 3.29, 2.89, 2.57, 2.30, 2.11, 1.96, and 1.84. The mean training losses of the model from the proposed method for each epoch are roughly 6.68, 4.55, 3.62, 2.89, 2.36, 1.97, 1.67, 1.46, 1.32, and 1.21, which are less than the mean training losses of the baseline after 3 epochs. The mean validation losses of the baseline are roughly 5.04, 4.06, 3.46, 3.01, 2.77, 2.49, 2.35, 2.25, 2.18, and 2.14. The mean validation losses of the model from the proposed method for each epoch are roughly 5.08, 3.93, 3.10, 2.51, 2.17, 1.88, 1.72, 1.63, 1.56, and 1.51, which are less than the mean validation losses of the baseline after 2 epochs. The comparison from the learning curves shows that the improvement of the training loss and validation loss of our model compared to the baseline is consistent for 10 different trials.

Figure: Figure 3: (a) The transparent blue dashed lines and the transparent blue solid lines are the training loss and the validation loss of the baseline transformer model on the Portuguese to English translation dataset from Ye et al. (), respectively. The transparent brown dashed lines and the transparent brown solid lines are the training loss and the validation loss of the proposed transformer model on the Portuguese to English translation dataset from Ye et al. (), respectively. Each line represent the training or validation loss for one trial. Each model are trained for 10 epochs or 12 hours. The baseline model has 4 trials that finished 10 epochs of training; 4 trails finished 9 epochs of training; 2 trials finished 8 epochs of training. The proposed model finished 10 epochs of training for all the trials. The training losses of the baseline have the mean of roughly 6.60, 4.55, 3.82, 3.29, 2.89, 2.57, 2.30, 2.11, 1.96, and 1.84 and the mean of the validation losses of the baseline are roughly 5.04, 4.06, 3.46, 3.01, 2.77, 2.49, 2.35, 2.25, and 2.18 for 10 different trials. The training losses of the proposed model have the mean of roughly 6.68, 4.55, 3.62, 2.89, 2.36, 1.97, 1.67, 1.46, 1.32, and 1.21 and the mean of the validation losses of the proposed model are roughly 5.08, 3.93, 3.10, 2.51, 2.17, 1.88, 1.72, 1.63, 1.56, and 1.51 for 10 different trials. The proposed model shows a lower mean of training losses after 3 epochs and a lower mean of validation losses after 2 epochs. (b) The bar plot shows the average training time per epoch for both the baseline and the proposed model on the Portuguese to English translation dataset from Ye et al. (). The average training time per epoch for the baseline is roughly 4297.79 seconds, which are higher than roughly 1352.27 seconds for the proposed model. In addition, The variance of the training time for the baseline is roughly 675.79 seconds, which are higher than roughly 144.50 seconds for the proposed model.
Refer to caption: extracted/5906057/Compare_pt2en.png

The average training time for the baseline model and the proposed model for 10 epochs or 12 hours and across 10 different trials is shown in Figure 3 (b). The average training time and the standard deviation of the training time for the proposed model are roughly 1352.27 seconds and 144.50 seconds, respectively, which are less than those in the baseline model with the average training time of 4297.79 seconds and the standard deviation of the training time of 675.79 seconds. The proposed model shows roughly a threefold reduction in training time compared to the baseline. This may be due to the approximately threefold reduction in the number of parameters in the proposed model, which has 2,809,634 parameters, compared to the baseline with 10,184,162 parameters. In summary, the results from Figure 3 show that the proposed model improves the training and validation loss while achieving a threefold reduction in training time.

To verify that the improvement of the training and validation loss and the improvement for the training time are consistent for different translation task, we compared the baseline the the proposed model after training of 10 epochs or 12 hours for one trial on fourteen different translation datasets from Ye et al. (). The number of training data and validation data for each dataset are show in the second and third column in Table 1, respectively. The results from the baseline are shown from the fourth to the seventh column and the results from the proposed model are shown from the eighth to eleventh column in Table 1. The fourth and the eighth columns show the number of epoch the corresponding model has completed in 12 hours. The training and validation loss of the baseline at the epoch indicated in the fourth column are shown in the fifth and the sixth column in Table 1. The training and validation loss of the proposed model at the epoch indicated in the fourth column are shown in the ninth and the tenth column in Table 1. The seventh and the the eleventh columns show the average computational time per epoch in seconds for the baseline and the proposed model in Table 1. The result in Table 1 shows that the proposed model achieved an improved training and validation loss compared to the baseline for 13 different translation datasets except for the Belarusian to English translation dataset (be_to_en). For the Belarusian to English translation dataset, the baseline shows a training loss of $4.56$ and a validation loss of $4.30$, which are lower than those in the proposed model with a training loss of $4.60$ and $4.35$. However, it is important to note that the proposed model outperformed the baseline in the majority of other thirteen datasets, achieving a reduction in training and validation loss. In addition, the training time for the proposed method is roughly a two to four times reduction from the training time of the baseline model for all the fourteen datasets.

**Table 1: Comparison of baseline and proposed transformer architectures**
| Dataset | Data (No.) | Baseline | Proposed method |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Train. | Val. | Ep. | Train.<br>loss | Val.<br>loss | Comp.<br>Time (s) | Ep. | Train.<br>loss | Val.<br>loss | Comp.<br>Time (s) |  |
| az_to_en | 5946 | 671 | 10 | 4.34 | 4.23 | 702 | 10 | 4.32 | 4.22 | 164 |
| aztr_to_en | 188396 | 671 | 2 | 3.59 | 3.47 | 15629 | 6 | 1.69 | 1.79 | 5628 |
| be_to_en | 4509 | 248 | 10 | 4.56 | 4.30 | 573 | 10 | 4.60 | 4.35 | 114 |
| beru_to_en | 212614 | 248 | 1 | 4.89 | 3.88 | 22045 | 6 | 1.64 | 1.68 | 6856 |
| es_to_pt | 44938 | 1016 | 10 | 2.15 | 2.41 | 3341 | 10 | 1.46 | 1.75 | 1157 |
| fr_to_pt | 43873 | 1131 | 10 | 2.33 | 2.80 | 3360 | 10 | 1.62 | 2.05 | 1139 |
| gl_to_en | 10017 | 682 | 10 | 3.77 | 3.75 | 1089 | 10 | 3.69 | 3.67 | 256 |
| glpt_to_en | 61802 | 682 | 9 | 1.92 | 2.48 | 4502 | 10 | 1.17 | 1.80 | 1644 |
| he_to_pt | 48511 | 1145 | 10 | 2.57 | 3.00 | 2970 | 10 | 1.68 | 2.07 | 1263 |
| it_to_pt | 46259 | 1162 | 9 | 2.45 | 2.66 | 4335 | 10 | 1.54 | 1.80 | 1258 |
| pt_to_en | 51785 | 1193 | 10 | 1.84 | 2.13 | 4049 | 10 | 1.22 | 1.53 | 1345 |
| ru_to_en | 208106 | 4805 | 3 | 2.94 | 2.70 | 11055 | 7 | 1.55 | 1.51 | 6122 |
| ru_to_pt | 47278 | 1184 | 8 | 3.22 | 3.43 | 5004 | 10 | 1.87 | 2.23 | 845 |
| tr_to_en | 182450 | 4045 | 2 | 3.60 | 3.03 | 19445 | 6 | 1.66 | 1.57 | 4971 |

## 4 Conclusion

The proposed Transformer architecture, which normalizes the token embedding matrix to have zero mean and unit variance, concatenates token embeddings with positional embeddings before the first encoder and decoder blocks, and uses the normalized token embedding matrix as the value in the attention layer, demonstrates improvements in training loss, validation loss, and computational time across thirteen translation datasets from Ye et al. () compared to the baseline model with the same embedding dimension but with double the number of encoder and decoder blocks. The proposed model achieves a training loss with a mean of roughly 1.21, a validation loss with a mean of roughly 1.51, an average training time of 1352.27 seconds per epoch for 10 epochs of training across 10 trials, which are all lower than those from the baseline which have a training loss with a mean of roughly 1.96, a validation loss with a mean of roughly 2.18, an average training time of 4297.79 seconds per epoch for 10 epochs or 12 hours of training across 10 trials on the Portuguese to English translation dataset from Ye et al. (). The proposed model shows a lower training and validation loss compared to the baseline on thirteen out of fourteen different translation datasets from Ye et al. () and shows around two to four times reduction on the training time on all the fourteen datasets from Ye et al. (). The proposed model shows promising improvements on thirteen out of fourteen datasets, suggesting that its improvements over the baseline may be robust across a range of translation tasks.

#### Acknowledgments

The computational work for this project was conducted using resources provided by the Storrs High-Performance Computing (HPC) cluster. We extend our gratitude to the UConn Storrs HPC and its team for their resources and support, which aided in achieving these results.

## Appendix A Baseline transformer architecture

We introduce the architecture of the baseline Transformer we used for showing the improvement of our proposed model. The baseline model architecture is similar to the Transformer architecture in the ”Neural machine translation with a Transformer and Keras” tutorial from the TensorFlow Abadi et al. () website.

### A.1 Encoder architecture in the baseline transformer

Given a sequence $\bm{x}=\begin{bmatrix}x_{1}&x_{2}&\cdots&x_{n}\end{bmatrix}\in\mathbb{R}$ with $n$ elements, the output of a token embedding is a $n\times m$ matrix and it can be written as

$$ $\bm{X}_{e}=\begin{bmatrix}\bm{e}^{(1)}&\bm{e}^{(2)}&\cdots&\bm{e}^{(n)}\end{ bmatrix}^{T}\bm{W}_{1},$ (16) $$

where $\bm{W}_{1}\in\mathbb{R}^{|Vocab|\times m}$ is an embedding matrix, $|Vocab|$ is the number of tokens in the vocabulary set, $m$ is the dimension of an embedding vector, and $\bm{e}^{(i)}\in\mathbb{R}^{m\times 1}$ is a unit column vector with a $1$ at position $i$ and 0 0 elsewhere.

For the positional encoding, given the sequence length $n$ and the embedding dimension $m$, the output is $\bm{P}\in\mathbb{R}^{n\times m}$, where

$$ $\bm{P}_{i,j}=\begin{cases}sin(i/(10000^{2j/m}))&\text{if $1\leq j\leq m/2$}\\ cos(i/(10000^{2j/m}))&\text{if $m/2+1\leq j\leq m$}\\ \end{cases}$ (17) $$

for $1\leq i\leq n$ and $1\leq j\leq m$.

The input for the first encoder block is

$$ $\bm{X}_{I}=\frac{1}{\sqrt{m}}\bm{X}_{e}+\bm{P}.$ (18) $$

For each encoder block, it contains a multi-head attention layer and a feed forward network layer. For the multi-head attention layer in the $u$-th encoder block, given three matrices $\bm{Q}_{1},\bm{K}_{1},\bm{V}_{1}\in\mathbb{R}^{n\times m}$, the output is computed by

$$ $MHA(\bm{Q}_{1},\bm{K}_{1},\bm{V}_{1},u)=\begin{bmatrix}\bm{H}_{1}&\bm{H}_{2}& \cdots&\bm{H}_{p}\end{bmatrix}\bm{W}^{O_{u}},$ (19) $$

where

$$ $\bm{H}_{k}=softmax(\frac{\bm{Q}_{1}\bm{W}^{Q_{1}}_{k}(\bm{K}_{1}\bm{W}^{K_{1}} _{k})^{T}}{\sqrt{m}})\bm{V}_{1}\bm{W}^{V_{1}}_{k},$ (20) $$

$u=1,2,\cdots,N$, $k=1,2,\cdots,p$, $N$ is the total number of encoders, $p$ is the number of heads, $\bm{W}^{O_{u}}\in\mathbb{R}^{qp\times m}$, $\bm{W}^{Q_{1}}_{k}\in\mathbb{R}^{m\times r}$, $\bm{W}^{K_{1}}_{k}\in\mathbb{R}^{m\times r}$, $\bm{W}^{V_{1}}_{k}\in\mathbb{R}^{m\times q}$, $r$ is the dimension of the query and key projection, $q$ is the dimension of the value projection, and $softmax:\mathbb{R}^{n\times s}\rightarrow[0,1]^{n\times s}$ for $s\in\mathbb{N^{+}}$ is a Softmax function with $softmax(\bm{A})_{i,j}=\frac{\bm{A}_{i,j}}{\sum_{j=1}^{s}e^{\bm{A}_{i,j}}}$ for $\bm{A}\in\mathbb{R}^{n\times s}$.

For the feed forward layer in the $u$-th encoder block, given a matrix $\bm{X}_{f}\in\mathbb{R}^{n\times m}$, the output of the feed forward layer is

$$ $FFN(\bm{X}_{f},u)=ReLU(\bm{X}_{f}\bm{W}_{2_{u}}+\bm{B}_{1_{u}})\bm{W}_{3_{u}}+ \bm{B}_{2_{u}},$ (21) $$

where $ReLU$ is an element-wise rectified linear unit (ReLU) activation function Nair & Hinton () $ReLU(\bm{A}_{i,j})=max(0,\bm{A}_{i,j})$ for $\bm{A}\in\mathbb{R}^{n\times m}$, $\bm{W}_{2_{u}}\in\mathbb{R}^{m\times s}$, $\bm{W}_{3_{u}}\in\mathbb{R}^{s\times m}$, $\bm{B}_{1_{u}}\in\mathbb{R}^{n\times s}$, and $\bm{B}_{2_{u}}\in\mathbb{R}^{n\times m}$.

Given a matrix $\bm{X}\in\mathbb{R}^{n\times m}$, the output of a function for layer normalization Ba () can be written as

$$ $LN(\bm{X})=\begin{bmatrix}\bm{x}_{1}&\bm{x}_{2}&\cdots&\bm{x}_{n}\end{bmatrix} ^{T},$ (22) $$

where

$$ $\bm{x}_{i}=(\frac{\bm{X}_{i,:}-E[\bm{X}_{i,:}]}{\sqrt{Var[\bm{X}_{i,:}]}})^{T},$ (23) $$

$$ $E[\bm{X}_{i,:}]=\frac{1}{m}\sum_{j=1}^{m}\bm{X}_{i,j},$ (24) $$

$$ $Var[\bm{X}_{i,:}]=E[(\bm{X}_{i,:}-E[\bm{X}_{i,:}])^{2}],$ (25) $$

and $\bm{X}_{i,:}$ is the $i$-th row of matrix $\bm{X}$ for $1\leq i\leq n$ and $1\leq j\leq m$.

The output of the first encoder block can be written as

$$ $\bm{X}_{E_{1}}=LN(FFN(\bm{X}_{M_{1}})+\bm{X}_{M_{1}}),$ (26) $$

where $\bm{X}_{M_{1}}=LN(MHA(\bm{X}_{I},\bm{X}_{I},\bm{X}_{I},1)+\bm{X}_{I})$.

The output of the $N$ stacks of encoder blocks can be written as

$$ $\bm{X}_{E_{N}}=LN(FFN(\bm{X}_{M_{N}})+\bm{X}_{M_{N}}),$ (27) $$

where $\bm{X}_{M_{N}}=LN(MHA(K,Q,V,N)+\bm{X}_{E_{N-1}})$ and $K=Q=V=\bm{X}_{E_{N-1}}$ for $N>1$.

### A.2 Decoder architecture in the baseline transformer

Given a sequence $\bm{y}=\begin{bmatrix}y_{1}&y_{2}&\cdots&y_{n}\end{bmatrix}\in\mathbb{R}^{n}$ with $n$ elements, the output of a token embedding is a $n\times m$ matrix and it can be written as

$$ $\bm{Y}_{e}=\begin{bmatrix}\bm{e}^{(1)}&\bm{e}^{(2)}&\cdots&\bm{e}^{(n)}\end{ bmatrix}^{T}\bm{W}_{4},$ $$

where $\bm{W}_{4}\in\mathbb{R}^{|Vocab|\times m}$ is an embedding matrix, $|Vocab|$ is the number of tokens in the vocabulary set, $m$ is the dimension of an embedding vector, and $\bm{e}^{(i)}\in\mathbb{R}^{m\times 1}$ is a unit column vector with a $1$ at position $i$ and 0 0 elsewhere. The input for the first decoder block is

$$ $\bm{Y}_{I}=\frac{1}{\sqrt{m}}\bm{Y}_{e}+\bm{P}.$ (28) $$

For each decoder block, it contains a masked multi-head attention layer, a multi-head cross attention, and a feed forward network layer. For the masked multi-head attention layer in the $u$-th decoder block, given three matrices $\bm{Q}_{2},\bm{K}_{2},\bm{V}_{2}\in\mathbb{R}^{n\times m}$, the output is computed by

$$ $MMHA(\bm{Q}_{2},\bm{K}_{2},\bm{V}_{2},u)=\begin{bmatrix}\bm{H}_{masked_{1}}& \bm{H}_{masked_{2}}&\cdots&\bm{H}_{masked_{p}}\end{bmatrix}\bm{W}^{O_{u}}_{ masked},$ (29) $$

where

$$ $\bm{H}_{masked_{k}}=softmax(\bm{M}+\frac{\bm{Q}_{2}\bm{W}^{Q_{2}}_{k}(\bm{K}_{ 2}\bm{W}^{K_{2}}_{k})^{T}}{\sqrt{m}})\bm{V}_{2}\bm{W}^{V_{2}}_{k}$ (30) $$

$u=1,2,\cdots,N$, $k=1,2,\cdots,p$, $N$ is the total number of decoders, $p$ is the number of heads, $\bm{W}^{O_{u}}_{masked}\in\mathbb{R}^{qp\times m}$, $\bm{W}^{Q_{2}}_{k}\in\mathbb{R}^{m\times r}$, $\bm{W}^{K_{2}}_{k}\in\mathbb{R}^{m\times r}$, $\bm{W}^{V_{2}}_{k}\in\mathbb{R}^{m\times q}$, $\bm{M}\in\mathbb{R}^{n\times n}$ is a matrix with $\bm{M}_{i,j}=\begin{cases}\infty&\text{if $j>i$}\\
0&\text{otherwise}\\
\end{cases}$ for $1\leq i\leq n$ and $1\leq j\leq n$, $r$ is the dimension of the query and key projection, $q$ is the dimension of the value projection, $k=1,2,\cdots,p$, and $softmax:\mathbb{R}^{n\times s}\rightarrow[0,1]^{n\times s}$ for $s\in\mathbb{N^{+}}$ is a Softmax function with $softmax(\bm{A})_{i,j}=\frac{\bm{A}_{i,j}}{\sum_{j=1}^{s}e^{\bm{A}_{i,j}}}$ for $\bm{A}\in\mathbb{R}^{n\times s}$.

For the multi-head cross attention layer in the $u$-th decoder block, given three matrices $\bm{Q}_{3},\bm{K}_{3},\bm{V}_{3}\in\mathbb{R}^{n\times m}$, the output is computed as

$$ $MHCA(\bm{Q}_{3},\bm{K}_{3},\bm{V}_{3},u)=\begin{bmatrix}\bm{H}_{1}&\bm{H}_{2}& \cdots&\bm{H}_{p}\end{bmatrix}\bm{W}^{O_{u}}_{causal},$ (31) $$

where

$$ $\bm{H}_{k}=softmax(\frac{\bm{Q}_{3}\bm{W}^{Q_{3}}_{k}(\bm{K}_{3}\bm{W}^{K_{3}} _{k})^{T}}{\sqrt{m}})\bm{V}_{3}\bm{W}^{V_{3}}_{k},$ (32) $$

$u=1,2,\cdots,N$, $k=1,2,\cdots,p$, $N$ is the total number of encoders, $p$ is the number of heads, $\bm{W}^{O_{u}}_{causal}\in\mathbb{R}^{qp\times m}$, $\bm{W}^{Q_{3}}_{k}\in\mathbb{R}^{m\times r}$, $\bm{W}^{K_{3}}_{k}\in\mathbb{R}^{m\times r}$, $\bm{W}^{V_{3}}_{k}\in\mathbb{R}^{m\times q}$, $r$ is the dimension of the query and key projection, and $q$ is the dimension of the value projection.

The output of the first decoder block can be written as

$$ $\bm{Y}_{D_{1}}=LN(FFN(\bm{Y}_{C_{1}})+\bm{Y}_{C_{1}}),$ (33) $$

where $\bm{Y}_{C_{1}}=LN(MHCA(\bm{Y}_{M_{1}},\bm{X}_{E_{N}},\bm{X}_{E_{N}},1)+\bm{Y}_
{M_{1}})$, $\bm{Y}_{M_{1}}=LN(MMHA(\bm{Y}_{I},\bm{Y}_{I},\bm{Y}_{I},1)+\bm{Y}_{I})$, and $\bm{X}_{E_{N}}$ is the output of the last encoder block defined in (27).

The output of the $N$ stacks of decoder blocks can be written as

$$ $\bm{Y}_{D_{N}}=LN(FFN(\bm{Y}_{C_{N}})+\bm{Y}_{C_{N}}),$ (34) $$

where $\bm{Y}_{C_{N}}=LN(MHCA(\bm{Y}_{M_{N}},\bm{X}_{E_{N}},\bm{X}_{E_{N}},N)+\bm{Y}_
{M_{N}})$, and $\bm{Y}_{M_{N}}=LN(MMHA(\bm{Y}_{D_{N-1}},\bm{Y}_{D_{N-1}},\bm{Y}_{D_{N-1}},N)+
\bm{Y}_{D_{N-1}})$.

## Peer Reviews

confidence: 3
contribution: 2
initial_rating: 3
presentation: 3
questions: More of two detailed comments rather than questions:

- Why is (13) written as an addition inside the softmax? Isn't it easier to formulate it as a binary multiplicative mask outside the softmax?
- The introduction of the architecture could be more concise and focus more on the difference from the baseline. For example, I'm not sure if there's any difference from the baseline that has been introduced in Section 2.4.
soundness: 2
strengths: 1. To the best of my knowledge, there is no research paper formally discussing the comparison between concatenating vs. adding positional embedding to the token embedding. Although there is interest and several informal discussions among the community about it (e.g.: https://www.reddit.com/r/MachineLearning/comments/cttefo/comment/exoz3uy/). Hence, this contribution of this paper may fill some relatively important void in the space.
2. The approach is very straightforward, and the description of the method is mostly easy-to-follow.
summary: This paper studies how to make transformer architecture more efficient by reinforcing the positional embedding. To achieve this goal, the paper makes three proposals: (1) concatenating the token and positional embedding matrices (instead of adding as usual) (2) normalizing the token embedding matrix across tokens (3) use the normalized token embedding matrix as value in the attention layer. Experiments on a few small MT test sets show that the proposed architecture achieves lower training/validation loss with ~30% of the parameter of the baseline model.
weaknesses: 1. My biggest complaints are the evaluation of this paper. Further broken down below:
1.1 The title of this paper states that this is "for language models", yet there is no experiment in a T5-like multi-task language model setup.
1.2 Even the machine translation experiments were not well-done. All experiments are done with toy-scale datasets. Evaluation is limited to perplexity, with no actual task-specific evaluation such as BLEU/COMET.
1.3 I'm not sure why we are not using a comparable hyper-parameter setup between the baseline vs. proposed model. It's impressive to see that you are able to achieve comparable performance with a smaller model, but why not use the same number of encoder/decoder blocks/dimensions etc. and show you can perform better with comparable number of parameters?
2. The proposal should also be benchmarked against more up-to-date baselines such as RoPE (https://arxiv.org/abs/2104.09864).
3. The paper lacks a survey of the existing literature.

---

confidence: 4
contribution: 3
initial_rating: 5
presentation: 3
questions: 1. Why did the authors select the Portuguese-English translation task as the experimental task? Hope to see more mainstream translation tasks, such as English-French, English-German, and English-Chinese.

2. Why is only the loss used as the evaluation metric in the experiment? Although the lower the training loss, the better the model performance, it is still desirable to see a comparison of translation-related evaluation scores, such as BLEU.

3. Why did the authors only use the sine-cosine based positional encodings as the baseline? Hope to give a comparison with more work about the improved positional embedding, such as RoPE, Relative Position Representations (https://arxiv.org/abs/2104.09864, https://arxiv.org/pdf/1803.02155)

4. Why did the authors only choose the translation tasks to validate the proposed method? The baseline model is not very strong, as the proposed method can also be applied to decoder-only architectures. I suggest that the authors verify the effectiveness of the proposed method on some large language models such as LLaMA.

5. The authors use a lot of content to explain the principle of reinforced position embedding, but I don't know why this method can improve the convergence speed of the model ? And why can the model converge to a better solution space than the Transformer baseline model? Adjusting the learning rate or warm-up, or using LayerNorm could also achieve the similar effect, right ?
soundness: 2
strengths: 1. This paper introduces an efficient transformer architecture employing reinforced positional embedding to achieve superior performance with half the number of encoder and decoder layers. Compared to the baseline model, it attains lower training and validation losses while utilizing only a quarter of the training time.

2. The proposed method is mainly aimed at improving the position embedding layer. There are no similar approaches from previous researchers, so the proposed method is novel.
summary: The modified Transformer architecture about position embedding improves training loss, validation loss, and computational time across 13 translation datasets compared to a baseline model. It normalizes token embeddings, concatenates them with positional embeddings, and uses the normalized embeddings in the attention layer. The proposed model achieves lower training and validation losses, with an average training time of 1352.27 seconds per epoch, compared to the baseline's 4297.79 seconds per epoch. These improvements suggest the model's robustness across various translation tasks.
weaknesses: The experimental results are not solid enough:

1. The baseline model is not very strong, as the proposed method can also be applied to decoder-only architectures. The authors should also use a more recent large language model, such as LLaMA, as the baseline model to validate the proposed method.

2. The evaluation score is not very reasonable, it is recommended to use task-related evaluation scores, rather than just using loss as the final evaluation metric

3. Some related works are missing. Hope to give a comparison with more work about the improved positional embedding, such as RoPE, Relative Position Representations (https://arxiv.org/abs/2104.09864, https://arxiv.org/pdf/1803.02155)

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 2
questions: See weakness part
soundness: 2
strengths: Give the author credit to carry out extensive results across 14 diverse translation datasets, and shows promising results over the baseline.
summary: The paper proposes modifications to the standard Transformer architecture by introducing reinforced positional embeddings and alterations in the attention mechanism. 
The author propose to concatenate token embeddings with positional embeddings before the first
encoder and decoder blocks, and uses the normalized token embedding matrix as the value in the
attention layer.
Experimental results demonstrates improvements in training loss, validation loss, and computational time
across thirteen translation datasets
weaknesses: 1.The presentation of section 2.3 and 2.4 is a little confusing. And the notation is not quite consistent, making it hard follow. So I feel the proposed modifications to the Transformer architecture are not entirely clear from the provided content. The paper needs to provide a more thorough explanation of how the reinforced positional embedding works and how it interacts with the Transformer's attention mechanism. Additionally, the mathematical formulation and the pseudo-code or diagrams could be more explicit to help reviewers and readers replicate the study.

2.Lack of related work comparison, I think more related work comparison should be added to show the merits of this work, not just compare to baseline transformer.

3. The English writing could be improved. There are some very long sentences which is hard to read, e.g. the second sentence in abstract, maybe the author could chop it into smaller sentences.

---

confidence: 5
contribution: 2
initial_rating: 3
presentation: 2
questions: 1)	Why is the model scale so small? The baseline is 10M parameters and the proposed method uses 3M parameters, yet it mentions using an A100 40G GPU. For the experiments to be convincing, the model scale should be at least in the billions of parameters - this would better match both the GPU resources mentioned and current research requirements.
2)	While the paper references Ke et al.'s work, it appears to lack comprehensive comparisons with other works on improving Transformer positional encoding. Should experimental comparisons with other related methods be added?
3)	Would this work be useful for more modern positional encodings, such as ROPE  and ALIBI?
soundness: 2
strengths: This work presents two interesting counter-intuitive points:

1)	There has been extensive academic research demonstrating that deep-narrow networks perform better than shallow-wide networks, particularly in the translation field. However, this paper essentially doubles the feature width while reducing the number of layers, yet achieves better results.
2)	The rationale for directly adding positional encodings and token embeddings in vector space has been discussed for many years, and remains the mainstream approach even today. However, this paper suggests that separating these spaces is necessary.
summary: The positional encoding is attached to the token embeddings after batch normalization along the feature dimension. To accommodate the residual connections, most weight matrices in the model are doubled in size (except for W_v). The model was trained on a translation dataset, and its performance was compared with the traditional transformer structure through loss comparison.
weaknesses: 1)	Most of the cited works in this paper are from 2017/2018, and there is a noticeable lack of citations. This raises concerns about whether the authors have a thorough understanding and have conducted comprehensive research in this field. To be honest, this work's issue goes beyond simply needing a few more citations - the paper only has 5 references total, which is far too sparse. While I need to provide some suggestions, I would recommend looking into works in the field of positional encoding, or research analyzing embeddings. I'm afraid I can't offer much more advice to the author beyond that.
2)	The experiments are not solid, with few experimental trials and superficial analysis. It only compares the advantages over the standard transformer model in a very toy setting. I suggest that this work should be combined with more efforts to improve relative/absolute positional encoding, and increase the model's parameter count to enhance soundness. Additionally, conclusions should not be drawn solely from one translation dataset - the effectiveness of this architectural modification should be evaluated on a broader range of tasks including QA, summarization, information extraction, etc.
3)	The structure of the paper is poorly organized. Given that transformer architecture is now very fundamental, making minor modifications only requires brief explanations or could be placed in the appendix - dedicating two and a half pages in the main text is excessive. This leads to a lack of ablation studies, such as analyzing the impact of applying batch normalization to token embeddings and expanding feature dimensions on performance. The paper should include more feasibility analysis from the perspective of vector characteristics. Since there has long been an established understanding that various embeddings can be directly added together, this paper's counter-intuitive approach should dedicate more space to explaining its feasibility. I suggest adding a Related Work section, condensing the lengthy equations and oversized figures in the Methods section, and including an ablation analysis chapter in the Experiments section.
Writing errors or typos:
1)	'Previous work by Ke et al. Ke et al. (2021)' in L33
2)	The order of Figures 1 (b) and 1 (c) has been wrong reversed in L41-42
3)	The model parameter configurations from L300 to L307 were written messily, maybe put it in a table.
4)	'training losses of the baseline have the mean of roughly 6.60, 4.55, 3.82,3.29, 2.89, 2.57, 2.30, 2.11, 1.96, and 1.84 and the mean of the validation losses of the baseline are roughly 5.04, 4.06, 3.46, 3.01, 2.77, 2.49, 2.35, 2.25, and 2.18 for 10 different trials. '  from L315-322,L339-L356. It is not recommended to present such a large amount of numerical results in the main text/figure captions
5)	In Table 1, the three-line table format is not used.