# On Exposure Bias, Hallucination and Domain Shift
in Neural Machine Translation

Chaojun Wang1 Rico Sennrich2,1

1School of Informatics, University of Edinburgh
2Department of Computational Linguistics, University of Zurich
<EMAIL_ADDRESS><EMAIL_ADDRESS>
###### Abstract

The standard training algorithm in neural machine translation (NMT) suffers
from exposure bias, and alternative algorithms have been proposed to mitigate
this. However, the practical impact of exposure bias is under debate. In this
paper, we link exposure bias to another well-known problem in NMT, namely the
tendency to generate hallucinations under domain shift. In experiments on
three datasets with multiple test domains, we show that exposure bias is
partially to blame for hallucinations, and that training with Minimum Risk
Training, which avoids exposure bias, can mitigate this. Our analysis explains
why exposure bias is more problematic under domain shift, and also links
exposure bias to the beam search problem, i.e. performance deterioration with
increasing beam size. Our results provide a new justification for methods that
reduce exposure bias: even if they do not increase performance on in-domain
test sets, they can increase model robustness to domain shift.

##  1 Introduction

Neural Machine Translation (NMT) has advanced the state of the art in MT
(Sutskever et al., 2014; Bahdanau et al., 2015; Vaswani et al., 2017), but is
susceptible to domain shift. Koehn and Knowles (2017) consider out-of-domain
translation one of the key challenges in NMT. Such translations may be fluent,
but completely unrelated to the input (hallucinations), and their misleading
nature makes them particularly problematic.

We hypothesise that exposure bias (Ranzato et al., 2016), a discrepancy
between training and inference, makes this problem worse. Specifically,
training with teacher forcing only exposes the model to gold history, while
previous predictions during inference may be erroneous. Thus, the model
trained with teacher forcing may over-rely on previously predicted words,
which would exacerbate error propagation. Previous work has sought to reduce
exposure bias in training (Bengio et al., 2015; Ranzato et al., 2016; Shen et
al., 2016; Wiseman and Rush, 2016; Zhang et al., 2019). However, the relevance
of error propagation is under debate: Wu et al. (2018) argue that its role is
overstated in literature, and that linguistic features explain some of the
accuracy drop at higher time steps.

Previous work has established a link between domain shift and hallucination in
NMT (Koehn and Knowles, 2017; Müller et al., 2019). In this paper, we will aim
to also establish an empirical link between hallucination and exposure bias.
Such a link will deepen our understanding of the hallucination problem, but
also has practical relevance, e.g. to help predicting in which settings the
use of sequence-level objectives is likely to be helpful. We further
empirically confirm the link between exposure bias and the ‘beam search
problem’, i.e. the fact that translation quality does not increase
consistently with beam size (Koehn and Knowles, 2017; Ott et al., 2018;
Stahlberg and Byrne, 2019).

We base our experiments on German$\to$English IWSLT’14, and two datasets used
to investigate domain robustness by Müller et al. (2019): a selection of
corpora from OPUS (Lison and Tiedemann, 2016) for German$\to$English, and a
low-resource German$\to$Romansh scenario. We experiment with Minimum Risk
Training (MRT) (Och, 2003; Shen et al., 2016), a training objective which
inherently avoids exposure bias.

Our experiments show that MRT indeed improves quality more in out-of-domain
settings, and reduces the amount of hallucination. Our analysis of translation
uncertainty also shows how the MLE baseline over-estimates the probability of
random translations at all but the initial time steps, and how MRT mitigates
this problem. Finally, we show that the beam search problem is reduced by MRT.

##  2 Minimum Risk Training

The de-facto standard training objective in NMT is to minimize the negative
log-likelihood $\mathcal{L}(\bm{\theta})$ of the training data $D$111This is
equivalent to maximizing the likelihood of the data, hence Maximum Likelihood
Estimation (MLE).:

$\mathcal{L}(\bm{\theta})=\sum_{(\mathbf{x},\mathbf{y})\in
D}\sum_{t=1}^{|\mathbf{y}|}-\log
P\left(\mathbf{y}_{t}|\mathbf{x},\mathbf{y}_{<t};\bm{\theta}\right)$ (1)

where $\mathbf{x}$ and $\mathbf{y}$ are the source and target sequence,
respectively, $\mathbf{y}_{t}$ is the tth token in $\mathbf{y}$, and
$\mathbf{y}_{<t}$ denotes all previous tokens. MLE is typically performed with
teacher forcing, where $\mathbf{y}_{<t}$ are ground-truth labels in training,
which creates a mismatch to inference, where $\mathbf{y}_{<t}$ are model
predictions.

Minimum Risk Training (MRT) is a sequence-level objective that avoids this
problem. Specifically, the objective function of MRT is the expected loss
(risk) with respect to the posterior distribution:

$\mathcal{R}(\bm{\theta})=\sum_{(\mathbf{x},\mathbf{y})\in
D}\sum_{\mathbf{\tilde{y}}\in\mathcal{Y}(\mathbf{x})}P\left(\mathbf{\tilde{y}}|\mathbf{x};\bm{\theta}\right)\Delta\left(\mathbf{\tilde{y}},\mathbf{y}\right)$
(2)

in which the loss $\Delta\left(\mathbf{\tilde{y}},\mathbf{y}\right)$ indicates
the discrepancy between the gold translation $\mathbf{y}$ and the model
prediction $\mathbf{\tilde{y}}$. Due to the intractable search space, the
posterior distribution $\mathcal{Y}(\mathbf{x})$ is approximated by a subspace
$\mathcal{S}(\mathbf{x})$ by sampling a certain number of candidate
translations, and normalizing:

$\tilde{P}\left(\mathbf{\tilde{y}}|\mathbf{x};\bm{\theta},\alpha\right)=\frac{P\left(\mathbf{\tilde{y}}|\mathbf{x};\bm{\theta}\right)^{\alpha}}{\sum_{\mathbf{y}^{\prime}\in\mathcal{S}\left(\mathbf{x}\right)}P\left(\mathbf{y}^{\prime}|\mathbf{x};\bm{\theta}\right)^{\alpha}}$
(3)

where $\alpha$ is a hyperparameter to control the sharpness of the subspace.
Based on preliminary results, we use random sampling to generate candidate
translations, and following Edunov et al. (2018), do not add the reference
translation to the subspace.

##  3 Experiments

###  3.1 Data

To verify the effectiveness of our MRT implementation on top of a strong
Transformer baseline (Vaswani et al., 2017), we first conduct experiments on
IWSLT’14 German$\to$English (DE$\to$EN) (Cettolo et al., 2014), which consists
of $180\,000$ sentence pairs. We follow previous work for data splits (Ranzato
et al., 2016; Edunov et al., 2018).

For experiments with domain shift, we use data sets and preprocessing as
Müller et al. (2019)222https://github.com/ZurichNLP/domain-robustness. For
DE$\to$EN, data comes from OPUS (Lison and Tiedemann, 2016), and is comprised
of five domains: medical, IT, law, koran and subtitles. We use medical for
training and development, and report results on an in-domain test set and the
four other domains (out-of-domain; OOD). German$\to$Romansh (DE$\to$RM) is a
low-resource language pair where robustness to domain shift is of practical
relevance. The training data is from the Allegra corpus (Scherrer and Cartoni,
2012) (law domain) with $100\,000$ sentence pairs. The test domain are blogs,
using data from Convivenza333https://www.suedostschweiz.ch/blogs/convivenza.
We have access to 2000 sentences for development and testing, respectively, in
each domain.

We tokenise and truecase data sets with Moses (Koehn et al., 2007), and use
shared BPE with $32\,000$ units (Sennrich et al., 2016).

###  3.2 Model

We implement444Code available at https://github.com/zippotju/Exposure-Bias-
Hallucination-Domain-Shift MRT in the Nematus toolkit (Sennrich et al., 2017).
All our experiments use the Transformer architecture (Vaswani et al., 2017).
Following Edunov et al. (2018), we use $1$-$\text{BLEU}_{\text{smooth}}$ (Lin
and Och, 2004) as the MRT loss. Models are pre-trained with the token-level
objective MLE and then fine-tuned with MRT. Hyperparameters mostly follow
previous work (Edunov et al., 2018; Müller et al., 2019); for MRT, we conduct
limited hyperparameter search on the IWSLT’14 development set, including
learning rate, batch size, and the sharpness parameter $\alpha$. We set the
number of candidate translations for MRT to 4 to balance effectiveness and
efficiency. Detailed hyperparameters are reported in the Appendix.

###  3.3 Evaluation

For comparison to previous work, we report lowercased, tokenised BLEU
(Papineni et al., 2002) with multi-bleu.perl for IWSLT’14, and cased,
detokenised BLEU with SacreBLEU (Post, 2018)555Signature:
BLEU+c.mixed+#.1+s.exp+tok.13a+v.1.4.2 otherwise. For settings with domain
shift, we report average and standard deviation of 3 independent training runs
to account for optimizer instability.

The manual evaluation was performed by two native speakers of German who
completed bilingual (German/English) high school or University programs. We
collected $\sim$3600 annotations in total, spread over 12 configurations. We
ask annotators to evaluate translations according to fluency and adequacy. For
fluency, the annotator classifies a translation as fluent, partially fluent or
not fluent; for adequacy, as adequate, partially adequate or inadequate. We
report kappa coefficient ($K$) (Carletta, 1996) for inter-annotator and intra-
annotator agreement in Table 1, and assess statistical significance with
Fisher’s exact test (two-tailed).

| inter-annotator | intra-annotator
---|---|---
annotation | $P(A)$ | $P(E)$ | $K$ | $P(A)$ | $P(E)$ | $K$
fluency | 0.66 | 0.38 | 0.44 | 0.87 | 0.42 | 0.77
adequacy | 0.82 | 0.61 | 0.54 | 0.93 | 0.66 | 0.79
Table 1: Inter-annotator (N=250) and intra-annotator agreement (N=617) of
manual evaluation.

###  3.4 Results

Table 2 shows results for IWSLT’14. We compare to results by Edunov et al.
(2018), who use a convolutional architecture (Gehring et al., 2017), and Wu et
al. (2019), who report results with Transfomer-base and dynamic convolution.

system | BLEU
---|---
ConvS2S (MLE) (Edunov et al., 2018) | 32.2 (+0.6)
ConvS2S (MRT) (Edunov et al., 2018) | 32.8 (+0.6)
Transformer (MLE) (Wu et al., 2019) | 34.4 (+0.6)
DynamicConv (MLE) (Wu et al., 2019) | 35.2 (+0.6)
MLE | 34.7 (+0.6)
MRT | 35.2 (+0.5)
Table 2: Results for IWSLT’14 DE$\to$EN with MLE and MRT (in brackets, improvement over MLE). | DE$\to$EN | DE$\to$RM
---|---|---
system | in-domain | average OOD | in-domain | average OOD
SMT Müller et al. (2019) | 58.4 ($\pm$0.00) | 11.8 ($\pm$0.00) | 45.2 ($\pm$0.00) | 15.5 ($\pm$0.00)
NMT Müller et al. (2019) | 61.5 ($\pm$0.00) | 11.7 ($\pm$0.00) | 52.5 ($\pm$0.00) | 18.9 ($\pm$0.00)
NMT+RC+SR+NC Müller et al. (2019) | 60.8 ($\pm$0.00) | 13.1 ($\pm$0.00) | 52.4 ($\pm$0.00) | 20.7 ($\pm$0.00)
MLE w/o LS | 58.3 ($\pm$0.53) |  09.7 ($\pm$0.25) | 52.2 ($\pm$0.19) | 15.8 ($\pm$0.39)
+MRT | 58.4 ($\pm$0.39) | 10.2 ($\pm$0.26) | 52.1 ($\pm$0.08) | 15.9 ($\pm$0.28)
MLE w/ LS | 58.9 ($\pm$0.45) | 11.2 ($\pm$0.16) | 53.9 ($\pm$0.16) | 18.0 ($\pm$0.17)
+MRT | 58.8 ($\pm$0.36) | 12.0 ($\pm$0.29) | 53.9 ($\pm$0.12) | 18.7 ($\pm$0.09)
Table 3: Average BLEU and standard deviation on in-domain and out-of-domain
test sets for models trained on OPUS (DE$\to$EN) and Allegra (DE$\to$RM). RC:
reconstruction; SR: subword regularization, NC: noisy channel.

With 34.7 BLEU, our baseline is competitive. We observe an improvement of 0.5
BLEU from MRT, comparable to Edunov et al. (2018), although we start from a
stronger baseline (+2.5 BLEU).

Table 3 shows results for data sets with domain shift. To explore the effect
of label smoothing (Szegedy et al., 2016), we train baselines with and without
label smoothing. MLE with label smoothing performs better by itself, and we
also found MRT to be more effective on top of the initial model with label
smoothing. For DE$\to$EN, MRT increases average OOD BLEU by 0.8 compared to
the MLE baseline with label smoothing; for DE$\to$RM the improvement is 0.7
BLEU. We note that MRT does not consistently improve in-domain performance,
which is a first indicator that exposure bias may be more problematic under
domain shift.

Our OOD results lag slightly behind those of Müller et al. (2019), but note
that the techniques employed by them, namely reconstruction (Tu et al., 2017;
Niu et al., 2019), subword regularization (Kudo, 2018), and noisy channel
modelling (Li and Jurafsky, 2016) are orthogonal to MRT. We leave the
combination of these approaches to future work.

##  4 Analysis

BLEU results indicate that MRT can improve domain robustness. In this section,
we report on additional experiments to establish more direct links between
exposure bias and domain robustness, hallucination, and the beam search
problem. Experiments are performed on DE$\to$EN OPUS data.

###  4.1 Hallucination

We manually evaluate the proportion of hallucinated translations on out-of-
domain and in-domain test sets. We follow the definition and evaluation by
Müller et al. (2019), considering a translation a hallucination if it is
(partially) fluent, but unrelated in content to the source text (inadequate).
We report the proportion of such hallucinations for each system.

| % hallucinations (BLEU)
---|---
system | out-of-domain | in-domain
MLE w/o LS | 35% 0(9.7) | 2% (58.3)
+MRT | 29% (10.2) | -
MLE w/ LS | 33% (11.2) | 1% (58.9)
+MRT | 26% (12.0) | -
Table 4: Proportion of hallucinations and BLEU on out-of-domain and in-domain
test sets. DE$\to$EN OPUS.

Results in Table 4 confirm that hallucinations are much more pronounced in
out-of-domain test sets (33–35%) than in in-domain test sets (1–2%). MRT
reduces the proportion of hallucinations on out-of-domain test sets (N=500 for
each system; reductions statistically significant at $p<0.05$) and improves
BLEU. Note that the two metrics do not correlate perfectly: MLE with label
smoothing has higher BLEU (+1) than MRT based on MLE without label smoothing,
but a similar proportion of hallucinations. This indicates that label
smoothing increases translation quality in other aspects, while MRT has a
clear effect on the number of hallucinations, reducing it by up to 21%
(relative).

source |  Wir haben ihn gefunden.
---|---
reference |  We found him.
MLE | Do not pass it.
MRT |  We have found it.
source | So höre nicht auf die Ableugner.
reference | So hearken not to those who deny (the Truth).
MLE |  Do not drive or use machines.
MRT |  Do not apply to dleugner.
Table 5: Out-of-domain translation examples. MLE hallucinates in both
examples; MRT was rated more adequate in top example, less fluent in bottom
one.

A closer inspection of segments where the MLE system was found to hallucinate
shows that some segments were scored higher in adequacy with MRT, others lower
in fluency. One example for each case is shown in Table 5. Even the example
where MRT was considered disfluent and inadequate actually shows an attempt to
cover the source sentence: the source word ‘Ableugner’ (denier) is
mistranslated into ‘dleugner’. We consider this preferable to producing a
complete hallucination.

###  4.2 Uncertainty Analysis

Figure 1: Per-token probability of out-of-domain reference translations and
in-domain distractors (first two graphs share legend). Rightmost plot shows
direct comparison for MLE baseline and final MRT model. DE$\to$EN OPUS .

Inspired by Ott et al. (2018), we analyse the model’s uncertainty by computing
the average probability at each time step across a set of sentences. Besides
the reference translations, we also consider a set of ‘distractor’
translations, which are random sentences from the in-domain test set which
match the corresponding reference translation in length.

In Figure 1, we show out-of-domain results for an MLE model and multiple
checkpoints of MRT fine-tuning. The left two graphs show probabilities for
references and distractors, respectively. The right-most graph shows a direct
comparison of probabilities for references and distractors for the MLE
baseline and the final MRT model. The MLE baseline assigns similar
probabilities to tokens in the references and the distractors. Only for the
first time steps is there a clear preference for the references over the
(mostly random!) distractors. This shows that error propagation is a big risk:
should the model make a wrong prediction initially, this is unlikely to be
penalised in later time steps.

MRT tends to increase the model’s certainty at later time steps666The
uncertainty of the baseline is due to label smoothing., but importantly, the
increase is sharper for the reference translations than for the distractors.
The direct comparison shows a widening gap in certainty between the reference
and distractor sentences.777For intermediate checkpoints, see Appendix, Figure
2. In other words, producing a hallucination will incur a small penalty at
each time step (compared to producing the reference), presumably due to a
higher reliance on the source signal, lessening the risk of error propagation
and hallucinations.

Our analysis shows similar trends on in-domain references. However, much
higher probabilities are assigned to the first few tokens of the references
than to the distractors. Hence, it is much less likely that a hallucination is
kept in the beam, or will overtake a good translation in overall probability,
reducing the practical impact of the model’s over-reliance on its
history.888Figures are shown in the Appendix (Figure 3).

###  4.3 Beam Size Analysis

Figure 1 shows that with MLE, distractor sentences are assigned lower
probabilities than the references at the first few time steps, but are
assigned similar, potentially even higher probabilities at later time steps.
This establishes a connection between exposure bias and the beam search
problem, i.e. the problem that increasing the search space can lead to worse
model performance.999The beam search problem has previously been linked to
length bias (Yang et al., 2018; Murray and Chiang, 2018) and the copy mode
(Ott et al., 2018). We consider hallucinations another result of using large
search spaces with MLE models. With larger beam size, it is more likely that
hallucinations survive pruning at the first few time steps, and with high
probabilities assigned to them at later time steps, there is a chance that
they become the top-scoring translation.

We investigate whether the beam search problem is mitigated by MRT. In Table
6, we report OOD BLEU and the proportion of hallucinations with beam sizes of
1, 4 and 50. While MRT does not eliminate the beam search problem, performance
drops less steeply as beam size increases. With beam size 4, our MRT models
outperform the MLE baseline by 0.5-0.8 BLEU; with beam size 50, this
difference grows to 0.6-1.5 BLEU. Our manual evaluation (N=200 for each system
for beam size 1 and 50) shows that the proportion of hallucinations increases
with beam size, and that MRT consistently reduces the proportion by 11-21%
(relative). For the system with label smoothing, the relative increase in
hallucinations with increasing beam size is also smaller with MRT (+33%) than
with MLE (+44%).

| BLEU (% hallucinations)
---|---
system | $k=1$ | $k=4$ | $k=50$
MLE w/o LS |  08.9 (28%) |  09.7 (35%) |  09.3 (37%)
+MRT |  09.1 (24%) | 10.2 (29%) |  09.9 (33%)
MLE w/ LS | 10.6 (27%) | 11.2 (33%) |  09.4 (39%)
+MRT | 11.3 (24%) | 12.0 (26%) | 10.9 (32%)
Table 6: Average OOD BLEU and proportion of hallucinations with different beam
sizes $k$. DE$\to$EN OPUS.

##  5 Conclusions

Our results and analysis show a connection between the exposure bias due to
MLE training with teacher forcing and several well-known problems in neural
machine translation, namely poor performance under domain shift, hallucinated
translations, and deteriorating performance with increasing beam size. We find
that Minimum Risk Training, which does not suffer from exposure bias, can be
useful even when it does not increase performance on an in-domain test set: it
increases performance under domain shift, reduces the number of hallucinations
substantially, and makes beam search with large beams more stable.

Our findings are pertinent to the academic debate how big of a problem
exposure bias is in practice – we find that this can vary substantially
depending on the dataset –, and they provide a new justification for sequence-
level training objectives that reduce or eliminate exposure bias. Furthermore,
we believe that a better understanding of the links between exposure bias and
well-known translation problems will help practitioners decide when sequence-
level training objectives are especially promising, for example in settings
where the test domain is unknown, or where hallucinations are a common
problem.

## Acknowledgments

Chaojun Wang was supported by the UK Engineering and Physical Sciences
Research Council (EPSRC) fellowship grant EP/S001271/1 (MTStretch). Rico
Sennrich acknowledges support of the Swiss National Science Foundation
(MUTAMUR; no. 176727). This project has received support from Samsung
Electronics Polska sp. z o.o. - Samsung R&D Institute Poland.

## References

  * 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, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings_.
  * Bengio et al. (2015) Samy Bengio, Oriol Vinyals, Navdeep Jaitly, and Noam Shazeer. 2015.  Scheduled sampling for sequence prediction with recurrent neural networks.  In C. Cortes, N. D. Lawrence, D. D. Lee, M. Sugiyama, and R. Garnett, editors, _Advances in Neural Information Processing Systems 28_ , pages 1171–1179. Curran Associates, Inc.
  * Carletta (1996) Jean Carletta. 1996.  Assessing agreement on classification tasks: The kappa statistic.  _Computational Linguistics_ , 22(2):249–254.
  * Cettolo et al. (2014) Mauro Cettolo, Jan Niehues, Sebastian Stüker, Luisa Bentivogli, and Marcello Federico. 2014.  Report on the 11th iwslt evaluation campaign, iwslt 2014.  In _Proceedings of the International Workshop on Spoken Language Translation, Hanoi, Vietnam_ , page 57.
  * Edunov et al. (2018) Sergey Edunov, Myle Ott, Michael Auli, David Grangier, and Marc’Aurelio Ranzato. 2018.  Classical structured prediction losses for sequence to sequence learning.  In _Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long Papers)_ , pages 355–364, New Orleans, Louisiana. Association for Computational Linguistics.
  * Gehring et al. (2017) Jonas Gehring, Michael Auli, David Grangier, Denis Yarats, and Yann N. Dauphin. 2017\.  Convolutional sequence to sequence learning.  In _Proceedings of the 34th International Conference on Machine Learning - Volume 70_ , ICML’17, pages 1243–1252. JMLR.org.
  * Koehn et al. (2007) Philipp Koehn, Hieu Hoang, Alexandra Birch, Chris Callison-Burch, Marcello Federico, Nicola Bertoldi, Brooke Cowan, Wade Shen, Christine Moran, Richard Zens, Chris Dyer, Ondřej Bojar, Alexandra Constantin, and Evan Herbst. 2007\.  Moses: Open source toolkit for statistical machine translation.  In _Proceedings of the 45th Annual Meeting of the Association for Computational Linguistics Companion Volume Proceedings of the Demo and Poster Sessions_ , pages 177–180, Prague, Czech Republic. Association for Computational Linguistics.
  * Koehn and Knowles (2017) Philipp Koehn and Rebecca Knowles. 2017.  Six challenges for neural machine translation.  In _Proceedings of the First Workshop on Neural Machine Translation_ , pages 28–39, Vancouver. Association for Computational Linguistics.
  * Kudo (2018) Taku Kudo. 2018.  Subword regularization: Improving neural network translation models with multiple subword candidates.  In _Proceedings of the 56th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)_ , pages 66–75, Melbourne, Australia. Association for Computational Linguistics.
  * Li and Jurafsky (2016) Jiwei Li and Dan Jurafsky. 2016.  Mutual information and diverse decoding improve neural machine translation.
  * Lin and Och (2004) Chin-Yew Lin and Franz Josef Och. 2004.  ORANGE: a method for evaluating automatic evaluation metrics for machine translation.  In _COLING 2004: Proceedings of the 20th International Conference on Computational Linguistics_ , pages 501–507, Geneva, Switzerland. COLING.
  * Lison and Tiedemann (2016) Pierre Lison and Jörg Tiedemann. 2016.  OpenSubtitles2016: Extracting large parallel corpora from movie and TV subtitles.  In _Proceedings of the Tenth International Conference on Language Resources and Evaluation (LREC 2016)_ , pages 923–929, Portorož, Slovenia. European Language Resources Association (ELRA).
  * Müller et al. (2019) Mathias Müller, Annette Rios, and Rico Sennrich. 2019.  Domain robustness in neural machine translation.
  * Murray and Chiang (2018) Kenton Murray and David Chiang. 2018.  Correcting length bias in neural machine translation.  In _Proceedings of the Third Conference on Machine Translation: Research Papers_ , pages 212–223, Belgium, Brussels. Association for Computational Linguistics.
  * Niu et al. (2019) Xing Niu, Weijia Xu, and Marine Carpuat. 2019.  Bi-directional differentiable input reconstruction for low-resource neural machine translation.  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 442–448, Minneapolis, Minnesota. Association for Computational Linguistics.
  * Och (2003) Franz Josef Och. 2003.  Minimum error rate training in statistical machine translation.  In _Proceedings of the 41st Annual Meeting of the Association for Computational Linguistics_ , pages 160–167, Sapporo, Japan. Association for Computational Linguistics.
  * Ott et al. (2018) Myle Ott, Michael Auli, David Grangier, and Marc’Aurelio Ranzato. 2018.  Analyzing uncertainty in neural machine translation.  In _Proceedings of the 35th International Conference on Machine Learning_ , Stockholm, Sweden.
  * Papineni et al. (2002) Kishore Papineni, Salim Roukos, Todd Ward, and Wei-Jing Zhu. 2002.  Bleu: a method for automatic evaluation of machine translation.  In _Proceedings of the 40th Annual Meeting of the Association for Computational Linguistics_ , pages 311–318, Philadelphia, Pennsylvania, USA. Association for Computational Linguistics.
  * 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.
  * Ranzato et al. (2016) Marc’Aurelio Ranzato, Sumit Chopra, Michael Auli, and Wojciech Zaremba. 2016.  Sequence level training with recurrent neural networks.  In _4th International Conference on Learning Representations, ICLR 2016, San Juan, Puerto Rico, May 2-4, 2016, Conference Track Proceedings_.
  * Scherrer and Cartoni (2012) Yves Scherrer and Bruno Cartoni. 2012.  The trilingual ALLEGRA corpus: Presentation and possible use for lexicon induction.  In _Proceedings of the Eighth International Conference on Language Resources and Evaluation (LREC-2012)_ , pages 2890–2896, Istanbul, Turkey. European Languages Resources Association (ELRA).
  * Sennrich et al. (2017) Rico Sennrich, Orhan Firat, Kyunghyun Cho, Alexandra Birch, Barry Haddow, Julian Hitschler, Marcin Junczys-Dowmunt, Samuel Läubli, Antonio Valerio Miceli Barone, Jozef Mokry, and Maria Nădejde. 2017.  Nematus: a toolkit for neural machine translation.  In _Proceedings of the Software Demonstrations of the 15th Conference of the European Chapter of the Association for Computational Linguistics_ , pages 65–68, Valencia, Spain. 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)_ , pages 1715–1725, Berlin, Germany. Association for Computational Linguistics.
  * Shen et al. (2016) Shiqi Shen, Yong Cheng, Zhongjun He, Wei He, Hua Wu, Maosong Sun, and Yang Liu. 2016\.  Minimum risk training for neural machine translation.  In _Proceedings of the 54th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)_ , pages 1683–1692, Berlin, Germany. Association for Computational Linguistics.
  * Stahlberg and Byrne (2019) Felix Stahlberg and Bill Byrne. 2019.  On NMT search errors and model errors: Cat got your tongue? In _Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP)_ , pages 3347–3353, Hong Kong, China. Association for Computational Linguistics.
  * Sutskever et al. (2014) Ilya Sutskever, Oriol Vinyals, and Quoc V Le. 2014.  Sequence to sequence learning with neural networks.  In Z. Ghahramani, M. Welling, C. Cortes, N. D. Lawrence, and K. Q. Weinberger, editors, _Advances in Neural Information Processing Systems 27_ , pages 3104–3112. Curran Associates, Inc.
  * Szegedy et al. (2016) C. Szegedy, V. Vanhoucke, S. Ioffe, J. Shlens, and Z. Wojna. 2016.  Rethinking the inception architecture for computer vision.  In _2016 IEEE Conference on Computer Vision and Pattern Recognition (CVPR)_ , pages 2818–2826.
  * Tu et al. (2017) Zhaopeng Tu, Yang Liu, Lifeng Shang, Xiaohua Liu, and Hang Li. 2017.  Neural machine translation with reconstruction.  In _Proceedings of the Thirty-First AAAI Conference on Artificial Intelligence_ , AAAI’17, pages 3097–3103. AAAI Press.
  * Vaswani et al. (2017) Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Ł ukasz Kaiser, and Illia Polosukhin. 2017.  Attention is all you need.  In I. Guyon, U. V. Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, editors, _Advances in Neural Information Processing Systems 30_ , pages 5998–6008. Curran Associates, Inc.
  * Wiseman and Rush (2016) Sam Wiseman and Alexander M. Rush. 2016.  Sequence-to-sequence learning as beam-search optimization.  In _Proceedings of the 2016 Conference on Empirical Methods in Natural Language Processing_ , pages 1296–1306, Austin, Texas. Association for Computational Linguistics.
  * Wu et al. (2019) Felix Wu, Angela Fan, Alexei Baevski, Yann Dauphin, and Michael Auli. 2019.  Pay less attention with lightweight and dynamic convolutions.  In _International Conference on Learning Representations_.
  * Wu et al. (2018) Lijun Wu, Xu Tan, Di He, Fei Tian, Tao Qin, Jianhuang Lai, and Tie-Yan Liu. 2018\.  Beyond error propagation in neural machine translation: Characteristics of language also matter.  In _Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing_ , pages 3602–3611, Brussels, Belgium. Association for Computational Linguistics.
  * Yang et al. (2018) Yilin Yang, Liang Huang, and Mingbo Ma. 2018.  Breaking the beam search curse: A study of (re-)scoring methods and stopping criteria for neural machine translation.  In _Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing_ , pages 3054–3059, Brussels, Belgium. Association for Computational Linguistics.
  * Zhang et al. (2019) Wen Zhang, Yang Feng, Fandong Meng, Di You, and Qun Liu. 2019.  Bridging the gap between training and inference for neural machine translation.  In _Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics_ , pages 4334–4343, Florence, Italy. Association for Computational Linguistics.

##  Appendix A Appendix

| IWSLT | OPUS/Allegra
---|---|---
General hyperparameters |  |
embedding layer size | 512
hidden state size | 512
tie encoder decoder embeddings | yes
tie decoder embeddings | yes
loss function | per-token-cross-entropy (MRT)
label smoothing | 0.1
optimizer | adam
learning schedule | transformer (constant)
warmup steps | 4000 | 6000
gradient clipping threshold | 1 | 0
maximum sequence length | 100
token batch size | 4096
length normalization alpha | 0.6 | 1
encoder depth | 6
decoder depth | 6
feed forward num hidden | 1024 | 2048
number of attention heads | 4 | 8
embedding dropout | 0.3 | 0.1
residual dropout | 0.3 | 0.1
relu dropout | 0.3 | 0.1
attention weights dropout | 0.3 | 0.1
beam size | 4
| beam search sampling | random sampling
MRT-revelant hyperparameters |  |
learning rate | 0.00003 | 0.00001
batch size | 8192 (tokens) | 10 (sentences)
sharpness alpha | 0.005 | 0.005

Table 7: Configurations of NMT systems used to pre-train and fine-tune over
three datasets. Note in general hyperparameters, the items in brackets denote
the options that will be used in MRT fine-tuning.

Figure 2: Per-token probability of out-of-domain reference translations and
in-domain distractors for different checkpoints in MRT training, showing a
widening gap between references and distractors. DE$\to$EN OPUS. Figure 3:
Per-token probability of in-domain reference translations and distractors.
Rightmost plot shows direct comparison for MLE baseline and final MRT model.
DE$\to$EN OPUS.


# An End-to-End Network Pruning Pipeline with Sparsity Enforcement

Evan Dogariu
Department of Computer Science
Princeton University
<EMAIL_ADDRESS>

###### Abstract

Neural networks have emerged as a powerful tool for solving complex tasks
across various domains, but their increasing size and computational
requirements have posed significant challenges in deploying them on resource-
constrained devices. Neural network sparsification, and in particular pruning,
has emerged as an effective technique to alleviate these challenges by
reducing model size, computational complexity, and memory footprint while
maintaining competitive performance. However, many pruning pipelines modify
the standard training pipeline at only a single stage, if at all. In this
work, we look to develop an end-to-end training pipeline that befits neural
network pruning and sparsification at all stages of training. To do so, we
make use of nonstandard model parameter initialization, pre-pruning training
methodologies, and post-pruning training optimizations. We conduct experiments
utilizing combinations of these methods, in addition to different techniques
used in the pruning step, and find that our combined pipeline can achieve
significant gains over current state of the art approaches to neural network
sparsification.

##  1 Introduction

Neural networks have revolutionized the field of machine learning by achieving
state-of-the-art performance in various complex tasks. However, the growing
demand for deploying these models on resource-constrained devices, such as
mobile phones, embedded systems, and Internet of Things (IoT) devices, makes
it valuable to investigate methods to decrease inference time and memory
footprint without sacrificing performance. Various techniques that fall under
the category of neural network sparsification have been developed to tackle
this problem, and one of the most promising subfields is that of neural
network pruning.

The goal of pruning is to determine which subset of the network parameters are
needed for good performance, and to retrain this subnetwork. This standard
pruning pipeline is motivated by the empirical support for the Lottery Ticket
Hypothesis lth , which suggests that within an overparameterized neural
network, there exist sparse subnetworks, or "winning tickets," that can
achieve comparable performance to the original dense network when trained in
isolation. This groundbreaking work demonstrated the potential for pruning to
discover highly efficient neural network architectures without losing
accuracy.

##  2 Related Works

In this section, we relate various prior methods for network pruning and for
the training of sparse neural networks. We will make use of these techniques
as building blocks of our pipeline, baselines to compare against, and
motivations for our research.

###  2.1 Neural Network Pruning

Neural network pruning is a technique that involves selectively removing
connections, neurons, or filters from a trained neural network, resulting in a
sparser architecture. More precisely, given a neural network $f(x,\theta)$
parameterized by $\theta$, the objective is to generate a mask
$M\in\\{0,1\\}^{|\theta|}$ (here, $|\theta|$ is the number of parameters) with
the following two desirable properties:

  1. 1.

the mask is sparse; that is, $||M||_{0}\ll|\theta|$ (here, $||\cdot||_{0}$ is
the $L^{0}$ "norm", measuring the number of nonzero elements)

  2. 2.

applying the mask to the parameters of the network doesn’t compromise
performance; that is, the model $f(x,\theta\odot M)$ performs well (here,
$\odot$ is the element-wise product)

Techniques for building the mask $M$ that make use of training a model are
called training-based. In such methods, we first train a model $f(x,\theta)$
to convergence, and then apply our pruning method to this trained model
(perhaps iteratively). These methods are generally the most common, and also
what motivated the Lottery Ticket Hypothesis lth . There are also newer
methods for constructing the mask $M$ that do not require any learning updates
on the model, but are instead based solely on the model’s architecture and the
initialization of the parameters $\theta$; we call these methods training-
free. We recount methods of both types below.

Magnitude-Based Pruning. The main method in this class is magnitude-based
pruning, in which weights with the lowest magnitude are assumed to be
negligible and can be pruned away liang2021pruning ; importantly, this method
is one-shot and only requires one round of training. Furthermore, the
technique is interpretable, efficient, and easy to implement.

Iterative Pruning. Another technique is iterative pruning, in which magnitude-
based pruning is applied in rounds, pruning a small subset of the weights each
iteration paganini2020iterative . This approach can often demonstrate better
performance than a single round of magnitude-based pruning, but is much more
computationally burdensome.

Random Pruning. In random pruning, networks are pruned randomly at
initialization. This is the most naive way to approach network pruning, and
generally has poor performance.

SNIP. In single-shot network pruning (SNIP) snip , networks are pruned at
initialization based on a saliency criterion that determines structurally
important connections. The motivation is that connections that are
architecturally important on initialization become the connections that are
needed in the trained model.

GraSP. In gradient signal preservation (GraSP) grasp , networks are pruned at
initialization in a fashion that preserves the overall gradient flow through
the network. The idea here is that the gradient flow is what dominates
training trajectories, and so weights that don’t contribute to the gradient
flow will not be learned well.

SynFlow. In iterative synaptic flow pruning (SynFlow) synflow , weights are
pruned on initialization in a way that preserves the total flow of synaptic
strengths, avoiding the layer collapse that comes about from gradient-based
methods.

###  2.2 Sparse Training Methods

Below, we mention several works that are tailored specifically to training
neural networks that either are or will be sparse.

ZerO Initialization. One such method is the ZerO initialization zhao2022zero ,
which makes use of a theoretically-justified deterministic initialization that
begins as low rank as possible. This sharply deviates from the tried and true
practice of random initialization, which begins with maximal rank
transformations at each layer, but it has been empirically shown to produce
networks that can be more effectively pruned.

Learnable Masks. Another approach, detailed in srinivas2016training , is to
learn a sparsity mask alongside the weights of the model, and to prune
according to this learned mask. During the pre-pruning training, an extra
regularization term is added to the loss to ensure that the learned mask is
indeed sparse. This technique, while certainly improving pruning capabilities,
effectively doubles the memory footprint of the pre-pruning training. However,
later stages of training and the resulting final model do not suffer this
memory penalty.

ToST. The sparse training toolkit ToST jaiswal2022training  is a collection of
small optimizations and changes to the traditional training methodology that
is curated specifically for training sparse neural networks. It proposes
annealed activation smoothing and decaying skip connection injection, as well
as a layer-wise rescaled initialization. Code is provided for each of these
methods.

RigL. In the paper "Rigging the Lottery: Making All Tickets Winners" (RigL)
evci2021rigging , Evci et al. design a training procedure that dynamically
updates the sparse network topology, resulting in a fully-trained, sparse
model in a single training stage. This methodology has been shown to
accomplish state of the art on sparse training, making it a formidible
baseline to compare against the pipeline that we develop. RigL operates by
growing significant connections and pruning insignificant ones many times
throughout the course of training.

#####

##  3 Methodology

In this paper, we build an end-to-end pipeline for training sparse neural
networks that is based on the training-based pruning paradigm. Given a
sparsity parameter $s$ denoting the desired sparsity level ($s$ is the
proportion of nonzero parameters to total parameters), we proceed as follows:

  1. 1.

In the pre-pruning stage, we train a model $f(x,\theta)$ to convergence using
the methods detailed in 3.1.

  2. 2.

Using the trained model $f(x,\theta)$, we apply magnitude-based pruning on the
model parameters $\theta$ in order to construct a mask
$M\in\\{0,1\\}^{|\theta|}$ such that $||M||_{0}\leq s|\theta|$.

  3. 3.

In the post-pruning stage, given the sparse mask $M$, we train the sparse
model $f(x,\theta\odot M)$ to convergence using the methods detailed in 3.2,
resulting in a fully-trained, sparse model.

This pipeline has the goal of pruning and network sparsity in mind from the
beginning to the end, with each design decision curated to this specific
setting. Furthermore, most prior approaches tend to seek undamaged performance
for sparsity settings on the order of $s\sim 0.1$. We aim to push that an
order of magnitude lower.

###  3.1 Pre-Pruning Training

The end goal of the pre-pruning training stage is a neural network that is
conducive to magnitude-based pruning. As such, we wish to deviate from the
traditional deep learning training pipeline in ways that will help enforce
sparsity of the model, measured by weight magnitude. Two such techniques will
be useful for this purpose.

#### ZerO Initialization

Normally, neural networks are initialized randomly with Gaussian weights
sampled from $\mathcal{N}\left(0,\frac{C}{n_{\ell}}\right)$ for some constant
$C$ with $n_{\ell}$ as either the fan-in or fan-out number. Clearly, this
initialization begins with all weights having the same magnitude in
expectation. Also, using the setting of a multilayer perceptron as a schematic
example, each weight matrix initialized in this way is full-rank. As such,
this initialization scheme (which is designed for general training) is not
immediately specialized for the goal of training a network to be sparse. It
then makes sense to apply an initialization scheme that is specialized for
such a task. For this purpose, we use the recently developed ZerO
initialization zhao2022zero , which we detail below for linear layers. First,
we define the types of matrices with which this initialization is built,
following the notation in zhao2022zero .

######  Definition 1 (Partial identity matrix).

Let $\mathbf{I}_{n}$ denote the $n\times n$ identity matrix, and let
$\mathbf{0_{r,c}}$ denote the $r\times c$ zero matrix. For given dimensions
$r,c\in\mathbb{N}$, the $r\times c$ partial identity matrix
$\mathbf{I}^{*}_{r,c}$ is defined as

$\mathbf{I}^{*}:=\begin{cases}(\mathbf{I}_{r},\mathbf{0}_{r,c-r})&r<c\\\
\mathbf{I}_{r}&r=c\\\
(\mathbf{I}_{c},0_{c,r-c})^{T}&r>c\end{cases}\in\mathbb{R}^{r\times c}$

######  Definition 2 (Hadamard matrix).

Define $\mathbf{H}_{0}:=[1]\in\mathbb{R}^{1\times 1}$ to be the $0^{th}$
Hadamard matrix. Then, for all $m\in\mathbb{N}$, we define the $n^{th}$
Hadamard matrix recursively by

$\mathbf{H}_{m}:=\begin{bmatrix}\mathbf{H}_{m-1}&\mathbf{H}_{m-1}\\\
\mathbf{H}_{m-1}&-\mathbf{H}_{m-1}\end{bmatrix}\in\mathbb{R}^{2^{m}\times
2^{m}}$

Now, the ZerO initialization algorithm for a multilayer perceptron is as
follows (the algorithm is very similar for a convolutional layer, except that
we treat the kernel as a $2\times 2$ matrix by only initializing the
diagonal):

Algorithm 1 ZerO initialization for MLPs.

A MLP with depth $L$ and layer widths $n_{0},...,n_{L}$, where $n_{0}$ is the
input dimension and $n_{L}$ is the output dimension. Let
$W^{(l)}\in\mathbb{R}^{n_{l}\times n_{l+1}}$ denote the $l^{th}$ weight matrix
for $0\leq l<L$.

for $l\in\\{0,...,L-1\\}$ do

if $n_{l}=n_{l+1}$ then

$W^{(l)}\leftarrow I_{n_{l}}$

else if $n_{l}<n_{l+1}$ then

$W^{(l)}\leftarrow I_{n_{l},n_{l+1}}^{*}$

else if $n_{l}>n_{l+1}$ then

$W^{(l)}\leftarrow cI_{n_{l},2^{m}}^{*}H_{m}I_{2^{m},n_{l+1}}^{*}$, where
$m:=\lceil\log_{2}(n_{l})\rceil$ and $c:=2^{-(m-1)/2}$

end if

end for

It can be shown (see Lemma 1 in zhao2022zero ) that the weight matrices
$W^{(l)}$ that we attain has a bounded rank on initialization. Furthermore, it
is empirically shown that the weight matrices that are learned from this
initialization monotonically increase in rank to their final state. This means
that this initialization is well suited to learning low-rank transformations,
which makes it conducive to sparisification via magnitude-based pruning (we
expect a low rank matrix to have a few large elements along several
directions).

#### Learning a Sparsity Mask

In addition to the initialization presented above, we will also learn a
sparsity mask throughout the pre-pruning stage. Precisely, we keep track of
parameters $\theta$ of the network and a mask $M\in\\{0,1\\}^{|\theta|}$. We
interpret each element $M_{i}$ of $M$ to be a Bernoulli random variable
determining whether the parameter $\theta_{i}$ will be used during each
forward pass during training; in other words, $1-M_{i}$ is the probability
that $\theta_{i}$ is zeroed out during the forward pass. This will assign a
sparse structure to the model immediately, and (if we allow gradients to
backpropagate through the masking procedure) training of the weights $\theta$
will proceed as normally. However, backpropagation through this network will
also train the mask $M$ in such a way that minimizes the loss. If we apply
proper regularization, we can promote masks $M$ that have the two desired
properties of sparse masks.

More thoroughly, let $\ell(\widehat{y},y)$ be the differentiable loss function
that measures performance. Initialize $M$ to be $1/2$ in all entries. Define
the regularized loss at a data point $(x,y)$ to be

$L_{reg}(x,y):=\ell(f(x,\theta\odot
M),y)+\lambda_{1}\sum_{i=1}^{|\theta|}M_{i}(1-M_{i})+\lambda_{2}\sum_{i=1}^{|\theta|}M_{i},$

where $f(x,\theta)$ is the output of the network for input $x$ with parameters
$\theta$, $M$ is the mask, and $\lambda_{1},\lambda_{2}$ are regularization
hyperparameters. Note that the $\lambda_{1}$ term is minimized when $M$ is
composed of only 0’s and 1’s; this causes the optimization to prefer bimodal
masks that are confident. The $\lambda_{2}$ term forces the optimization to
prefer masks with more 0’s and less 1’s. Together, the added regularization
terms enforce property 1 of sparse masks: that they are sparse. Meanwhile, the
$\ell(f(x,\theta\odot M),y)$ term requires that the masked model’s inference
performs well.

This scheme to learn a sparsity mask will ensure that the mask befits the
model’s weights, as they were jointly optimized. After pre-pruning training,
we set $\theta\leftarrow\theta\odot M$, such that magnitude-based pruning
preferentially selects along the mask $M$ that we just learned.

###  3.2 Post-Pruning Training

After pruning, we now have a sparse model $f(x,\theta\odot M)$ that we wish to
train to convergence (by this, we mean that only the elements
$\\{\theta_{i}:M_{i}\neq 0\\}$ get trained). We apply several optimizations to
the post-pruning training scheme to improve the training of this model. As one
of the main difficulties of training sparse models lies in handling
discontinuity and sharp transitions, most of these techniques encourage
initial smoothness of the optimization trajectory.

#### Label Smoothing

Label smoothing is a well-known technique that modifies the target labels by
interpolating them with the uniform label. More specifically, let $K$ be the
number of classes and let $\alpha$ be the label smoothing parameter. Then, for
a given target label $y$, we produce the modified target label

$\widetilde{y}:=(1-\alpha)\cdot y+\alpha\cdot\frac{1}{K}$

This procedure smooths the labels (hence the name) against the uniform
distribution. When cross-entropy is applied on top of this technique, the
training optimization prefers smoother logit prediction as opposed to
overconfident guesses.

Motivated by this technique, and following prior works, we apply a decayed
label smoothing to the post-pruning training stage. By this, I mean that given
hyperparameters $\alpha_{0}$ and a decay time $T_{\alpha}$ (measured in
epochs), we set $\alpha$ for each epoch $t\in\\{0,...,T_{\alpha}\\}$ to be

$\alpha_{t}:=\alpha_{0}\left(1-\frac{t}{T_{\alpha}}\right)$

As such, $\alpha$ decays linearly from $\alpha_{0}$ to 0 in $T_{\alpha}$
epochs. This ensures that at the beginning of the post-pruning stage, the
optimization trajectory begins smoother but eventually behaves normally (i.e.
after $T_{\alpha}$ epochs, the loss is computed without any label smoothing).

#### Soft Activations

Following the work of jaiswal2022training  and others, we make use of a
smoothed version of the ReLU activation in place of ReLU. This is because the
gradient discontinuity of ReLU at 0 makes the optimization trajectories of
sparse models sharp and tricky. To combat this, we apply a similar trick as we
did in label smoothing. In particular, we define the parameterized Swish
activation $\verb|PSwish|$ to be given by

$\verb|PSwish|(t):=t\cdot\verb|sigmoid|(\beta t)=\frac{t}{1+e^{-\beta t}}$

for a temperature parameter $\beta$. Note that when $\beta=0$ we have the
(rescaled) identity, but as $\beta\rightarrow\infty$ we approach ReLU. This
"soft activation" has continuous gradients and induces a smoother optimization
trajectory, and so we would like to initially soften the activation and then
increase $\beta$ during training. To this end, we define an initial
$\beta_{0}$ and decay time $T_{\beta}$ (measured in epochs) as
hyperparameters, and for each epoch $t\in\\{0,...,T_{\beta}\\}$ we set

$\beta_{t}:=\beta_{0}\left(1-\frac{t}{T_{\beta}}\right)+\beta_{max}\frac{t}{T_{\beta}}$

#### Soft Skip Connections

Another point of failure of sparse network training in the initial phase is
that, because of the sparse topology of the network, information propagation
through the forward and backward passes can be halted or forced to pass
through weird paths. Since we have removed a majority of the possible paths
for gradients to flow, we would like to create some artificial ones in the
beginning of the post-pruning training stage. To accomplish this, we create
skip connections, as is done in jaiswal2022training . Between each constituent
block of the network, we add a weighted residual connection, parameterized by
$\gamma$. More precisely, if $z$ is the intermediate tensor being passed into
a block $g(\cdot;\theta_{g})$, then we modify the output of the block to be

$\widetilde{g}(z,\theta_{g}):=g(z,\theta_{g})+\gamma\cdot z$

(If the dimensions of $g(z,\theta_{g})$ and $z$ do not add up, we simply
resize and apply interpolation). As always, we introduce hyperparameters
$\gamma_{0},T_{\gamma}$ and for each epoch $t\in\\{0,...,T\\}$ we define

$\gamma_{t}:=\gamma_{0}\left(1-\frac{t}{T_{\gamma}}\right)$

and use $\gamma_{t}$ as the parameter for that epoch. This allows the network
to initially optimize as though it had many skip connections and allow
gradients to flow reasonably. However, later on in training (after
$T_{\gamma}$ epochs), we recover the intital topology of the model so that we
don’t sacrifice final performance.

##  4 Experiments

###  4.1 Settings

Baselines. We compare the proposed methodology against training-free pruning
pipelines (i.e. random, SNIP, GraSP, and SynFlow snip ; grasp ; synflow ), as
well as against RigL, a state of the art sparse training pipeline that does
not use a pruning phase. For RigL, we use the suggested hyperparameters.

Datasets. We evaluate our proposed pipeline in two multi-class classification
settings:

  1. 1.

on the MNIST dataset with a fully-connected 6-layer MLP with 100 hidden units
per layer

  2. 2.

on the CIFAR-10 dataset with a VGG-16 model

Hyperparameters. We set the following training hyperparameters for all
experiments:

  * •

`batch_size` = 256

  * •

`learning_rate` = 0.001

  * •

`weight_decay` = $1\times 10^{-6}$

  * •

$\verb|num_epochs|=\begin{cases}20&\text{MNIST}\\\
100&\text{CIFAR-10}\end{cases}$

  * •

We use the Adam optimizer.

Furthermore, for the hyperparameters that are introducted in our methodology,
for the MNIST experiments we set:

  * •

$\lambda_{1}=0.001$

  * •

$\alpha_{0}=N/A$

  * •

$\beta_{0}=N/A$

  * •

$\gamma_{0}=N/A$

  * •

$\lambda_{2}=0.05$

  * •

$T_{\alpha}=N/A$

  * •

$T_{\beta}=N/A$

  * •

$T_{\gamma}=N/A$

and for the CIFAR-10 experiments we set:

  * •

$\lambda_{1}=0.01$

  * •

$\alpha_{0}=0.3$

  * •

$\beta_{0}=1.0$

  * •

$\gamma_{0}=0.4$

  * •

$\lambda_{2}=0.1$

  * •

$T_{\alpha}=50$

  * •

$T_{\beta}=60$

  * •

$T_{\gamma}=60$

The above hyperparameters were determined on a validation dataset. (By $N/A$
for the MNIST experiments, we mean that we saw no observable differences when
turning on label smoothing, soft activations, nor soft skip connections, and
so we leave them off for the experiments below).

###  4.2 Results

In the plots that follow, the following information is relevant:

  * •

$s$ is the sparsity parameter that dictates the proportion of parameters that
is left nonzero after sparsification. The related "compression" value, which
is sometimes used to parameterize this, is related by
$compression=-\log_{10}(s)$.

  * •

$\verb|training-free|$ denotes the result of the pruning pipeline where the
best training-free pruning technique was used, and then the resulting sparse
model was trained for `num_epochs` epochs; in other words, this plots the
corresponding values for
$\max\\{\verb|random|,\verb|SNIP|,\verb|GraSP|,\verb|SynFlow|\\}$ for each
experiment (it was usually SNIP).

  * •

`RigL` denotes the result of training for $\verb|num_epochs|$ using the RigL
training scheme with the desired sparsity $s$.

  * •

All other runs followed the procedure of

train for num_epochs $\rightarrow$ magnitude-based pruning $\rightarrow$ train
for num_epochs

with the combination of techniques given in the legend.

  * •

We use `ToST` to abbreviate all of the post-pruning training optimizations,
since we only tried turning them on together.

  * •

Dashed lines are used to indicate baselines.

  * •

`vanilla` is used to indicate vanilla magnitude pruning.

Figure 1: Top-1 accuracies on the MNIST dataset using the MLP architecture. We see that our proposed methods offer improvement over the baselines, with the learned mask having the largest impact. Furthermore, the addition of the post-pruning training optimizations seems to help when applied, but very barely. Figure 2: Top-1 accuracies on the CIFAR-10 dataset using the VGG-16 architecture. We see that our proposed methods offer improvement over the baselines, with the learned mask having the largest impact. Furthermore, the addition of the post-pruning training optimizations seems to help when applied, but very barely. | $s=10\%$ | $s=3.2\%$ | $s=1\%$ | $s=0.32\%$
---|---|---|---|---
`training-free` | 97.3 | 93.2 | 57.0 | 11.4
`vanilla` | 88.9 | 75.2 | 11.4 | 11.4
`RigL` | 97.1 | 92.8 | 40.1 | 11.4
`+ ZerO` | 97.2 | 97.1 | 66.2 | 17.9
`+ Mask` | 97.0 | 96.5 | 88.9 | 73.5
`+ ZerO, Mask` | 97.2 | 87.8 | 89.0 | 65.0
Table 1: Reported Top-1 Accuracies of MLP architecture trained on the MNIST dataset. We find that the ZerO initialization and learned masks seem to have the largest effects on pruned accuracy for sparsities on the order of 3%, but for anything less it is the learned mask that has the largest impact. | $s=10\%$ | $s=3.2\%$ | $s=1\%$ | $s=0.32\%$ | $s=0.1\%$
---|---|---|---|---|---
`training-free` | 87.3 | 77.7 | 75.8 | 55.2 | 39.7
`vanilla` | 87.1 | 34.0 | 10.0 | 10.0 | 10.0
`RigL` | 86.5 | 82.4 | 64.2 | 44.7 | 29.5
`+ ZerO` | 87.1 | 86.7 | 82.1 | 57.4 | 27.6
`+ Mask` | 87.5 | 86.0 | 81.1 | 59.1 | 45.7
`+ ToST` | 87.2 | 39.1 | 10.0 | 10.0 | 10.0
`+ ZerO, ToST` | 86.9 | 86.8 | 83.1 | 59.4 | 27.0
`+ Mask, ToST` | 87.1 | 86.0 | 83.1 | 60.1 | 42.2
`+ All` | 86.8 | 86.2 | 72.4 | 39.6 | 22.8
Table 2: Reported Top-1 Accuracies of VGG-16 architecture trained on the
CIFAR-10 dataset. We find that the ZerO initialization and learned masks seem
to have the largest effects on pruned accuracy.

The above results demonstrate several things. Firstly, our combined approach
does significantly better than the baselines when the models become very
sparse. Indeed, we are able to retain full performance with only 3% of the
parameters unpruned, which is better than both the training-free pruning
baselines and RigL. At $\leq 1\%$ of the parameters left, our methodology
seems to dominate.

There are, however, some interesting and worrisome trends to note. Firstly, it
appears the ZerO initialization and learned masks don’t play well with each
other; we discuss this in Section 5. Furthermore, we find that the gap between
our pipeline and the baselines is much larger in MNIST than in CIFAR-10.
Perhaps things look quite different if we continue to scale the experiments;
this would be interesting to investigate. Lastly, the combination of label
smoothing, soft activations, and soft skip connections seems to have a very
small positive effect, smaller even than what is reported in
jaiswal2022training  (though they weren’t testing at such small sparsity
levels).

##  5 Discussion

Our combined methodology appears to beat the well-established previous methods
and also a fairly new state of the art in very sparse settings. Interestingly,
it appears that much of the heavy lifting is done by the learnable masks, and
so further research in this direction seems fitting.

As mentioned earlier, the results indicate that the ZerO initialization is
combined with the learnable mask do not combine in a way that improves final
performance. Upon further examination, this is indeed the case for the
following reason: on initialization and early in the pre-pruning training
stage, the gradients of the cross-entropy with respect to the mask are wildly
undirected. As such, the $\lambda_{2}$ regularization term that attempts to
sparsify the mask sends the entire mask to 0 and training collapses
irreparably; this happens for any $\lambda_{2}>0$. If we run the same
experiment but instead set $\lambda_{2}=0$, training no longer collapses, but
instead the learned mask is all 1’s and we recover the same pruned performance
that ZerO accomplishes on its own. Further work in investigating the interplay
between these two methods seems promising, as they both individually offer
fantastic improvement and appear to operate in complementary ways. Perhaps the
scheme with which we initialize the mask could be modified to fit these two
techniques together?

Lastly, the results for very small sparsities are tantalizing; for a sense of
scale, recall that 1% sparsity of VGG-16 is 147k parameters remaining out of
the original 14.7M. All of the works cited in this paper do not attempt
sparsities this small, and so it is useful knowledge just to understand how
pruning looks at this scale. That being said, it would be very interesting to
know more about how methods such as the proposed pipeline scale to larger
models and more complex settings. Operating with such small sparsities isn’t
all that useful for VGG-16 (the difference between 1 and 3% is not much), but
such sparsities would be crucially important if they could be usefully reached
in much larger models. In particular, it is important to see if the same
techniques apply to large language models and large models that mix
convolution and attention, since this can be where the pruning of parameters
finds biggest impact. Furthermore, in the ZerO initialization paper
zhao2022zero  the author claims to have reached similar results for the
initialization of attention layers, though I have not been able to replicate
this nor find an implementation. If the results of our methodology here carry
over to attention layers (which have inference time $n^{2}$), the inference
footprint rewards of our increased training efforts could be very large.

##  6 Acknowledgements

Thank you so much to Professor Kai Li and to Yushan for a fantastic course! I
loved the opportunity to learn about a topic I didn’t know all that much
about, and the paper selection was incredible. Thank you :)

## References

  * [1] Jonathan Frankle and Michael Carbin.  The lottery ticket hypothesis: Finding sparse, trainable neural networks, 2019.
  * [2] Tailin Liang, John Glossner, Lei Wang, Shaobo Shi, and Xiaotong Zhang.  Pruning and quantization for deep neural network acceleration: A survey, 2021.
  * [3] Michela Paganini and Jessica Forde.  On iterative neural network pruning, reinitialization, and the similarity of masks, 2020.
  * [4] Namhoon Lee, Thalaiyasingam Ajanthan, and Philip H. S. Torr.  Snip: Single-shot network pruning based on connection sensitivity, 2019\.
  * [5] Chaoqi Wang, Guodong Zhang, and Roger Grosse.  Picking winning tickets before training by preserving gradient flow, 2020\.
  * [6] Hidenori Tanaka, Daniel Kunin, Daniel L. K. Yamins, and Surya Ganguli.  Pruning neural networks without any data by iteratively conserving synaptic flow, 2020.
  * [7] Jiawei Zhao, Florian Schäfer, and Anima Anandkumar.  Zero initialization: Initializing neural networks with only zeros and ones, 2022.
  * [8] Suraj Srinivas, Akshayvarun Subramanya, and R. Venkatesh Babu.  Training sparse neural networks, 2016.
  * [9] Ajay Jaiswal, Haoyu Ma, Tianlong Chen, Ying Ding, and Zhangyang Wang.  Training your sparse neural network better with any mask, 2022.
  * [10] Utku Evci, Trevor Gale, Jacob Menick, Pablo Samuel Castro, and Erich Elsen.  Rigging the lottery: Making all tickets winners, 2021.