Improving sentence representation for vietnamese natural language understanding using optimal transport

2.1 Multilingual pre-trained language models

Multilingual pre-trained LMs have shown great success in various NLP tasks thanks to their cross-lingual effectiveness. Following the success of Transformer-based architecture [13], several efforts have trained large scale LMs on multiple languages. mBERT is a multilingual variant of BERT [1] that adopts the same training regime by pre-training on 104 languages. Conneau and Lample proposed XLM [6] – a multilingual model pre-trained on a parallel corpus using the translation language modeling objective. XLM-R [7] improves XLM with enormous pre-training data and a better vocabulary; achieving further performance gain on downstream tasks [14, 15]. Recently, mDeBERTa [16] – a multilingual variant of DeBERTa [17], which is pre-trained on the same amount of data as XLM-R. mDeBERTa improves the BERT and RoBERTa models by using disentangled attention and enhanced mask decode which demonstrates its prominent capability on cross-lingual transfer benchmarks. In the meantime, there are many empirical studies that analyze the multilinguality and cross-lingual transfer ability of the multilingual models. They show that a fine-tuned model in one language can be used for another language without relying on any direct cross-lingual supervision [4, 5].

However, many studies point out fundamental limitations of the multilingual models. Conneau et al. [7] observed that scaling the multilingual LMs to a high number of languages only improves the performance to a certain point, after that the performance drops as per language capacity degrades. This is termed as “the curse of multilinguality” or “transfer-interference trade-off”. Also due to this defect, many languages in multilingual models are underrepresented and lag behind their monolingual counterparts in terms of performance on downstream tasks. This can be alleviated by extending per language capacity [7, 18] or further pre-training procedures [19, 20]. Since the model capacity is limited, scaling a bigger model to represent all languages is not practical. To address this, recent works mainly focus on improving performance on one or a set of languages. They have pointed out that monolingual language-specific models pre-trained from scratch achieve better performance than multilingual models, namely Vietnamese [9, 21], Arabic [22], French [23], Finnish [10], Dutch [24], etc. There are some reasons attributed to the increase in performance regarding the multilingual models. For example, monolingual LMs use tailored vocabularies, reduce the need of splitting a word into multiple sub-words, thus avoiding sub-optimal decompositions for other languages [25]. Monolingual LMs often use more pre-training data and this leads to significant improvement in performance, as observed in Li et al. [3].

Recently, the field of natural language understanding has witnessed rapid development, because of the potent ability of LMs in representing textual data. They have a wide range of applications in natural language understanding, including sentiment analysis, question answering, natural language inference, etc. Through the utilization of pre-trained LMs, natural language understanding is becoming more accessible and effective on comprehension and interpretation of human language.

2.2 Knowledge distillation

Knowledge Distillation (KD) is an effective technique that transfers knowledge from a model to another, was first introduced by Hinton et al. [26]. KD aims to train a model, called the student, by exploiting valuable information provided by soft label distribution from another model, called the teacher model. KD has been widely used in a variety of applications in NLP, including model compression [27–29], multi task learning [30, 31], etc. Recent works have investigated methods that align feature spaces between the student and teacher models for a better knowledge transfer. In particular, Wang et al. [32] proposed an application of the KD method that distills the structural knowledge from several monolingual teacher models to the unified multilingual student model in order to solve the problem of the performance gap between monolingual and multilingual model due to capacity limitation. Li et al. [19] focused on learning semantic structure of representation by adopting KD framework to learn rich knowledge from English BERT to improve the multilingual LM. Khanuja et al. [33] enhanced the generalization ability of student models for resource-poor languages by distilling knowledge from multiple multilingual teachers in a task-agnostic setting.

2.3 Optimal transport

Optimal Transport (OT) [34, 35] has become popular in NLP due to its ability to compare two different distributions. It has been successfully applied to many applications. For example, Jianqiao Li et al. [36] proposed using OT to capture positional and contextual information of tokens and tackle the problem of exposure bias in text generation tasks. Peggy Tang et al. [37] formulated text summarization as an OT problem, they considered optimizing transportation cost from an optimal summary to a document based on their semantic distributions. Kyle Swanson et al. [38] employed OT as an objective to align inputs in the text matching problem. Recently, OT has been used for KD to transfer knowledge across models. In this line of work, OT is utilized to compute the optimal value to map between semantic spaces and update the student model afterwards. Thong Nguyen et al. [11] investigated transferring knowledge from a monolingual teacher to a multilingual student for the cross-lingual summarization task. Specifically, they proposed a Knowledge Distillation loss using Sinkhorn Divergence for the transfer process. Tulika Bose et al. [39] introduced a new framework that distills the natural language semantic knowledge from multiple teacher networks to a student network using OT.

In our setup, OT distance is employed as a distillation objective to align the feature spaces of the teacher models and the student model, ensuring that OT enables the positive transfer of knowledge between models.

This work focuses on two Vietnamese NLU tasks, which consists of: machine reading comprehension (MRC) and natural language inference (NLI). The two tasks are formulated as follows:

MRC: the MRC problem can be characterized as a triple 〈C, Q, A〉 where C denotes context, Q denotes question and A denotes answer. A MRC model must be capable of reading and extracting a relevant span of text in the context to answer the question. Table 1 shows an example from UIT-ViQuAD dataset.

OPEN IN VIEWER

Table 2

An example of ViNLI dataset on NLI task (English translation included)

As shown in Figure 1, we present the inputs as a single packed sequence x of length L, where x consists of two segments for the tasks. We then use pre-trained LMs (PLM) to obtain the contextualized embedding H of the input:

H = PLM ( x ) = [ h 1 , . . . h L ](1)

Where H ∈ ℝ L × d ; d is the hidden dimension of the PLM. At the output level, a task-specific layer is utilized by taking the contextualized embedding of the final Transformer layer H for MRC or the representation of the first token in the input sequence [CLS] for NLI. In order to output the predictions, a task-specific layer is followed by an additional Softmax layer. Denoting D as the training data; θ as model parameters, the student model is fine-tuned with negative log-likelihood loss:

L task = - ∑ ( x , y ) ∈ D log p ( y ∣ x , θ )(2)

Let D = { D 1 , . D k } denotes the set of training data on k languages, we have D i is the corpus on language i with m examples: D i = { ( x i j , y i j ) } j = 1 m i · T = { T 1 , . T k } denotes the set of teacher models pre-trained on k languages. Given the set of training data and teacher models, we first fine-tune the monolingual teachers on the corresponding training datasets. The fine-tuning process of teachers will give us the access to exploit the task and language-specific features for the knowledge transfer process. For each training iteration, we select a random language l and sample a batch of dataset B i = { ( x i l , y i l ) } from the corresponding dataset D l . We then feed the batch B to the respective monolingual teacher and multilingual student concurrently to get the hidden semantic representation of the student and the teachers:

H S = PLM student ( x i l )(3) H T = PLM T l ( x i l )(4)

Notice that only the student model parameters θ _S are updated during knowledge distillation. So, we detach H _T from the computational graph.

Previous works [40, 41] have shown the effectiveness of leveraging the hidden representation of the teacher in the knowledge distillation process. However, utilizing the hidden states presents a difficulty in our setting as the teachers and the student were pre-trained on different kinds of data and belong to two feature spaces. This results in the fact that they cannot be projected one-to-one since disjoined vocabularies. In order to make monolingual and multilingual feature space compatible, we employ the Sinkhorn Divergence as the KD objective:

L KD = dist ( H T , H S )(5) Where dist is the OT distance between 2 probability measures. We present this in the next section. Finally, the student jointly learns from the gold targets and the soft targets by minimizing the following objective function:

L = λ L KD + L task(6) Where λ is the hyperparameter that controls the contribution of KD loss and needs to be tuned.

For the purpose of taking the supervision signal from English and Vietnamese monolingual models, we also have to fine-tune the multilingual student model using cross-lingual data. Specifically, at the beginning of the fine-tuning process, we concatenate the original Vietnamese training data with the English training data as the final training set for the student model. Several works have investigated the cross-lingual transfer effect, showing that leveraging one or more source languages can improve the performance of the target language [42–44]. In our setting, the student model is supervised by the training signal from the hard targets of the training data and the soft targets of monolingual models.

Since the representation of the monolingual teacher models and the multilingual student model lie on two different feature spaces, we propose to use OT to measure the distance between them. We also consider this distance as a transportation cost between two probability measures. OT is a powerful method for transferring probability mass from one distribution to another. Formally, given two distributions: the source distribution α and the target distribution β over the domains 𝕋 and 𝕊 , our goal is to optimize θ _s such that the student model distribution matches the teacher model distribution. We use the Sinkhorn Divergence [12] which interpolate between OT and Maximum Mean Discrepancy.

Let α; β are probability distributions that take the form of a sum of Diracs:

α = ∑ i = 1 L T α i δ h i t , α i > 0(7) β = ∑ i = 1 L S β i δ h i s , β i > 0(8)

Note that each α; β must sum to 1: ∑ i = 1 L T α i = 1 ; ∑ i = 1 L S β i = 1 . Let C ∈ ℝ L T × L S be the cost matrix, where c_ij specifies the cost of transferring probability mass from a point in source distribution to a point in target distribution. In our setting, we optimize the student model parameter θ _S so that the transportation cost from H _T to H _S is minimized. Each element of Π ∈ ℝ L T × L S denotes how much probability mass from a point in the source distribution is assigned to a point in the target distribution. Cuturi [45] proposed an entropy regularized approximation of the OT distance that can be solved quickly by adding an entropy penalty term. Using above notations, the entropy-regularized optimal transport problem is formulated as follows:

dist ( H S , H S ) = OT ϵ ( α , β ) - 1 2 OT ϵ ( α , α )(9) - 1 2 OT ϵ ( β , β )

OT ϵ ( α , β ) = min Π 〈 Π , C 〉 + ϵ KL ( Π , α ⊗ β )(10)

s . t . Π ≥ 1 , Π × 1 = α , Π ⊺ × 1 = β(11)

Note that in equation (9), α and β are initialized with uniform distribution; ϵ > 0 controls the amount of entropy regularization to interpolate between OT and Maximum Mean Discrepancy. The OT plan can be efficiently approximated by the Sinkhorn-Knopp algorithm [45, 46]. The Sinkhorn-Knopp algorithm is differentiable, making it ideal for any neural network architectures.

4.1 Datasets

We conduct experiments on two Vietnamese benchmarks datasets for NLU, namely:

Vietnamese Question Answering Dataset (UIT-ViQuAD) [49] contains human-generated question-answer pairs from Vietnamese articles on Wikipedia for Machine Reading Comprehension (MRC). Given a question and a context, The objective of MRC task is to read and extract the answer which is a text span provided in the context.

4.2 Model configurations

We employ transformer-based pre-trained LMs with two versions: base and large. All of the student and teacher models share the same architecture: 12 layers, 8 attention heads; 24 layers, 16 attention heads for base version and large version respectively. For the monolingual teachers, we use RoBERTa [2] on English and PhoBERT [9] on Vietnamese. All of the teachers are already fine-tuned on downstream tasks for learning language-specific features. The student model is initialized from XLM-R [7] and followed by a task-specific layer. During the fine-tuning process, we keep the student model parameters trainable, while freezing the teachers’ parameters. For both tasks, we optimize the student model by the AdamW optimizer and search for the best learning rate in the set {1e - 5, 2e - 5, 3e - 5} and the number of epochs is 2. The batch size is 32. We also employ the learning rate scheduler with a linear decay after 10% of training iterations. We set the entropic regularization parameter ϵ to a relatively small number: 0.005, the value of λ parameter in equation (6) is set to 0.3. All the experiments are performed using two T4 GPUs.

We report the results across evaluation benchmarks. Evaluations are performed on fine-tuned models using the test set, based on their performance on the development set. We name our models XLMR-KDDA. We compare our models with XLM-R, XLM-R with cross-lingual data augmentation (XLMR-DA) and the Vietnamese model PhoBERT on both tasks. Specifically, for the MRC task we also compare our model to two other baselines: XLMRQA [47] and ViReader [48] two open-domain MRC systems for Vietnamese.

To measure the performance of our approach, we adapt the commonly used F1 score and Accuracy for the NLI task. For the MRC task, we measure the F1 and Exact Match (EM) scores. In detail, the F1 score measures the number of overlapping tokens (partial match) between the ground truth and the prediction, while the EM score measures the exact matches. We run all tasks three times with different seeds and report the average scores. More detailed results are presented in the following section.

4.3 Results

ViQuAD In Table 4 we compare the results of our methods with baselines on MRC. It can be observed that our XLMR-KDDA method outperforms other baselines in all evaluation sets, providing a +5.11% and +4.25% improvement in EM over the original multilingual XLM-R fine-tuning method for the base and large versions, respectively. Our gains compared to ViReader and XLMRQA are +2.40% and +2.94% in EM score, and an improvement of +0.51% and +3.19% in F1-score, respectively. The monolingual PhoBERT performs almost the same as the multilingual XLM-R with analogous training data for the base and large versions. We observe that in this setting, all baseline multilingual fine-tuning methods benefit from extra training data from other languages. The performance gains in EM score are up to +2.89% and +3.80% for the base and large versions of XLM-R, respectively. Moreover, we improve EM score from 65.89% to 68.11%, F1 score from 84.60% to 85.37% on the base version with the help of KD. Similarly, on the large version, our proposed KDDA slightly improves the EM and F1 scores by +0.45% and +0.84% respectively, compared to DA. This demonstrates the positive effectiveness of our method.

ViNLI As shown in Table 5, we present our evaluation results on the ViNLI dataset. It is important to note that due to the consistency of the English dataset, we do not take into account the “other” label in the evaluation dataset. Once again, XLMR-KDDA provides better results than other baseline methods. In particular, we obtain improvement on the base version (76.59% /76.59%) and on the large version (85.24% /85.25%) of Accuracy and F1 score, respectively. Compared to the XLM-R _large model, our model shows improvements of +3.89% and +3.94%. Similar to the MRC setting, when we combine KD and DA, the model performs superior to DA in all cases. The performance of PhoBERT is similar to XLM-R; except for the evaluation results of PhoBERT _base on the development set (+3.05% /+3.09% improvement in Accuracy and F1 score).

Comparing our method with other baselines, we can conclude that our method notably improves the results on both ViQuAD and ViNLI datasets. The results show that in this setting, the XLMR-KDDA enables positive transfer in three aspects. First, we observe that the performance benefits from additional training data from another language, particularly the XLM-R _large gets the highest improvement. All the results from the benchmarks suggest that the multilingual model benefits from positive transfer, i.e., training a multilingual model on multiple language datasets yields greater improvements than training it on a monolingual dataset. This demonstrates that cross-lingual transfer may act as regularization strategy, making the model robust to noise and generalize on the training dataset, as proven in [53]. Second, our KD method utilizing OT improves the results and outperforms other baselines in all settings. The monolingual teacher models provide additional training signals, including language-specific and task-specific features, to the multilingual student model, thereby further boosting its performance. In our case, the student and the teacher features are determined by two feature spaces which are Multilingual and Monolingual spaces, while the OT loss function promotes the learning of a student model that minimizes the transportation cost between the feature sets. This provides an efficient approach for the distillation process without the need of additional mapping between two spaces. Third, our findings suggest that the overall performance of the multilingual student model can be enhanced by letting it learn from multiple monolingual teachers simultaneously. These monolingual supervising signals contribute to a diverse and comprehensive understanding of the relationships between different languages from various perspectives, thus helping the generalization process. Our study has led us to the conclusion that by simultaneously modeling the soft label distribution from the monolingual teachers with their corresponding hard label distribution, results in a more flexible and positive transfer. This transfer is learned without making any modifications to the original model architecture, thereby demonstrating the effectiveness of the proposed approach.

It is also worth mentioning that the monolingual models perform worse than multilingual models using the same training data. Except for the evaluation of the PhoBERT-base on ViNLI development dataset, we only observe deterioration in performance with XLM-R. One possible reason for this is that the superior performance of XLM-R is attributed to the enormous amount of Vietnamese pre-training data (7 times bigger than PhoBERT), which aligns with the findings in the previous work [3].

In order to shed more light on the contribution of each component in the OT-based knowledge distillation settings, in this section, we study the performance of the student model under different configurations. We inspect (1) the student’s performance with different teachers, (2) the distillation strategies, and (3) the trade-off between hard targets and soft targets.

Effect of teacher models performance. To investigate the performance of the student model under different teachers, we use only one English teacher and fine-tune seven Vietnamese teachers with random seeds. Subsequently, we conduct the knowledge transfer process using the English teacher and various Vietnamese teachers. As shown in Figure 2 & Figure 3, KDDA yields significantly varied results under different teachers, the EM score is concentrated in range 73.3% to 76.5% and peeks at 76.33% on the ViQuAD dataset. On the ViNLI dataset, we observe a smaller performance gap. The accuracy score peaks at 86.39%. Most of the student models demonstrate an improvement in performance when receiving supervision signals from the teacher models. This clearly highlights that, given fixed data and student model capacity, the performance of the student is sensitive to the quality of the teacher models. We can conclude that stronger teachers result in better student performance.

OPEN IN VIEWER

Fig. 2

Exact Match score on the UIT-ViQuAD dataset for different Vietnamese teacher models.

OPEN IN VIEWER

Fig. 3

Accuracy score on the ViNLI dataset for different teacher.

Effect of using different model output to calculate OT objective Here, we conduct a study on how the performance of the student model is affected by using different distillation scenarios, namely using the hidden states or the logits as an input for calculating the distillation objective. Since previous studies [11, 40, 41] regard teacher’s hidden states as containing supervision signals, we also investigate the effect of utilizing the teacher’s logits. Specifically, we take the unnormalized predictions generated after the task-specific layer and examine how different signals would affect the performance of KDDA. The results on two datasets in Table 6 indicate that using hidden states provides more meaningful representation for the knowledge transfer process compared to using logits. Especially, EM score and Accuracy slightly increase by 2.78% EM and 1.02% accuracy on ViQuAD and ViNLI, respectively. Both distillation scenarios exhibit improvements in performance. This clearly highlights the importance of choosing appropriate KD objectives in the framework.

Effect of lambda parameter Our KD framework depends on one additional parameter λ, which controls the contribution of KD loss in the fine-tuning objective. We present the effect of changing lambda in Table 7 on development set results while fixing other hyperparameters. When the lambda is set to 0.3, we obtain the best performance on the development set. KDDA consistently outperforms other baselines, indicating that the student model benefits from KD on other teacher models.