PaperlayerView exact revision on arXiv

Source-linked AI summary

Improving Non-autoregressive Translation Quality with Pretrained Language Model, Embedding Distillation and Upsampling Strategy for CTC

Shen-sian Syu, Juncheng Xie, Hung-yi Lee

arXiv:2306.06345v3cs.CL

TL;DR

NAT can generate translations in parallel, but its quality can lag autoregressive models. This paper combines PMLM fine-tuning with CTC, upsampling, and embedding distillation; it reports stronger benchmark performance than Transformer base alongside a 16.35-times speed improvement.

  • Problem

    NAT’s parallel generation can reduce inference latency, but its quality is lower than autoregressive models and its independence assumption poses translation-quality challenges.

    Source: B. Non-Autoregressive NMT · p. 2

  • Method

    CTCPMLM fine-tunes PMLMs with CTC for one-pass NAT, using MASK insertion and adaptive upsampling alongside sequence-level and embedding distillation.

    Source: §I · pp. 1–2

  • Results

    CTCPMLM outperforms Transformer base on WMT’14 DE↔EN, WMT’16 RO↔EN, and IWSLT’14 DE↔EN, with a 16.35-times speed improvement over the autoregressive model.

    Source: §I · p. 2

  • Takeaways & Limitations

    The authors report state-of-the-art performance among NAT models and results above existing NAT models on WMT’14 En↔De and WMT’16 En↔Ro.

    Source: §I · p. 2, §VII · pp. 11–12

Abstract

from arXiv · show

Non-autoregressive approaches aim to improve the inference speed of translation models, particularly those that generate output in a one-pass forward manner. However, these approaches often suffer from a significant drop in translation quality compared to autoregressive models. This paper introduces a series of innovative techniques to enhance the translation quality of Non-Autoregressive Translation (NAT) models while maintaining a substantial acceleration in inference speed. We propose fine-tuning Pretrained Multilingual Language Models (PMLMs) with the CTC loss to train NAT models effectively. Furthermore, we adopt the MASK insertion scheme for up-sampling instead of token duplication, and we present an embedding distillation method to further enhance performance. In our experiments, our model outperforms the baseline autoregressive model (Transformer \textit{base}) on multiple datasets, including WMT'14 DE$\leftrightarrow$EN, WMT'16 RO$\leftrightarrow$EN, and IWSLT'14 DE$\leftrightarrow$EN. Notably, our model achieves better performance than the baseline autoregressive model on the IWSLT'14 En$\leftrightarrow$De and WMT'16 En$\leftrightarrow$Ro datasets, even without using distillation data during training. It is worth highlighting that on the IWSLT'14 DE$\rightarrow$EN dataset, our model achieves an impressive BLEU score of 39.59, setting a new state-of-the-art performance. Additionally, our model exhibits a remarkable speed improvement of 16.35 times compared to the autoregressive model.

I. INTRODUCTION

NAT enables parallel translation but faces quality and target-length challenges. The introduction presents CTCPMLM as a one-pass approach combining pretrained multilingual models, CTC, upsampling, and distillation, with reported gains over Transformer base.

I. INTRODUCTION

  • NAT’s independent target-token predictions can produce multimodal outputs and repetition, while sequence-level distillation replaces training targets with an autoregressive model’s outputs.The introduction also notes that iterative architectures use previously predicted tokens to add dependence.Source: §I · p. 1
  • The proposed CTCPMLM is a one-pass NAT approach combining encoder-based PMLM initialization, CTC loss, sequence-level distillation from a strong AT teacher, upsampling, and embedding distillation.The authors describe the goal as integrating these techniques to improve NAT performance.Source: §I · p. 1
  • Unlike prior approaches that add modules to PMLMs, CTCPMLM fine-tunes PMLMs directly with CTC without adding parameters.The authors characterize this as efficient and effective in their experiments.Source: §I · p. 2
  • For CTC upsampling, the paper proposes inserting MASK tokens and adapting the upsampling rate, and reports searching among schemes to improve accuracy.The introduction identifies target sequence length as a major challenge.Source: §I · pp. 1–2
  • Embedding distillation uses a frozen PMLM to regularize contextualized embeddings, aiming to retain target-language information learned during pretraining.The paper presents this as a way to preserve PMLM knowledge in the NAT model.Source: §I · p. 2
  • CTCPMLM outperforms Transformer base across WMT’14 DE↔EN, WMT’16 RO↔EN, and IWSLT’14 DE↔EN, with a 16.35-times speed improvement; IWSLT’14 DE→EN reaches 39.93 BLEU.The authors also report performance above baseline autoregressive models on IWSLT’14 En↔De and WMT’16 En↔Ro with raw, undistilled training data.Source: §I · p. 2

II. BACKGROUND

The background contrasts autoregressive generation, which conditions on previously generated tokens, with NAT’s parallel, conditionally independent predictions. It then describes CTC as an approach to target-length handling that requires inputs at least as long as outputs.

A. Autoregressive Neural Machine Translation

  • Autoregressive NMT factors target-sequence probability left to right, conditioning each token on earlier generated tokens and the source; this makes parallelization difficult and increases latency.The model parameterizes each token’s probability using the previous target tokens, source, and parameters.Source: A. Autoregressive Neural Machine… · p. 2

B. Non-Autoregressive NMT

  • NAT emits target tokens simultaneously, but its independence assumption harms translation quality and is associated with multimodality and inflexible prefixed output length.The background identifies multimodality and output-length inflexibility as significant NAT limitations.Source: B. Non-Autoregressive NMT · p. 2

C. Connectionist Temporal Classification (CTC)

  • CTC handles target-length prediction through alignments with repeated and blank tokens, but requires the source sequence to be at least as long as the target.The alignment is post-processed by removing blanks and collapsing consecutive repeated tokens.Source: C. Connectionist Temporal… · p. 2

D. Upsampling Method

CTCPMLM combines pretrained language-model initialization, CTC training, and upsampling strategies to support one-pass NAT translation. Its approach uses insertion-based upsampling and dynamically adjusted ratios to address CTC length requirements.

III. APPROACH

  • The framework maps source x to upsampled source x̂ and prediction ŷ, while a frozen PMLM provides target-side representations for embedding alignment.The figure labels the source target y and the NAT and teacher hidden states used in training.Source: §III · p. 3

A. PMLM Initialization

  • CTCPMLM fine-tunes an encoder-only pretrained language model with CTC loss, using its word embedding, self-attention, and output projection layers.The framework also uses a frozen PMLM teacher to align the NAT model's embedding space.Source: A. PMLM Initialization · p. 3, §III · p. 3

B. Upsampling Strategy for CTC

  • CTC requires input sequences longer than the output, so the method offers fixed and dynamic upsampling ratios.The dynamic ratio adjusts to source length to keep long upsampled inputs below the position limit; fixed-ratio positions are clipped at the model maximum.Source: B. Upsampling Strategy for CTC · p. 3
  • Insertion upsampling expands the source by duplicating tokens (IT) or inserting [MASK] tokens (IM), rather than using projection or SoftCopy.With s = 2, the example maps [A, B, C, D] to [A, M, B, M, C, M, D, M] under IM.Source: B. Upsampling Strategy for CTC · p. 4

C. Embedding Distillation (ED)

Embedding distillation aligns NAT representations with those of a frozen PMLM teacher. Because the representations may differ in length and token order, the method first finds an alignment for comparing them.

C. Embedding Distillation (ED)

  • The NAT model is trained to maximize cosine similarity with a frozen copy of the PMLM, approximating the teacher's target-language contextualized representations.The NAT representation comes from the source sequence, while the teacher representation comes from the target sequence.Source: C. Embedding Distillation (ED) · p. 4
  • Because NAT and teacher representations can differ in length and token order, the method aligns predicted and target sequences before comparing their representations.It forms Q from the probability of each target token at each predicted position, then uses the Hungarian algorithm to find optimal matching pairs.Source: C. Embedding Distillation (ED) · p. 4
  • The resulting alignment is used to calculate the similarity between NAT and frozen-PMLM representations.Source: C. Embedding Distillation (ED) · p. 4

D. Final Objective

The final training objective combines the proposed objectives at every training step, with λ set to 0 or 1 in the experiments. The setup evaluates the models on three translation corpora and specifies their training and evaluation procedures.

D. Final Objective

  • The model jointly trains with the proposed objectives at each step, with λ set to 0 or 1 in the experiments.For IWSLT’14 and WMT’14, λ changes from 0 to 1 during training; WMT’16 uses the same staged schedule with different step counts.Source: D. Final Objective · p. 4, A. Experiment Setup · p. 5

A. Experiment Setup

  • Experiments use IWSLT’14 En↔De, WMT’16 En↔Ro, and WMT’14 En↔De; IWSLT’14 supports hyperparameter search and ablations.WMT’16 uses newsdev-2016 for validation and newstest-2016 for testing; WMT’14 uses newstest2012 and newstest2013 for validation and newstest2014 for testing.Source: A. Experiment Setup · p. 4
  • The setup uses BiBERT for IWSLT’14 and WMT’14, mBERT for WMT’16, and dataset-specific teacher models for distillation.The IWSLT’14 and WMT’14 strong teacher is the model from; WMT’16 uses the strong model from, trained with back-translation data.Source: A. Experiment Setup · p. 5
  • NAT models use dropout 0.1 and learning rate 10^-4, with batch sizes and training steps set separately for each dataset.BLEU is used for evaluation, and latency is measured on a single Nvidia 3090 GPU with one sentence at a time.Source: A. Experiment Setup · p. 5

B. Hyperparameter Search

Hyperparameter search on IWSLT’14 DE→EN selected different upsampling ratios for BiBERT and mBERT. The experiments also favored frozen word embeddings, trainable output projection, mask insertion, and dynamic upsampling.

B. Hyperparameter Search

  • A frozen word-embedding layer with a trainable output projection performed better and became the standard configuration.These settings were tested with upsampling ratio s = 2 and FR.Source: B. Hyperparameter Search · p. 5, Figure 3 · p. 6
  • IM consistently outperformed IT, while DR was particularly beneficial at higher upsampling rates.BLEU varied more across ratios when using IT.Source: B. Hyperparameter Search · pp. 5–6, Figure 4 · p. 6
  • Extremely low or high upsampling rates did not yield favorable results.For BiBERT, the selected configuration used s = 4.0 with DR and IM, incorporating ED.Source: B. Hyperparameter Search · p. 6
  • s = 4.0 with BiBERT and s = 3.0 with mBERT were selected; the BiBERT model combined DR, IM, and ED.The mBERT ratio was chosen after BLEU rose across ratios 1.5 to 3.0, plateaued, then declined at higher ratios.Source: B. Hyperparameter Search · p. 6
  • For mBERT, BLEU rose gradually across upsampling ratios 1.5 to 3.0, plateaued, then declined as the ratio increased further.The selected mBERT upsampling ratio was s = 3.0.Source: B. Hyperparameter Search · p. 6, Figure 5 · p. 6

C. Main Result

CTCPMLM outperforms autoregressive and other NAT systems on several translation benchmarks, while retaining substantial inference speedups. Results vary by dataset and decoding setup, with CTC beam search not consistently improving performance.

Beam search

  • CTC beam search does not consistently improve performance; on IWSLT’14, it leaves En→De BLEU unchanged and substantially lowers De→En performance.The authors suggest the tested beam-search hyperparameters may not be optimal for this setting.Source: C. Main Result · pp. 6, 8

IWSLT’14 De↔En

  • On IWSLT’14 En↔De, CTCPMLM surpasses one-pass and iterative NAT models, outperforms autoregressive models including pretrained ones, and exceeds its teacher.The IWSLT’14 BLEU results are reported in Table III.Source: C. Main Result · pp. 6–7, Table III · p. 7

Inference speed

  • The model achieves a 16.35× speedup over the baseline autoregressive model; with beam search, the speedup is 7.16×.Both speedups are reported for the IWSLT’14 results.Source: C. Main Result · p. 7

WMT’14 En↔De

  • CTCPMLM surpasses existing NAT models and Transformer-base on WMT’14 En↔De, and sets a NAT state of the art there.The WMT’14 results are reported alongside WMT’16 En↔Ro in Table IV.Source: C. Main Result · p. 7, Table IV · p. 8

WMT’16 En↔Ro

  • CTCPMLM surpasses existing NAT models and Transformer-base on WMT’16 En↔Ro, setting a state of the art among NAT models.Table IV reports results for WMT’16 En↔Ro as well as WMT’14 En↔De.Source: C. Main Result · p. 7, Table IV · p. 8

Teacher and distillation comparisons

  • On IWSLT’14, distillation from base or strong teachers performs better than Transformer-base, and the strong-teacher-trained model also outperforms the strong teacher.Without knowledge distillation, the model still outperforms Transformer-base on IWSLT’14 and WMT’16, and surpasses the current NAT model on IWSLT’14.Source: C. Main Result · p. 7

D. Ablation Study

The ablations find gains from knowledge distillation, PMLM initialization, mask insertion, and embedding distillation in the appropriate settings. Backbone choice and the teacher layer used for embedding distillation also affect results.

Impact of variant techniques

  • Knowledge distillation provides performance gains across the ablation blocks, while PMLM initialization yields cumulative improvements and, alone, outperforms the baseline autoregressive teacher.Table VI reports the technique combinations on IWSLT’14 De→En.Source: D. Ablation Study · p. 9, Table VI · p. 9
  • Mask insertion (IM) improves CTCPMLM performance, including without PMLM initialization, and outperforms token duplication in the reported comparison.The comparison of IM and IT is also reported for WMT’16 En↔Ro.Source: D. Ablation Study · p. 9
  • Embedding distillation improves performance when paired with knowledge distillation and PMLM initialization, but worsens it when the model lacks PMLM initialization.The authors attribute the latter result to a mismatch between the model and PMLM representations; the combined setup performs best across the reported datasets and directions.Source: D. Ablation Study · p. 9

Different backbone models

  • BiBERT achieves the highest BLEU among the compared backbones with either base- or strong-teacher distillation; DistilBERT is faster, while XLMR is slowest.The paper attributes DistilBERT’s speed to its six layers and XLMR’s slower speed partly to its vocabulary being 2.4 times larger than mBERT’s.Source: D. Ablation Study · pp. 9–10

ED with Different Distillation Layers

  • Embedding distillation improves BLEU across the tested teacher layers, with scores from 39.8 to 39.94 and the highest score using the last two layers.Figure 6 compares distillation layers on the IWSLT’14 DE→EN test set; “None” denotes no embedding distillation.Source: D. Ablation Study · p. 10, Figure 6 · p. 10

V. RELATED WORK UTILIZING PLM INITIALIZATION IN NAT

Prior NAT systems use pretrained language models with additional modules, whereas this approach uses the pretrained model directly and adds embedding distillation to preserve its knowledge.

Prior NAT methods

  • Several NAT approaches add mechanisms to pretrained models, including adapters and a length token, a CRF, or latent transformation modules.The cited examples include AB-Net, Bert+CRF-NAT, DEER, and XLM-D.Source: §V · p. 10

Contrast with this approach

  • This approach directly uses pretrained language models without adapters, which the authors say preserves their layer-wise capabilities and reduces inference and storage overhead.The method also uses CTC to avoid explicit sequence-length prediction and embedding distillation to preserve pretrained knowledge.Source: §V · p. 10

Design rationale

  • CTC removes the need for explicit sequence-length prediction, while embedding distillation is used to retain knowledge from the pretrained language model.The authors frame these as distinct design choices in their approach.Source: §V · p. 10

A. Does the dependency on KD become looser as the model becomes stronger ?

With PMLM initialization, stronger knowledge distillation raises BLEU, while embedding distillation adds only a small training-time increase and no inference-time cost. Vocabulary pruning also improves mBERT’s inference speed and performance.

A. Does the dependency on KD become looser as the model becomes stronger ?

  • With PMLM initialization, BLEU rises from 37.85 without KD to 39.93 with stronger KD.The authors note that PMLM initialization alone outperforms NAT models with KD listed in Table III, while retaining KD as a performance-enhancing component.Source: A. Does the dependency on KD become… · p. 10

B. How about the efficiency when utilizing ED ?

speed while maintaining performance?

D. Results on Low-Resource Language Pairs

The reported results cover low-resource IWSLT language pairs and test whether CTCPMLM’s gains persist across parameter sizes. The conclusion also summarizes its performance against NAT and autoregressive systems.

D. Results on Low-Resource Language Pairs

E. Is the performance of the model due to the increase in the number of parameters?

VII. CONCLUSION

  • The conclusion highlights dynamic CTC upsampling, better performance from Inserting Masks than Inserting Tokens, and embedding distillation of contextualized PMLM target-space information.Source: §VII · p. 11
  • CTCPMLM achieves state-of-the-art results on IWSLT’14 De↔En, outperforms all existing NAT models on WMT’14 En↔De and WMT’16 En↔Ro, and is competitive with many strong autoregressive models.Source: §VII · p. 12
Loading 2306.06345v3…