# Auto Correcting in the Process of Translation — Multi-task Learning Improves
Dialogue Machine Translation

Tao Wang1, Chengqi Zhao1, Mingxuan Wang1, Lei Li1, Deyi Xiong2
1ByteDance AI Lab
2College of Intelligence and Computing, Tianjin University, Tianjin, China
{wangtao.960826, zhaochengqi.d<EMAIL_ADDRESS>
<EMAIL_ADDRESS>
<EMAIL_ADDRESS>
###### Abstract

Automatic translation of dialogue texts is a much needed demand in many real
life scenarios. However, the currently existing neural machine translation
delivers unsatisfying results. In this paper, we conduct a deep analysis of a
dialogue corpus and summarize three major issues on dialogue translation,
including pronoun dropping (ProDrop), punctuation dropping (PunDrop), and
typos (DialTypo). In response to these challenges, we propose a joint learning
method to identify omission and typo, and utilize context to translate
dialogue utterances. To properly evaluate the performance, we propose a
manually annotated dataset with 1,931 Chinese-English parallel utterances from
300 dialogues as a benchmark testbed for dialogue translation. Our experiments
show that the proposed method improves translation quality by 3.2 BLEU over
the baselines. It also elevates the recovery rate of omitted pronouns from
26.09% to 47.16%. We will publish the code and dataset publicly at
https://github.com/rgwt123/DialogueMT.

##  1 Introduction

Remarkable progress has been made in Neural Machine Translation (NMT) Bahdanau
et al. (2015); Wu et al. (2016); Lin et al. (2020); Liu et al. (2020) in
recent years, which has been widely applied in everyday life. A typical
scenario for such application is translating dialogue texts, in particular the
record of group chats or movie subtitles, which helps people of different
languages understand cross-language chat and improve their comprehension
capabilities.

(1) |  | Nancy怎么了?
---
${[}$她${]}_{drop}$是不是哭了啊。
MT |  | What happened to Nancy?
---
Did you cry?
REF |  | What happened to Nancy?
---
Did she cry?
(2) | Nancy怎么了${[}$?${]}_{drop}$是不是哭了啊。
MT | Did Nancy cry?
REF | What happened to Nancy? Did she cry?
(3) | Nancy怎么${[}$乐${]}_{typo}$?
MT | How happy is Nancy?
REF | What happened to Nancy?
Table 1: Examples of ProDrop (1), PunDrop (2) and DialTypo (3). MT is
translation results from Google Translate while REF is references.

However, traditional NMT models translate texts in a sentence-by-sentence
manner and focus on the formal text input, such as WMT news translation
Barrault et al. (2020), while the translation of dialogue must take the
meaning of context and the input noise into account. Table 1 shows examples of
dialogue fragment in Chinese and their translation in English. Example (1)
demonstrates that the omission in traditional translation (e.g., dropped
pronouns in Chinese) leads to inaccurate translation results.

Despite its vast potential application, efforts of exploration into dialogue
translation are far from enough. Existing works Wang et al. (2016); Maruf et
al. (2018) focus on either extracting dialogues from parallel corpora, such as
OpenSubtitles Lison et al. (2019), or leveraging speaker information for
integrating dialogue context into neural models. Also, the lack of both
training data and benchmark test set makes establishing current dialogue
translation models far from satisfying and need to be further improved.

In this paper, we try to alleviate the afore-mentioned challenges in dialogue
translation. We first analyze a fraction of a dialogue corpus and summarize
three critical issues in dialogue translation, including ProDrop, PunDrop, and
DialTypo. Then we design a Multi-Task Learning (MTLdial) approach that learns
to self-correct sentences in the process of translating. The model’s encoder
part automatically learns how to de-noise the noise input via explicit
supervisory signals provided by additional contextual labeling. We also
propose three strong baselines for dialogue translation, including repair
(REPAIRdial) and robust (ROBUSTdial) model. To alleviate the challenges
arising from the scarcity of dialogue data, we use sub-documents in the
bilingual parallel corpus to enable the model to learn from cross-sentence
context.

Additionally as for evaluation, the most commonly used BLEU metric Papineni et
al. (2001) for NMT is not good enough to provide a deep look into the
translation quality in such a scenario. Thus, we build a Chinese-English test
set containing sentences with the issues in ProDrop, PunDrop and DialTypo,
attached with the human translation and annotation. Finally, we get a test set
of 300 dialogues with 1,931 parallel sentences. The main contributions of this
paper are as follows:  a) We analyze three challenges ProDrop, PunDrop and
DialTypo, which greatly impact the understanding and translation of a
dialogue.  b) We propose a contextual multi-task learning method to tackle the
analyzed challenges.  c) We create a Chinese-English test set specifically
containing those problems and conduct experiments to evaluate proposed method
on this test set.

##  2 Analysis on Dialogue Translation

Figure 1: Overall approach to NMTdial. The left part of the figure is data
generation, and the right part is three proposed methods. ①/②/③ represents
REPAIRdial, ROBUSTdial and MTLdial separately.

There were already some manual analyses of translation errors, especially in
the field of discourse translation. Voita et al. (2019) study English-Russian
translation and find three main challenges for discourse translation: deixis,
ellipsis, and lexical cohesion. For Chinese-English translation, tense
consistency, connective mismatch, and content-heavy sentences are the most
common issues Li et al. (2014).

Different from previous works, we mainly analyze the specific phenomena in
dialogue translation. We begin with a study on a bilingual dialogue corpus
Wang et al. (2018).111https://github.com/longyuewangdcu/tvsub We translate
source sentences into the target language at sentence level and compare
translation results with reference at dialogue level. Around 1,000 dialogues
are evaluated, and the results are reported in Table 2. From the statistic, we
observe two persistent dialogue translation problems: pronoun dropping
(ProDrop), punctuation dropping(PunDrop). The phenomenon is consistent with
the issue we collect in practical Instant Messaging (IM) chat scenarios,
except for typos since the analyzed dialogue corpus has been proofread to
remove typos.

Types of phenomena | Frequency
---|---
Correct | 88.1%
ProDrop | 4.3%
PunDrop | 3.2%
Incorrect segmentation | 2.4%
Other translation errors | 2.0%
Table 2: Manual evaluation of dialogue samples.

###  2.1 Pronoun Dropping

Pronouns are frequently omitted in pro-drop languages Huang (1989), such as
Chinese, Japanese, Korean, Vietnamese, and Slavic languages. Such phenomenon
are more frequent in dialogue, where the interlocutors are both aware of
what’s omitted in the context. However, when translating a pro-drop language
into a non-pro-drop language (e.g.,
English)222https://en.wikipedia.org/wiki/Pro-drop_language, it is hard to
translate those omitted pronouns, resulting in grammatical errors or semantic
inaccuracies in the target language. The first conversation in Table 1 is an
example.

###  2.2 Punctuation Dropping

In dialogue scenarios, such as IM software, punctuation is often omitted and
the users tend to segment sentences with spaces. The problem becomes much
serious in languages with no spaces, such as Chinese, Japanese, Korean, and
Thai. Table 1 shows this phenomenon in Example (2).

###  2.3 Dialogue Typos

Typos repairing is another fundamental but very challenging practical
problems. In scenarios of dialogue translation, typos or misspellings are very
common, which dramatically undermines the quality of translation output
produced by the machine translation system. Table 1 shows this phenomenon in
Example (3).

##  3 Approach to NMTdial

This section aims to propose a unified framework that facilitates NMT models
to correct the noise input in dialogue neural machine translation (NMTdial).
The framework includes three different methods, which are REPAIRdial,
ROBUSTdial and MTLdial.

###  3.1 Contextual Perturbation Example Generation

The most challenging problem for NMTdial is the data distribution gap between
training and inference stage, where the training data are clean sentence-level
pairs while the test data are noisy dialogue-level conversations.

To bridge the distribution gap, the first step is to generate perturbation
examples based on training instances. The data generation mainly consists of
two steps. The first step is to obtain sub-documents with cross-sentence
context, and the second step is to generate examples with word perturbations
within sub-documents. The left part of Figure 1 shows a complete process.

Cross-sentence Context It is difficult to acquire dialog-level parallel
training data. As an alternative approach, we use parallel document data to
catch the dependency across sentences.

Formally, Let $x_{d}=\\{x^{(1)},x^{(2)},\cdots,x^{(M)}\\}$ be a source-
language document containing $M$ source sentences. And
$y_{d}=\\{y^{(1)},y^{(2)},\cdots,y^{(M)}\\}$ is the corresponding target-
language document containing the same number of sentences as that of the
source document. To get more context information, we random sample consecutive
sub-document pairs $(x_{d},y_{d})$ of $N$ sentences (i.e., snippet pairs from
aligned documents, we set $N\in[1,10]$ in this paper).

We use a special token $<$sep$>$ as the separator to joint the sentence as the
parallel sub-document $\\{(x_{d},y_{d})\\}$.

Contextual Perturbation Then we consider generating perturbation example
$x_{d}^{\prime}$ from $x_{d}$ with respect to sub-document context. For
ProDrop, PunDrop and DialTypo, we build a Chinese pronoun table
$\text{T}_{\texttt{ProDrop}}$, a common punctuation table
$\text{T}_{\texttt{PunDrop}}$ and a Chinese homophone table
$\text{T}_{\texttt{DialTypo}}$ respectively.

With ProDrop and PunDrop, we traverse source sentences of $x_{d}$, discard
pronouns/punctuation in these sentences with a probability of 30% and record
deletion positions with corresponding labels (see details below); to construct
a typo, we choose a word with a probability of 1%, of which 80% is replaced
with one of its homophones according to $\text{T}_{\texttt{DialTypo}}$ and a
probability of 20% is replaced with another random word. We determine these
percentages by observing the generated perturbation data. For annotation
labels, we tag correct words with 0, words of DialTypo with 1, ProDrop words
with 2 and PunDrop words with 3.

Finally we get $x_{d}$, $x_{d}^{\prime}$ and their corresponding label
sequences $\ell_{x}$, $\ell_{x}^{\prime}$. $\ell_{x}$ is a sequence of all 0s.

###  3.2 NMTdial Base Models

With the created training data, we first introduce two methods for NMTdial as
our strong baselines, which will be elaborated here for model comparison.

REPAIRdial A natural way for NMTdial is to train a dialog repair model to
transform dialogue inputs into forms that an ordinary NMT system can deal
with. REPAIRdial involves training a repair model to transform
$x_{d}^{\prime}$ to $x_{d}$ and a clean translation model that translates
$x_{d}$ to $y_{d}$. As a pipeline method, REPAIRdial may suffer from error
propagation.

ROBUSTdial We extend the robust NMT Cheng et al. (2018) to dialogue-level
translation. Specifically, we take both the original $(x_{d},y_{d})$ and the
perturbation $(x_{d}^{\prime},y_{d})$ bilingual pairs as training instances.
So the model is more resilient on dialogue translation. For the inference
stage, the robust model directly translates the raw inputs into the target
language.

###  3.3 MTLdial

ROBUSTdial has the potential to handle translation problems caused by noisy
dialogue input. However, the internal mechanism is rather implicit and in a
black box. Therefore, the improvement is limited, and it is not easy to
analyze the improvement. To address this issue, we introduce a context-aware
multi-task learning method MTLdial for NMTdial.

As shown in ③ of Figure 1, the only difference is that we have a contextual
labeling module based on the encoder. We denote the final layer output of the
Transformer encoder as $H$. For each token $h_{i}$ in
$H=(h_{1},h_{2},...,h_{m})$, the probability of contextual labeling is defined
as:

$P(p_{i}=j|X)=softmax(W\cdot h_{i}+b)[j]$ (1)

where $X=(x_{1},x_{2},...,x_{m})$ is the input sequence, $P(p_{i}=j|X)$ is the
conditional probability that token $x_{i}$ is labeled as $j$ ($j\in{0,1,2,3}$
as defined above).

Here we make the labeling module as simple as possible, so that the
Transformer encoder can behave like BERT Devlin et al. (2019), learning more
information related to perturbation and guiding the decoder to find desirable
translations.

During the training phrase, the model takes
$(x_{d},x_{d}^{\prime},\ell_{x},\ell_{x}^{\prime},y_{d})$ as the training
data. The learning process is driven by optimizing two objectives,
corresponding to sequence labeling as auxiliary loss ($\mathcal{L}_{SL}$) and
machine translation as the primary loss ($\mathcal{L}_{MT}$) in a multi-task
learning framework.

$\displaystyle\mathcal{L}_{SL}=-log(P(\ell_{x}|x_{d})+P(\ell_{x}^{\prime}|x_{d}^{\prime}))$
(2)
$\displaystyle\mathcal{L}_{MT}=-log(P(y_{d}|x_{d})+P(y_{d}|x_{d}^{\prime}))$
(3)

The two objective are linearly combined as the overall objective in learning.

$\mathcal{L}=\mathcal{L}_{MT}+\lambda\cdot\mathcal{L}_{SL}$ (4)

$\lambda$ is coefficient. During experiments, we set as follows according the
best practice:

$\lambda=max(1.0-\frac{update\\_num}{10^{5}},0.2)$ (5)

where $update\\_num$ is updating steps during training.

We introduce multi-task learning for two reasons: 1) The labeling performance
reflects the model’s understanding of sentences containing the mentioned
phenomena. 2) Contextual Labeling can be seen as a pre-training process based
on the BERT-like model, and explicit guidance can enable the encoder to learn
more about the information we annotate.

###  3.4 Modeling Dialogue Context

The modes that we use dialogue context during decoding can be divided into
offline and online. For the offline setting, all sentences in a dialogue are
concatenated one by one with $<$sep$>$. The concatenated sequence is
translated, and the target translation for each sentence can be easily
detected according to the separator $<$sep$>$.

The offline mode can be used for dialogue translation where the entire source
dialogue has already been available before translation (e.g., movie
subtitles). However, we continuously get new source sentences for online chat
and need to generate corresponding translations immediately. We refer to this
mode as the online setting.

We experiment with two online methods. One is online-cut where the current
sentence is concatenated to the previous context with the separator $<$sep$>$.
The trained NMTdial model then translates the concatenated sequence and the
last target segment is used as the translation for the current source
sentence. The other is online-fd. Online-fd is a force decoding method. It
forces the decoder to use translated history and continues decoding instead of
re-translating the entire concatenated sequence. Online-fd brings more
consistent translation.

##  4 Experiments

###  4.1 Test Set

For better evaluation of NMTdial, we create a Chinese-English test set
covering all issues discussed above based on the corpus we analyze in the
second section.

Item | Count
---|---
#dialogues | 300
#sentence pairs | 1,931
#total tokens | 19,155/15,976
#average tokens | 9.92/8.27
#ProDrop | 299
#PunDrop | 542
#DialTypo | 203
Table 3: Statistics on the test set. “/” denote numbers in Chinese and English
separately.

Statistics on the built test set are displayed in Table 3. Building such a
test set is hard and time-consuming as we need to perform manual selection,
translation and annotation.

As for translation quality of the test set, we use other metrics in addition
to overall BLEU. For PunDrop and DialTypo, we evaluate BLEU scores on
sentences containing missing punctuation or typos according to the annotation
information. As for ProDrop, we evaluate the translation quality by the
percentage of correctly recovering and translating the dropped pronouns.

###  4.2 Settings

We adopt the Chinese-English corpus from WMT2020333This corpus includes News
Commentary, Wiki Titles, UN Parallel Corpus, CCMT Corpus, WikiMatrix and Back-
translated news., with about 48M sentence pairs, as our bilingual training
data $D$. We select newstest2019 as the development set. After splicing, we
get $D_{doc}$ with 1.2M pairs and corresponding perturbated dataset
$D^{\prime}$ and $D_{doc}^{\prime}$ with 48M and 1.2M pairs separately.

We use byte pair encoding compression algorithm (BPE) Sennrich et al. (2016)
to process all these data and restricted merge operations to a maximum of 30k
separately. In our studies, all translation models are Transformer-big,
including 6 layers for both encoders and decoders, 1024 dimensions for model,
4096 dimensions for FFN layers and 16 heads for attention.

During training, we use label smoothing = 0.1 Szegedy et al. (2016), attention
dropout = 0.1 and dropout Hinton et al. (2012) with a rate of 0.3 for all
other layers. We use Adam Kingma and Ba (2015) to train the NMT models.
${\beta}1$ and ${\beta}2$ of Adam are set to 0.9 and 0.98, the learning rate
is set to 0.0005, and gradient norm was set to 5. The models are trained with
a batch size of 32,000 tokens on 8 Tesla V100 GPUS during training. During
decoding, we employ beam search algorithm and set the beam size to 5. We use
sacrebleu Post (2018) to calculate uncased BLEU-4 Papineni et al. (2001).

Methods | Overall | Details
---|---|---
BLEU | ProDrop | PunDrop | DialTypo
BASE | 32.7 | 26.09% | 28.2 | 24.0
REPAIRdial | 34.0 | 29.77% | 31.2 | 27.4
ROBUSTdial | 34.1 | 45.48% | 33.0 | 28.8
MTLdial | 35.9 | 47.16% | 34.3 | 28.7
GOLD+BASE | 36.8 | 97.32% | 34.6 | 36.8

Table 4: Experiment results on our constructed dialogue translation test set
in offline setting. The GOLD+BASE represents translations of completely
correct inputs (without ProDrop, PunDrop or DialTypo) using BASE model, which
is used to show the oracle results with Transformer on the test set.

Methods | Overall | Details
---|---|---
BLEU | ProDrop | PunDrop | DialTypo
BASE | 32.8(+0.1) | 19.06%(-7.03%) | 28.1(-0.1) | 22.3(-1.7)
REPAIRdial | 33.8(-0.2) | 24.75%(-5.02%) | 32.0(+0.8) | 28.3(+0.9)
ROBUSTdial | 34.2(+0.1) | 36.79%(-8.69%) | 32.7(-0.3) | 28.9(-0.5)
MTLdial | 35.3(-0.6) | 34.78%(-12.38%) | 34.3(-0.0) | 28.6(-0.1)
GOLD+BASE | 37.1(+0.3) | 96.66%(-0.66%) | 35.3(+0.7) | 35.9(-0.9)

Table 5: Results on our constructed dialogue translation test set in online
setting at the sentence level. Figure 2: Overall BLEU and ProDrop recovery
performance (Accuracy) of MTLdial with different context length. Dash lines
are the offline results.

###  4.3 Results of Offline Setting

The offline mode aims at using the entire source dialogue for translation. We
experiment with all the methods in the offline setting, and the results are
shown in Table 4. BASE is a Transformer-big model trained with $D$ and
$D_{doc}$. GOLD+BASE represents the oracle result on this test set. We can see
that MTLdial has achieved the best results, reducing the gap between
$test_{wrong}$ and $test_{gold}$ from 4.1 to 0.9. Compared with ROBUSTdial and
MTLdial, REPAIRdial performs relatively poorly. We believe that this is due to
the error propagation caused by the pipeline.

From the specific indicators, we can draw the following conclusions: 1)
DialTypo has a very obvious impact on BLEU, and the gap between BASE and
GOLD+BASE is more than 12 points; 2) The recovery of ProDrop is a relatively
difficult task. Although compared with BASE, the current best result of 47.16%
has been greatly improved, but is still far away from 97.32%; 3) PunDrop seems
to be a relatively easy task for each method to address.

###  4.4 Results of Online Setting

The online mode only makes use of the previous context during translation. An
extreme situation of online setting is that there is no context, that is,
sentence-level translation. We show the results of all the methods on the test
set at the sentence level in Table 5. Despite the lack of context, our
approaches can still bring general benefits. We find that ProDrop relies
heavily on context, especially for MTLdial, where the absence of context
results in a 12.38% drop in performance. This is in line with our
expectations, as in many cases machine translation system heavily depends on
context to fulfill the dropped pronouns.

We further experiment on how context lengths can affect NMTdial. The results
are shown in Figure 2. In the online-cut setting, we can see that using the
previous few sentences as the context may improve overall BLEU score, but
continuously adding more preceding texts will lead to a continuous decline.
Online-fd performs well because using historical translation records to
continue decoding can bring more consistent translation results. For the
recovery accuracy of ProDrop, online-cut is better than online-fd in contrast,
because forced decoding may cause wrong pronoun transmission.

##  5 Analysis

Data |  | Precison | Recall | F1
---|---|---|---|---
validation | ProDrop | 61.3 | 48.7 | 54.3
PunDrop | 80.0 | 63.6 | 70.9
DialTypo | 85.3 | 64.2 | 73.2
test | ProDrop | 48.6 | 32.2 | 38.8
PunDrop | 96.6 | 87.9 | 92.1
DialTypo | 83.3 | 31.0 | 45.2
Table 6: Labeling performance on the validation/test set.

###  5.1 Labeling Performance

To better understand how our proposed MTLdial make sense, we calculate the
labeling performance on both validation and test set. Table 6 shows the
overall performance. The validation set follows the same processing progress
of training data, while the test set is the real dialogue data set built
manually.

The proposed model obtains 54.3% F1 score on the validation set for ProDrop,
70.9% for PunDrop, and 73.2% for DialTypo. When testing on the real test data,
the performance on ProDrop has declined a lot because of the difference
between synthetic training/validation data and real test data. Especially
noteworthy is the fact that F1 score of DialTypo drops the most, reaching 26%,
because of its low recall. It may be due to the considerable difference
between the typos generated by our automatic method and the actual
distribution.

| zh | en
---|---|---
(1) |  | 艾丽最近怎么样了
---
她已经不在我的律所了
什么 (她/she)为什么走了
(她/she)开了自己的律所
| What’s going on with Ellie?
---
She is no longer in my law firm.
What, why are you/is she  going?
I open my/She opens her  own law firm.
(2) | 琼斯 (我/I)问你件事 | Jones asked/, I want to ask  you something.
(3) |  | 他上次帮我私下搞
---
(他/He)差点工作都丢了
| He helped me out in private last time.
---
I/He  nearly lost my/his  job.

Table 7: Examples of ProDrop recovery errors. Figure 3: ProDrop recovery
performance of BASE and contextual MTLdial. Total means the total number of
corresponding pronouns in the test set. We ignore pronouns with a total number
less than 5.

###  5.2 Effects of Pronoun Correcting

We further explore the auto-correcting of specific pronouns. As shown in
Figure 3, we can find that pronouns such as I/you, which occur mostly in the
corpus, generally have a higher recovery success rate. We believe this is due
to the data imbalance. Compared with BASE, MTLdial has a much better
performance. While ProDrop recovery accuracy has been multiplied, it still has
not achieved 50%. The most common error is that the model does not capture any
context or captures previous inappropriate context. We summarize frequently-
occurring recovery errors in Table 7.

##  6 Related Work

Our work is related with dialogue translation and robust training.

Dialogue Translation There has been some work on building bilingual dialogue
data sets for the translation task in recent years. Wang et al. (2016) propose
a novel approach to automatically construct parallel discourse corpus for
dialogue machine translation and release around 100K parallel discourse data
with manual speaker and dialogue boundary annotation. Maruf et al. (2018)
propose the task of translating Bilingual Multi-Speaker Conversations. They
introduce datasets extracted from Europarl and Opensubtitles and explore how
to exploit both source and target-side conversation histories. Bawden et al.
(2019) present a new English-French test set for evaluating of Machine
Translation (MT) for informal, written bilingual dialogue. Recently WMT2020
has also proposed a new shared task - machine translation for
chats444http://www.statmt.org/wmt20/chat-task.html, focusing on bilingual
customer support chats Farajian et al. (2020).

Robust Training Since the popularity of deep learning, it has been deeply
affected by noisy issues. Many efforts Li et al. (2017); Sperber et al.
(2017); Vaibhav et al. (2019); Yang et al. (2020) focus on data augmentation
methods and alleviates the problem by adding synthetic noise data to the
training set. However, generating noise has always been a challenge, as
natural noise is always more diversified than artificially constructed noise
Belinkov and Bisk (2018); Anastasopoulos (2019); Anastasopoulos et al. (2019).

##  7 Conclusions

In this paper, we manually analyze challenges in dialogue translation and
detect three main problems. In order to tackle these issues, we propose a
multi-task learning method with contextual labeling. For deep evaluation, we
construct dialogues with translation and detailed annotations as a benchmark
test set. Our proposed model achieves substantial improvements over the
baselines. What is more, we further analyze the performance of contextual
labeling and pronoun recovery errors.

## References

  * Anastasopoulos (2019) Antonios Anastasopoulos. 2019.  An analysis of source-side grammatical errors in nmt.  In _Proceedings of the 2019 ACL Workshop BlackboxNLP: Analyzing and Interpreting Neural Networks for NLP_ , pages 213–223.
  * Anastasopoulos et al. (2019) Antonios Anastasopoulos, Alison Lui, Toan Q Nguyen, and David Chiang. 2019.  Neural machine translation of text from non-native speakers.  In _NAACL-HLT (1)_.
  * Bahdanau et al. (2015) Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. 2015.  Neural machine translation by jointly learning to align and translate.  In _3rd International Conference on Learning Representations, ICLR 2015_.
  * Barrault et al. (2020) Loïc Barrault, Magdalena Biesialska, Ondřej Bojar, Marta R Costa-jussà, Christian Federmann, Yvette Graham, Roman Grundkiewicz, Barry Haddow, Matthias Huck, Eric Joanis, et al. 2020.  Findings of the 2020 conference on machine translation (wmt20).  In _Proceedings of the Fifth Conference on Machine Translation_ , pages 1–55.
  * Bawden et al. (2019) Rachel Bawden, Sophie Rosset, Thomas Lavergne, and Eric Bilinski. 2019.  Diabla: A corpus of bilingual spontaneous written dialogues for machine translation.  _arXiv preprint arXiv:1905.13354_.
  * Belinkov and Bisk (2018) Yonatan Belinkov and Yonatan Bisk. 2018.  Synthetic and natural noise both break neural machine translation.  In _International Conference on Learning Representations_.
  * Cheng et al. (2018) Yong Cheng, Zhaopeng Tu, Fandong Meng, Junjie Zhai, and Yang Liu. 2018.  Towards robust neural machine translation.  In _Proceedings of the 56th Annual Meeting of the Association for Computational Linguistics, ACL 2018, Melbourne, Australia, July 15-20, 2018, Volume 1: Long Papers_ , pages 1756–1766. Association for Computational Linguistics.
  * Devlin et al. (2019) Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. 2019.  Bert: Pre-training of deep bidirectional transformers for language understanding.  In _NAACL-HLT (1)_.
  * Farajian et al. (2020) M Amin Farajian, António V Lopes, André FT Martins, Sameen Maruf, and Gholamreza Haffari. 2020.  Findings of the wmt 2020 shared task on chat translation.  In _Proceedings of the Fifth Conference on Machine Translation_ , pages 65–75.
  * Hinton et al. (2012) Geoffrey E. Hinton, Nitish Srivastava, Alex Krizhevsky, Ilya Sutskever, and Ruslan R. Salakhutdinov. 2012.  Improving neural networks by preventing co-adaptation of feature detectors.  _CoRR_ , abs/1207.0580.
  * Huang (1989) CT James Huang. 1989.  Pro-drop in chinese: A generalized control theory.  In _The null subject parameter_ , pages 185–214. Springer.
  * Kingma and Ba (2015) Diederik P. Kingma and Jimmy Ba. 2015.  Adam: A method for stochastic optimization.  In _Proceedings of ICLR_.
  * Li et al. (2014) Junyi Jessy Li, Marine Carpuat, and Ani Nenkova. 2014.  Assessing the discourse factors that influence the quality of machine translation.  In _Proceedings of the 52nd Annual Meeting of the Association for Computational Linguistics (Volume 2: Short Papers)_ , pages 283–288.
  * Li et al. (2017) Yitong Li, Trevor Cohn, and Timothy Baldwin. 2017.  Robust training under linguistic adversity.  In _Proceedings of the 15th Conference of the European Chapter of the Association for Computational Linguistics: Volume 2, Short Papers_ , pages 21–27.
  * Lin et al. (2020) Zehui Lin, Xiao Pan, Mingxuan Wang, Xipeng Qiu, Jiangtao Feng, Hao Zhou, and Lei Li. 2020.  Pre-training multilingual neural machine translation by leveraging alignment information.  _arXiv preprint arXiv:2010.03142_.
  * Lison et al. (2019) Pierre Lison, Jörg Tiedemann, Milen Kouylekov, et al. 2019.  Open subtitles 2018: Statistical rescoring of sentence alignments in large, noisy parallel corpora.  In _LREC 2018, Eleventh International Conference on Language Resources and Evaluation_. European Language Resources Association (ELRA).
  * Liu et al. (2020) Yinhan Liu, Jiatao Gu, Naman Goyal, Xian Li, Sergey Edunov, Marjan Ghazvininejad, Mike Lewis, and Luke Zettlemoyer. 2020.  Multilingual denoising pre-training for neural machine translation.  _Transactions of the Association for Computational Linguistics_ , 8:726–742.
  * Maruf et al. (2018) Sameen Maruf, André FT Martins, and Gholamreza Haffari. 2018.  Contextual neural model for translating bilingual multi-speaker conversations.  _WMT 2018_ , page 101.
  * Papineni et al. (2001) Kishore Papineni, Salim Roukos, Todd Ward, and Wei-Jing Zhu. 2001.  Bleu: a method for automatic evaluation of machine translation.  In _ACL_.
  * Post (2018) Matt Post. 2018.  A call for clarity in reporting BLEU scores.  In _Proceedings of the Third Conference on Machine Translation: Research Papers_ , pages 186–191, Belgium, Brussels. Association for Computational Linguistics.
  * Sennrich et al. (2016) Rico Sennrich, Barry Haddow, and Alexandra Birch. 2016.  Neural machine translation of rare words with subword units.  In _Proceedings of the 54th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)_ , volume 1, pages 1715–1725.
  * Sperber et al. (2017) Matthias Sperber, Jan Niehues, and Alex Waibel. 2017.  Toward robust neural machine translation for noisy input sequences.  In _International Workshop on Spoken Language Translation (IWSLT)_.
  * Szegedy et al. (2016) Christian Szegedy, Vincent Vanhoucke, Sergey Ioffe, Jon Shlens, and Zbigniew Wojna. 2016.  Rethinking the inception architecture for computer vision.  In _Proceedings of the IEEE conference on computer vision and pattern recognition_ , pages 2818–2826.
  * Vaibhav et al. (2019) Vaibhav Vaibhav, Sumeet Singh, Craig Stewart, and Graham Neubig. 2019.  Improving robustness of machine translation with synthetic noise.  In _Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers)_ , pages 1916–1920.
  * Voita et al. (2019) Elena Voita, Rico Sennrich, and Ivan Titov. 2019.  When a good translation is wrong in context: Context-aware machine translation improves on deixis, ellipsis, and lexical cohesion.  In _Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics_ , pages 1198–1212.
  * Wang et al. (2018) Longyue Wang, Zhaopeng Tu, Shuming Shi, Tong Zhang, Yvette Graham, and Qun Liu. 2018\.  Translating pro-drop languages with reconstruction models.  In _Thirty-Second AAAI Conference on Artificial Intelligence_.
  * Wang et al. (2016) Longyue Wang, Xiaojun Zhang, Zhaopeng Tu, Andy Way, and Qun Liu. 2016.  Automatic construction of discourse corpora for dialogue translation.
  * Wu et al. (2016) Yonghui Wu, Mike Schuster, Zhifeng Chen, Quoc V Le, Mohammad Norouzi, Wolfgang Macherey, Maxim Krikun, Yuan Cao, Qin Gao, Klaus Macherey, et al. 2016.  Google’s neural machine translation system: Bridging the gap between human and machine translation.  _arXiv preprint arXiv:1609.08144_.
  * Yang et al. (2020) Jiacheng Yang, Mingxuan Wang, Hao Zhou, Chengqi Zhao, Weinan Zhang, Yong Yu, and Lei Li. 2020.  Towards making the most of bert in neural machine translation.  In _Proceedings of the AAAI Conference on Artificial Intelligence_ , volume 34, pages 9378–9385.


# SBMLtoODEjax: efficient simulation and optimization of ODE SBML models in
JAX

Mayalen Etcheverry1,2,* Michael Levin3 Clément Moulin-Frier1 Pierre-Yves
Oudeyer1

1 INRIA, University of Bordeaux, Talence 33405, France
2 Poietis, Pessac 33600, France
3 Allen Discovery Center, Tufts University, Medford, MA, USA

* Correspondence: Mayalen Etcheverry<EMAIL_ADDRESS>
Keywords: SBML, Biological Network Analysis, Python, JAX, high performance
computing, parallel computing

## Summary

Developing methods to explore, predict and control the dynamic behavior of
biological systems, from protein pathways to complex cellular processes, is an
essential frontier of research for bioengineering and biomedicine [1]. Thus,
significant effort has gone in computational inference and mathematical
modeling of biological systems [2], [3]. This effort has resulted in the
development of large collections of publicly-available models, typically
stored and exchanged on online platforms (such as the BioModels Database [4],
[5]) using the Systems Biology Markup Language (SBML), a standard format for
representing mathematical models of biological systems [6], [7].

SBMLtoODEjax is a lightweight library that allows to automatically parse and
convert SBML models into python models written end-to-end in JAX, a high-
performance numerical computing library with automatic differentiation
capabilities [8]. SBMLtoODEjax is targeted at researchers that aim to
incorporate SBML-specified ordinary differential equation (ODE) models into
their python projects and machine learning pipelines, in order to perform
efficient numerical simulation and optimization with only a few lines of code.
Taking advantage of JAX’s core transformation features, one can easily boost
the speed of ODE models time-course simulations and perform efficient search
and optimization by running simulations in parallel and/or using automatic
differentiation to find derivatives. SBMLtoODEjax extends SBMLtoODEpy, a
python library developed in 2019 for converting SBML files into python files
written in Numpy/Scipy [9]. The chosen conventions for the generated variables
and modules are slightly different from the standard SBML conventions (used in
the SBMLtoODEpy library) with the aim to accommodate for more flexible
manipulations while preserving JAX-like functional programming style.

SBMLtoODEjax is available at https://github.com/flowersteam/sbmltoodejax.

## Statement of Need

Despite the wealth of available SBML models, scientists still lack an in-depth
understanding of the range of possible behaviors that these models can exhibit
under different initial data and environmental stimuli, and lack effective
ways to search and optimize those behaviors via external interventions. Except
for a subset of simple networks where system behavior and response to stimuli
can be well understood analytically (or with exhaustive enumeration methods),
onerous sampling of the parameter space and time-consuming numerical
simulations are often needed which remains a major roadblock for progress in
biological network analysis.

On the other hand, recent progress in machine learning (ML) has led to the
development of novel computational tools that leverage high-performance
computation, parallel execution and differentiable programming and that
promise to accelerate research across multiple areas of science, including
biological network analysis [10] and applications in drug discovery and
molecular medicine [11], [12]. However, to our knowledge, there is no software
tool that allows seamless integration of existing mathematical models of
cellular molecular pathways (SBML files constructed by biologists) with ML-
supported pipelines and programming frameworks. Whereas there exists many
software tools for manipulation and numerical simulation of SBML models, they
typically rely either on specialized simulation platforms limiting the
flexibility for customization and scripting (such as COPASI [13], [14],
Virtual Cell [15], [16] and Cell Designer [17], [18]) or provide scripting
interfaces in Python or Matlab but rely on backend engines that do not support
hardware acceleration or automatic differentiation (like Tellurium [19], [20]
and SBMLtoODEpy [9] python packages, or the Systems Biology Format Converter
(SBFC) which generates MATLAB and OCTAVE code [21]).

SBMLtoODEjax seeks to bridge that gap by bringing SBML simulation to the JAX
ecosystem, a thriving community of JAX libraries that aim to accelerate
research in machine learning and beyond, with diverse applications spanning
molecular dynamics [22], protein engineering [23], quantum physics [24],
cosmology [25], ocean modeling [26], photovoltaic research [27], acoustic
simulations [28] and fluid dynamics [29]. SBMLtoODEjax aims to integrate this
ecosystem and provide tools to accelerate research in biological network
analysis.

## Why use SBMLtoODEjax?

_Simplicity and extensibility_ SBMLtoODEjax retains the simplicity of the
SBMLtoODEPy library to facilitate incorporation and refactoring of the ODE
models into one’s own python projects. As shown in Figure 1, with only a few
lines of python code one can load and simulate existing SBML files.

Figure 1: Example code (left) and output snapshot (right) reproducing original
simulation results of Kholodenko 2000’s paper [30] hosted on BioModels
website.

_JAX-friendly_ The generated python models are tailored to take advantage of
JAX main features. Model rollouts use jit transformation and scan primitive to
reduce compilation and execution time of the recursive ODE integration steps,
which is particularly useful when running large numbers of steps (long
reaction times). Models also inherit from the Equinox module abstraction [31]
and are registered as PyTree containers, which facilitates the application of
JAX core transformations to any SBMLtoODEjax object.

_Efficiency simulation and optimization_ The application of JAX core
transformations, such as just-in-time compilation (jit), automatic
vectorization (vmap) and automatic differentiation (grad), to the generated
models make it very easy (and seamless) to efficiently run simulations in
parallel. For instance, as shown in Figure 2, with only a few lines of python
code one can vectorize calls to model rollout and perform batched computations
efficiently, which is particularly useful when considering large batch sizes.

Figure 2: (left) Example code to vectorize calls to model rollout (right)
Results of a (rudimentary) benchmark comparing the average simulation time of
models implemented with SBMLtoODEpy versus SBMLtoODEjax (for different number
of rollouts i.e. batch size). For additional details on the comparison, please
refer to our Benchmarking notebook.

As shown in Figure 3, SBMLtoODEjax models can also be integrated within Optax
pipelines, a gradient processing and optimization library for JAX [32],
allowing to optimize model parameters and/or external interventions with
stochastic gradient descent.

Figure 3: (left) Default simulation results of biomodel #145 which models ATP-
induced intracellular calcium oscillations, and (arbitrary) target sine-wave
pattern for Ca_Cyt concentration. (middle) Training loss obtained when running
the Optax optimization loop, with Adam optimizer, over the model kinematic
parameters $c$. (right) Simulation results obtained after optimization. The
full example is available at our Gradient Descent tutorial.

Altogether, the parallel execution capabilities and the differentiability of
the generated models opens interesting possibilities to design and optimize
intervention strategies.

_Current limitations_ SBMLtoODEjax is still in its early phase and does not
yet handle all possible cases of SBML files, but we welcome contributions and
aim to handle more cases in future releases. Similarly, SBMLtoODEjax only
integrates one ODE solver for now (jax.experimental.odeint), but could benefit
from more [33]. Finally, whereas SBMLtoODEjax is to our knowledge the first
software tool enabling gradient backpropagation through the SBML model
rollout, applying it in practice can be hard and other optimization methods,
such as evolutionary strategies, might be more adapted.

_Documentation_ Please refer to https://developmentalsystems.org/sbmltoodejax/
for additional details on SBMLtoODEjax’s main design principles, advantages
and limitations, and API docs as well as for various hands-on tutorials for
loading and simulating biomodels, parallel execution and gradient descent.

## Software requirements

SBMLtoODEjax is developed under the MIT license and available on PyPI via pip
install sbmltoodejax. It is written on top of SBMLtoODEpy [9], JAX (cpu) [8]
and Equinox [31], which are the main requirements.

## Acknowledgements

SBMLtoODEjax builds on SBMLtoODEpy’s parsing and conversion of SBML files [9],
JAX’s composable transformations [8], Equinox’s module abstraction [31] and
BasiCO’s access to the BioModels REST api [14].

This work has been funded by the biotechnology company Poietis and the French
National Association of Research and Technology (ANRT), as well as by the
French National Research Agency (ANR, project DeepCuriosity). It also
benefited from two mobility research scholarships given by the French Academy
(Jean Walter Zellidja scholarship) and the University of Bordeaux (UBGRS-Mob
scholarship). Finally, this work benefited from the use of the Jean Zay
supercomputer associated with the Genci grant A0091011996.

## References

re [1] H. Kitano, “Systems Biology: A Brief Overview,” _Science_ , vol. 295,
no. 5560, pp. 1662–1664, Mar. 2002, doi: 10.1126/science.1069492.

pre [2] H. de Jong, “Modeling and Simulation of Genetic Regulatory Systems: A
Literature Review,” _Journal of Computational Biology_ , vol. 9, no. 1, pp.
67–103, Jan. 2002, doi: 10.1089/10665270252833208.

pre [3] F. M. Delgado and F. Gómez-Vela, “Computational methods for Gene
Regulatory Networks reconstruction and analysis: A review,” _Artificial
Intelligence in Medicine_ , vol. 95, pp. 133–145, Apr. 2019, doi:
10.1016/j.artmed.2018.10.006.

pre [4] M. Glont _et al._ , “BioModels: Expanding horizons to include more
modelling approaches and formats,” _Nucleic Acids Research_ , vol. 46, no. D1,
pp. D1248–D1253, Jan. 2018, doi: 10.1093/nar/gkx1023.

pre [5] R. S. Malik-Sheriff _et al._ , “BioModels—15 years of sharing
computational models in life science,” _Nucleic Acids Research_ , vol. 48, no.
D1, pp. D407–D415, Jan. 2020, doi: 10.1093/nar/gkz1055.

pre [6] M. Hucka _et al._ , “The systems biology markup language (SBML): A
medium for representation and exchange of biochemical network models,”
_Bioinformatics_ , vol. 19, no. 4, pp. 524–531, Mar. 2003, doi:
10.1093/bioinformatics/btg015.

pre [7] M. Hucka _et al._ , “The Systems Biology Markup Language (SBML):
Language Specification for Level 3 Version 2 Core Release 2,” _Journal of
Integrative Bioinformatics_ , vol. 16, no. 2, p. 20190021, Jun. 2019, doi:
10.1515/jib-2019-0021.

pre [8] J. Bradbury _et al._ , “JAX: Composable transformations of
Python+NumPy programs.” 2018.

pre [9] S. M. Ruggiero and A. N. F. Versypt, “SBMLtoODEpy: A software program
for converting SBML models into ODE models in Python,” _Journal of Open Source
Software_ , vol. 5, no. 53, p. 1643, Sep. 2019, doi: 10.21105/joss.01643.

pre [10] G. Muzio, L. O’Bray, and K. Borgwardt, “Biological network analysis
with deep learning,” _Briefings in Bioinformatics_ , vol. 22, no. 2, pp.
1515–1530, Mar. 2021, doi: 10.1093/bib/bbaa257.

pre [11] D. M. Camacho, K. M. Collins, R. K. Powers, J. C. Costello, and J. J.
Collins, “Next-Generation Machine Learning for Biological Networks,” _Cell_ ,
vol. 173, no. 7, pp. 1581–1592, Jun. 2018, doi: 10.1016/j.cell.2018.05.015.

pre [12] M. AlQuraishi and P. K. Sorger, “Differentiable biology: Using deep
learning for biophysics-based and data-driven modeling of molecular
mechanisms,” _Nature Methods_ , vol. 18, no. 10, pp. 1169–1180, Oct. 2021,
doi: 10.1038/s41592-021-01283-4.

pre [13] S. Hoops _et al._ , “COPASI—a COmplex PAthway SImulator,”
_Bioinformatics_ , vol. 22, no. 24, pp. 3067–3074, Dec. 2006, doi:
10.1093/bioinformatics/btl485.

pre [14] F. T. Bergmann, “Copasi/basico: Release 0.48.” Mar. 2023.

pre [15] L. M. Loew and J. C. Schaff, “The Virtual Cell: A software
environment for computational cell biology,” _Trends in Biotechnology_ , vol.
19, no. 10, pp. 401–406, Oct. 2001, doi: 10.1016/S0167-7799(01)01740-1.

pre [16] B. M. Slepchenko, J. C. Schaff, I. Macara, and L. M. Loew,
“Quantitative cell biology with the Virtual Cell,” _Trends in Cell Biology_ ,
vol. 13, no. 11, pp. 570–576, Nov. 2003, doi: 10.1016/j.tcb.2003.09.002.

pre [17] A. Funahashi, M. Morohashi, H. Kitano, and N. Tanimura,
“CellDesigner: A process diagram editor for gene-regulatory and biochemical
networks,” _BIOSILICO_ , vol. 1, no. 5, pp. 159–162, Nov. 2003, doi:
10.1016/S1478-5382(03)02370-9.

pre [18] A. Funahashi, Y. Matsuoka, A. Jouraku, M. Morohashi, N. Kikuchi, and
H. Kitano, “CellDesigner 3.5: A Versatile Modeling Tool for Biochemical
Networks,” _Proceedings of the IEEE_ , vol. 96, no. 8, pp. 1254–1265, Aug.
2008, doi: 10.1109/JPROC.2008.925458.

pre [19] K. Choi _et al._ , “Tellurium: An extensible python-based modeling
environment for systems and synthetic biology,” _Biosystems_ , vol. 171, pp.
74–79, Sep. 2018, doi: 10.1016/j.biosystems.2018.07.006.

pre [20] J. K. Medley _et al._ , “Tellurium notebooks—An environment for
reproducible dynamical modeling in systems biology,” _PLOS Computational
Biology_ , vol. 14, no. 6, p. e1006220, Jun. 2018, doi:
10.1371/journal.pcbi.1006220.

pre [21] N. Rodriguez _et al._ , “The systems biology format converter,” _BMC
Bioinformatics_ , vol. 17, no. 1, p. 154, Apr. 2016, doi:
10.1186/s12859-016-1000-2.

pre [22] S. S. Schoenholz and E. D. Cubuk, “JAX, M.D.: A Framework for
Differentiable Physics.” arXiv, Dec. 2020. Accessed: Jun. 02, 2023. [Online].
Available: https://arxiv.org/abs/1912.04232

pre [23] E. J. Ma and A. Kummer, “Reimplementing Unirep in JAX.” bioRxiv, p.
2020.05.11.088344, May 2020. doi: 10.1101/2020.05.11.088344.

pre [24] G. Carleo _et al._ , “NetKet: A machine learning toolkit for many-
body quantum systems,” _SoftwareX_ , vol. 10, p. 100311, Jul. 2019, doi:
10.1016/j.softx.2019.100311.

pre [25] J.-E. Campagne _et al._ , “JAX-COSMO: An End-to-End Differentiable
and GPU Accelerated Cosmology Library,” _The Open Journal of Astrophysics_ ,
vol. 6, p. 10.21105/astro.2302.05163, Apr. 2023, doi:
10.21105/astro.2302.05163.

pre [26] D. Häfner, R. Nuterman, and M. Jochum, “Fast, Cheap, and
Turbulent—Global Ocean Modeling With GPU Acceleration in Python,” _Journal of
Advances in Modeling Earth Systems_ , vol. 13, no. 12, p. e2021MS002717, 2021,
doi: 10.1029/2021MS002717.

pre [27] S. Mann, E. Fadel, S. S. Schoenholz, E. D. Cubuk, S. G. Johnson, and
G. Romano, “$\partial$PV: An end-to-end differentiable solar-cell simulator,”
_Computer Physics Communications_ , vol. 272, p. 108232, Mar. 2022, doi:
10.1016/j.cpc.2021.108232.

pre [28] A. Stanziola, S. R. Arridge, B. T. Cox, and B. E. Treeby, “J-Wave: An
open-source differentiable wave simulator,” _SoftwareX_ , vol. 22, p. 101338,
May 2023, doi: 10.1016/j.softx.2023.101338.

pre [29] D. A. Bezgin, A. B. Buhendwa, and N. A. Adams, “JAX-Fluids: A fully-
differentiable high-order computational fluid dynamics solver for compressible
two-phase flows,” _Computer Physics Communications_ , vol. 282, p. 108527,
Jan. 2023, doi: 10.1016/j.cpc.2022.108527.

pre [30] B. N. Kholodenko, “Negative feedback and ultrasensitivity can bring
about oscillations in the mitogen-activated protein kinase cascades,”
_European Journal of Biochemistry_ , vol. 267, no. 6, pp. 1583–1588, Mar.
2000, doi: 10.1046/j.1432-1327.2000.01197.x.

pre [31] P. Kidger and C. Garcia, “Equinox: Neural networks in JAX via
callable PyTrees and filtered transformations,” _Differentiable Programming
workshop at Neural Information Processing Systems 2021_ , 2021.

pre [32] I. Babuschkin _et al._ , “The DeepMind JAX Ecosystem.” 2020.

pre [33] P. Städter, Y. Schälte, L. Schmiester, J. Hasenauer, and P. L.
Stapor, “Benchmarking of numerical integration methods for ODE models of
biological systems,” _Scientific Reports_ , vol. 11, no. 1, p. 2696, Jan.
2021, doi: 10.1038/s41598-021-82196-2.

p