#  Depth-Wise Attention (DWAtt):
A Layer Fusion Method for Data-Efficient Classification

Muhammad ElNokrashy,, 1 Badr AlKhamissi,111In this work, symbols in  red have
learned parameters., 2 Mona Diab 2 Equal Contribution and Correspondence.Work
started independently, and continued during residency at Meta AI.
<EMAIL_ADDRESS>
###### Abstract

Language Models pretrained on large textual data have been shown to encode
different types of knowledge simultaneously. Traditionally, only the features
from the last layer are used when adapting to new tasks or data. We put
forward that, when using or finetuning deep pretrained models, intermediate
layer features that may be relevant to the downstream task are buried too deep
to be used efficiently in terms of needed samples or steps. To test this, we
propose a new layer fusion method: Depth-Wise Attention (DWAtt), to help re-
surface signals from non-final layers. We compare DWAtt to a basic
concatenation-based layer fusion method (Concat), and compare both to a deeper
model baseline—all kept within a similar parameter budget. Our findings show
that DWAtt and Concat are more step- and sample-efficient than the baseline,
especially in the few-shot setting. DWAtt outperforms Concat on larger data
sizes. On CoNLL-03 NER, layer fusion shows $3.68-9.73\%$ F1 gain at different
few-shot sizes. The layer fusion models presented significantly outperform the
baseline in various training scenarios with different data sizes,
architectures, and training constraints.

##  1 Introduction

The Transformer architecture (Vaswani ., 2017) and variants (Fan ., 2020) have
become a mainstream, reliable choice for a wide range of Natural Language
Processing (NLP) tasks, such as sentence classification, token labeling,
retrieval, and question answering (A. Wang ., 2018, 2019). This is in part due
to architectural properties that enable enhanced parallelization and better
modeling of long-range dependencies. Several models have since surfaced,
setting new records on different benchmarks (Devlin ., 2019; Lewis ., 2020;
Brown ., 2020). In the vanilla architecture, a stack of Transformer blocks are
applied sequentially to refine the representation of an input sequence, which
is then fed to a task-specific module, such as a classifier head.

Recent works have shown that the hidden representations from intermediate
layers may benefit downstream tasks (Wallat ., 2021). Other works have tested
the fusion of hidden representations in tasks such as sequence-to-sequence
machine translation (Shen ., 2018; F. Liu ., 2020; X. Liu ., 2021).

In this work, we propose a method to combine the hidden representations of
encoder layers to more easily utilize the full model. Moreover, we further
investigate a simpler alternative as another compelling baseline. Some of the
inspiration for this work is based on Vertical Attention introduced in
AlKhamissi . (2021).

Figure 1: Basic architecture: The Mixer $H$ can be a Sum of Affine
Transformations, a Dot-Product Attention module, etc. Different variants may
define $K_{i}$, $V_{i}$, and $Q$ (key, value, and query transforms) and
utilize them differently.

##  2 Motivation & Hypothesis

A common goal when designing deeper networks is increasing their ability to
represent longer chains of abstraction in different data domains and training
objectives. As the model optimizes for a specific task, seemingly _unneeded_
information is ignored. More general patterns that could benefit different
downstream tasks become less likely to pass through to later layers and to the
model’s final representation.

Some properties of such base models (like large language models) may
implicitly mitigate this effect—by leveraging the larger parameter capacity
and the increased width. The additional parameters enable memorizing more
subtleties or idiosyncrasies of the training data, while the increased model
width allows for a less compact final representation, which may let through
patterns not _directly_ used by the general language modeling objective of
choice. Some tasks may need such low-level knowledge across a large example
space (X. Liu ., 2021). A similar problem is the strength and clarity of
gradients into earlier intermediate layers. Methods to alleviate this include
skip connections (He ., 2016), and alternatives to full back-propagation
(Nøkland, 2016).

In this spirit, we propose an add-on module that can be added to any
pretrained deep sequence model to combine the representations of its
intermediate layers to adapt better to novel tasks. To investigate our
proposed modules, we experiment with the following tasks: (a) Named Entity
Recognition (NER) in the few-shot setting on the CoNLL-03 and WikiAnn
datasets, (b) Masked Language Modeling (MLM) on a small text dataset
(WikiText-2) in two settings: Finetuning (FT) and Feature Extraction (FE).

#### Adapting to New Tasks.

For many pretraining projects, an important performance signal is often its
ability to adapt to novel language tasks, with the intent to exploit general
language understanding from general self-supervised tasks like MLM on large
datasets.

Metrics in this case can include: (a) Performance compared to a model trained
from scratch on task data, (b) the time to convergence, and (c) the size of
needed task training data to achieve the desired performance.

One may finetune the whole model (possibly infeasible), or use it as a feature
extractor and train only a task-specific module. A commonly researched
alternative is finding or introducing a subset of parameters specifically for
finetuning (Houlsby ., 2019; Ben-Zaken ., 2022). We test our proposed method
as well as another highlighted method in different adaptation settings with
multiple n-shot training sizes, max training steps, and tiers of availability
of language data during pretraining.

Figure 2: F1-Score on the CoNLL-03 dev set after training for $100$ epochs on
few-shot training datasets with N-Shot samples per class sampled uniformly
from the full training set. Both DWAtt and Concat improve on the additional
layers baseline R26. At $128$ samples per class, DWAtt improves on both Concat
and the enhanced baseline by $1.6\%$ and $5.28\%$ absolute, respectively.

##  3 Tasks, Datasets, and Raised Questions

On many common benchmarks, state-of-the-art performance is often near
saturation. We consider some useful synthesizable variants of the benchmarks
that test particular training settings or aspects of performance. Our
experiments aim to test the following aspects: adaptability (finetuning (FT)
versus feature extraction (FE)), sample efficiency (few-shot training),
training step efficiency (time to convergence), and effect of model depth.

###  3.1 Few-Shot Adaptation on NER

Dataset | Classes | Resources | Train | Dev
---|---|---|---|---
CoNLL-03 | 4 | High | 14k | 3250
WikiAnn | 3 | High | 20k | 10k
Mid | 1k–5k | 1k
Low | 100 | 100
Table 1: Train and Dev subset sizes in sentence count. For WikiAnn, we list
the average size of a language in the corresponding resource tier, determined
by train size.

#### CoNLL-2003

The CoNLL-03 dataset (Tjong Kim Sang De Meulder, 2003) provides a mainstream
NER benchmark in the English language. For few-shot trials, we sample training
points uniformly at random from the full training set.

#### WikiAnn

The WikiAnn dataset (Rahimi ., 2019) is a multilingual NER benchmark which we
use to test the effect of resource availability in the pretraining phase. For
few-shot experiments, we sample a fixed training set of size $N=100$ and train
for each language separately. We also compare performance on the English
subset between RoBERTa and the multilingual XLM-RoBERTa.

###  3.2 Masked Language Modeling

The WikiText-2 dataset (Merity ., 2017) is a subset of the English Wikipedia
developed for use in long-term dependency language modeling. It contains the
first $2$ million words of a dump of the English Wikipedia at the time of
creation. We report the perplexity on the dev and test sets.

##  4 Models

Let $L$ denote a deep network’s layer stack. Then
${\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}{H}}$111In
this work, symbols in  red have learned parameters. is a learned function that
mixes the layers’ intermediate representations $\\{{\mathbf{z}}_{n}\mid
n\in\lvert L\rvert\\}$, and ${{\mathbf{h}}}$ is the final representation.

${{\mathbf{h}}}={\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}{H}}~{}(\cdots,\big{\\{}(n,{\mathbf{z}}_{n})\mid
n\in\lvert L\rvert\big{\\}})$ (1)

The $H$ module is a function of the layer indices $n$ and the corresponding
representation vectors ${\mathbf{z}}_{n}$, but can take other signals, like
${\mathbf{z}}_{L}\mapsto{\mathbf{q}}$ for query. See diagram in Figure 1. We
consider the following layer fusion models.

#### Layer Concatenation.

A sum of linear transforms. Note that this is equivalent to concatenating all
$\\{{\mathbf{z}}_{n}\\}$ then transforming the concatenation into the model
width ${{d}_{\text{\,z}}}$.

${{\mathbf{h}}}={\textstyle\sum_{n}}\
{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}{{\mathbf{W}}_{n}}}({\mathbf{z}}_{n})$
(2)

#### Depth-Wise Attention.

DWAtt uses dot-product attention with keys ${\mathbf{k}}_{n}$, values
${\mathbf{v}}_{n}$, and query ${\mathbf{q}}$.

$\displaystyle{\mathbf{k}}_{n}$
$\displaystyle={\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}{\operatorname{\vphantom{fg}\mathsf{PE}}}}(n)$
(3) $\displaystyle{\mathbf{v}}_{n}$
$\displaystyle={\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}{\operatorname{\vphantom{fg}\mathsf{LN}}_{n}}}\left({\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}{f^{V}_{n}}}({\mathbf{z}}_{n})\right)$
(4) $\displaystyle{\mathbf{q}}$
$\displaystyle=1+\operatorname{\vphantom{fg}\mathsf{elu}}\left({\mathbf{z}}_{L}+{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}{f^{Q}}}({\mathbf{z}}_{L})\right)$
(5)

Where $\operatorname{\vphantom{fg}\mathsf{PE}}(n)$ is a learned positional
embedding vector for layers $\\{n\\}$.
$\operatorname{\vphantom{fg}\mathsf{LN}}$ is LayerNorm (Ba ., 2016). $f^{Q}$,
$f^{V}_{n}$ are MLPs with a bottleneck at $1/2$ the model width
${{d}_{\text{\,z}}}$. $\operatorname{\vphantom{fg}\mathsf{elu}}$ is from
Clevert . (2015).

$f({\mathbf{z}})={\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}{{\mathbf{W}}}}\cdot{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}{\operatorname*{\vphantom{fg}\mathsf{LN}}}}(\operatorname{\vphantom{fg}\mathsf{gelu}}({\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}{{\mathbf{U}}}}{\mathbf{z}}))$
(6)

It is then comparable to a single linear layer of size
${{d}_{\text{\,z}}}\times{{d}_{\text{\,z}}}$. See Section 9 for details. Note:
the scoring step in $\operatorname{\vphantom{fg}\mathsf{Attend}}$ reduces to a
single, vector-by-static matrix multiplication as the keys
$\\{{\mathbf{k}}_{n}\\}$ need no input. Assume a time step $t$, then let
${\mathbf{K}}=\\{{\mathbf{k}}_{n}\\}$ and
${\mathbf{V}}=\\{{\mathbf{v}}_{n}\\}$ be the matrix forms of the keys and
values for $n\in\lvert L\rvert$. Then:

$\displaystyle\operatorname*{\vphantom{fg}\mathsf{Score}}({\mathbf{q}},{\mathbf{K}})$
$\displaystyle=\operatorname{\vphantom{fg}\mathsf{softmax}}_{n}\left({\mathbf{q}}\cdot{\mathbf{K}}^{\top}\right),$
(7) $\displaystyle\operatorname*{\vphantom{fg}\mathsf{Attend}}(\cdots)$
$\displaystyle=\operatorname{\vphantom{fg}\mathsf{Score}}({\mathbf{q}},{\mathbf{K}})\cdot{\mathbf{V}},$
(8) $\displaystyle{{\mathbf{h}}}$
$\displaystyle={\mathbf{z}}_{L}+\operatorname{\vphantom{fg}\mathsf{Attend}}\big{(}{\mathbf{q}},\\{{\mathbf{k}}_{n}\\},\\{{\mathbf{v}}_{n}\\}\big{)}.$
(9)

###  4.1 Baselines

Name | New Layer | \+ $O(\cdot)$ | \+ #
---|---|---|---
Base | - | - | -
Base+n |  $n\times$ Base | $7n\ (d+d^{2})$ |  $25.19$M
Average | MeanPool | - | -
Concat | Affine | $Ld^{2}$ |  $25.18$M
DWAtt | DWAtt | $Ld^{2}+d^{2}+Ld$ |  $26.38$M
Table 2: The size and count of parameters of new layers in different
configurations. Counts are of models building on RoBERTa24. Base+n (deeper
baseline) is chosen to be close in size to DWAtt and Concat. For example, for
base models with 12 layers, only n=1 layers are added; n=2 for 24 layer
models. $L$ refers to the depth of the _base_ model, while $d$ refers to the
model width.

#### Base Model.

The base model (RoBERTa or XLM-RoBERTa in base or large sizes) is used as-is.
The task module is changed where needed. In FE only the task module (or
vocabulary module in MLM) is trained.

#### Extra Transformer Layers.

On top of the base model, we add $2$ more Transformer layers before the
classification head. In FE mode, we train only the added layers. We refer to
this model as R26 (RoBERTa) or XLM-R26 (XLM-RoBERTa).

##  5 Experimentation Setup

###  5.1 Base Models

Monolingual experiments using CoNLL-03, the English subset of WikiAnn, and
WikiText-2 all build on RoBERTa (R) (Y. Liu ., 2019) via the pretrained
roberta-large model. Multilingual experiments using WikiAnn build on XLM-
RoBERTa (XLM-R) via the pretrained xlm-roberta-large model. Experiments using
other variants are explicitly specified. Pretrained models are accessed via
the HuggingFace package (Wolf ., 2019).

###  5.2 Training

In figures where the epoch count is reported: Each chart shows an experiment
with the indicated ${N_{\text{epochs}}}$ max epochs, which starts LR at max
then decays it linearly to zero.

Few-shot experiments sample $NC$ points uniformly at random from the full
training set. Sampling is not stratified, so $N=8$ Shot for $C=4$ classes
means $32$ points in total sampled without replacement.

###  5.3 Evaluation

NER experiments report Micro-Average F1 using seqeval (Nakayama, 2018). Where
applicable, experiments are given 5 trials whose average and confidence
interval are reported by seaborn (Waskom, 2021). For easier reading, we report
scores that lie in $[0,1]$ as percentages $[0,100]\%$.

##  6 Results and Analysis

###  6.1 CoNLL: Few-Shot Adaptation

Figure 3: F1-Score on the CoNLL-03 devset. All pretrained weights are frozen
(FE). In each chart from left to right, training is constrained to $25$, $50$,
$75$, and $100$ max epochs. For each N-Shot experiment, $NC$ samples ($C=4$
classes) are randomly selected and trained on for the entire experiment. The
scores are averaged across $5$ trials with random initialization of weights
and data sampling. We report the best observed dev score from the full
training of each experiment and trial.

Micro-F1 is reported at each N-Shot in $\\{8,16,32,64,128\\}$. Each model in
Figure 3 adds the specified module on top of a pretrained RoBERTaLARGE. Concat
(orange) noticeably outperforms R26 at all few-shot settings, while DWAtt
(blue) outperforms Concat at the higher data sizes.

###  6.2 Step and Sample Efficiency

(a) Few-shot at $N=8$.

(b) Few-shot at $N=128$.

Figure 4: Training Behavior. Validation F1-Score across steps in FE training
on CoNLL-03. Concat generalizes more readily at smaller $N$ compared to DWAtt.
At both sizes, layer fusion methods are able to extract more from pretrained
models than traditional last-layer fitting.

In Figure 3, Concat outperforms noticeably at the lowest end ($N{=}8$ at
${N_{\text{epochs}}}{=}25$), improving on R26 by $\mathbf{58\%}$—possibly
owing to its simpler connectivity and gradient path. While DWAtt’s improvement
given increased data and training time may be explained by its being more
selective due to the attention module. In the largest setup ($N{=}128$,
${N_{\text{epochs}}}{=}100$; Figure 2), Concat improves on R26 by $3.68\%$,
while DWAtt improves on it by $\mathbf{5.28\%}$. R26 demonstrates the
difficulties of extracting the patterns needed by a downstream task from a
feature extractor that has already fit its last layer representation ($L{=}24$
for RoBERTaLARGE) for its pretraining task.

In Figure 4, we focus on model validation behavior during training with
${N_{\text{epochs}}}{=}50$. At $N{=}8$, only Concat manages to adapt well, and
rapidly, converging in $30\%$ of training time. At $N{=}128$, DWAtt has enough
data to reach similar performance at a similar pace, and quickly pulls ahead.

We repeat in Figure 8 the experiments from Figure 3 but on English only from
WikiAnn, on R26. Yet again Concat performs better with less data, while DWAtt
is better at the higher end ($N{=}128$, ${N_{\text{epochs}}}{=}100$) by
$+2.08\%$ and $+0.92\%$ over Concat and R26, respectively.

###  6.3 Feature Extractor Adaptability

Figure 5 shows the effect of finetuning (FT) all model parameters for each of
the same configurations used in other CoNLL-03 experiments. At $N=8$, Concat
in feature extraction (FE) training already gets most of the improvement
observable from full FT across the board. At $N=128$ the effect is stronger.
While layer fusion methods help close the gap—with DWAtt giving nearly half
the improvement of FT using only FE, at $+5.28\%$—FT still manages overshadow
all FE configurations even with only R26.

(a) $N=8$

(b) $N=128$

Figure 5: Feature Extraction (_FE_) versus Finetuning (_FT_). Best validation
F1-Score on CoNLL-03 within ${N_{\text{epochs}}}=50$ training epochs.
Finetuning consistently outperforms alternatives, but Concat and DWAtt
approach its performance even in FE training at lower data sizes.

###  6.4 Scaling by Depth

In Figure 6, we compare the same add-on configurations on the base and large
variants of RoBERTa. At $N=8$, Concat is an outlier in its performance gain
over the alternatives. In all $N\in\\{8,128\\}$, both Concat and DWAtt match
or beat the performance gain observed from RoBERTa+222RoBERTa+ base is
$L{=}12+1$ layers, while large is $L{=}24$+2. when changing from base to
large.

(a) $N=8$

(b) $N=128$

Figure 6: base versus large Pretrained Models. Validation F1-Score on CoNLL-03
on the DWAtt, Concat, and enhanced (+layers) configurations of RoBERTa base
and large. (a) At $N\\!=\\!8$, ConcatBASE shows a clear lead even on
RoBERTa+BASE. (b) At $N\\!=\\!128$, ConcatBASE outperforms RoBERTa+LARGE,
while DWAttLARGE outperforms ConcatLARGE. (c) This also supports the claim for
higher data and training efficiency from Concat and DWAtt in FE training
compared to traditional last-layer fitting.

###  6.5 WikiAnn: Adapting to High, Mid, and Low Resource Languages

Figure 7 shows average performance grouped by the token count of pretraining
data according to Conneau . (2019). Low-resource languages had $10$–$100$M
tokens, Mid-resource languages had $200$–$300$M tokens , while High-resource
languages had $2$–$20$G tokens. While there is a slight lead for layer fusion
methods in the Low-resource column, the difference does not translate to a
similar lead in Mid-resource scores over High-resource.

Compare the scores in both Figures 8 & 9. Starting from
${N_{\text{epochs}}}{=}25$ and only $N{=}32$ samples, performance using
RoBERTaLARGE as the base model beats the corresponding run on XLM-RLARGE at
${N_{\text{epochs}}}{=}25$ with $100$ samples.

Figure 7: F1-Score on the WikiAnn devset. Languages are grouped into three
tiers of resource size. See Figure 9 and Table 1 for details.

###  6.6 Masked Language Modeling

Model | Test Perplexity
---|---
RoBERTa+ | 2.98
Concat | 3.44
DWAtt | 2.91
Base | 2.90
Table 3: MLM perplexity on WikiText-2 test set.

inline]

  * •

Training loss

  * •

Dev loss

  * •

Final Test loss

We report the test perplexity of the checkpoint with the best validation loss
from FT training. DWAtt performs closer to the baseline Base than Concat and
R26.

Figure 8: F1-Score on the WikiAnn _English_ devset on the _RoBERTa LARGE_ base
model. At  Figure 9: F1-Score on the WikiAnn devset. Each language was trained
for ${N_{\text{epochs}}}=25$ and tested separately. Each model was initialized
from the pretrained XLM-RLARGE weights, augmented with the specified module
(DWAtt, Concat, or additional layers), then trained on exactly $100$
uniformly-sampled training points for each language. The languages are sorted
by token count of pretraining data according to Conneau . (2019).

##  7 Related Work

#### Adaptors.

Adaptors are layers introduced into a pretrained model layer stack and
finetuned on a specific downstream task, to avoid changing the (larger)
model’s weights. (Houlsby ., 2019), among others.

#### Layer Aggregation (static weights).

Bapna . (2018) defines, for each decoder layer, a trainable softmax-normalized
vector of weights to get the weighted sum of the encoder intermediates. X. Liu
. (2021) provides one static, learned vector for each intermediate layer’s
representation which acts as element-wise scaling. Both works are similar to
Concat in applying a static, learned mixing transform to all layers. For each
layer: The first provides a scalar weight, the second provides a vector for
element-wise multiplication, while Concat applies a linear transform.

In Shen . (2018), the DenseNMT is an encoder-decoder NMT architecture densely
connected in the style of DenseNet (Huang ., 2016). Each encoder layer takes a
concatenation of all previous layer representations. Similarly for decoder
layers. See also Q. Wang . (2019). These methods compare structurally to
Concat, which is applied only once on the full layer stack.

Q. Wang . (2018) generates the weight for an encoder layer via an MLP on the
layer’s representation, irrespective of the rest of the model:
${w}_{i}=f^{a}({\mathbf{z}}_{i})$ then
${{\mathbf{h}}}=\sum_{i}({w}_{i}{\mathbf{z}}_{i})$.

#### Dynamic Layer Mixing.

The following methods use signals from the input to change the transformation
itself dynamically. See also DWAtt which uses a Dot-Product Attention module
for comparison. Li . (2018) applies an LSTM depth-wise on the intermediate
vectors of the stack of LSTM cells applied on input sequences, as expected. We
attempted adding an LSTM cell applied depth-wise to the Transformer encoder
stack but observed lacking performance. Note that the referenced work utilizes
a non-basic cell with peep-hole expressions, and some architecture
connectivity that complicates experiments in the Feature Extraction context.
Dou . (2019) utilizes a dynamically-weighted routing mechanism to mix
transformations of each intermediate representation, then concatenate all
such.

##  8 Discussion

###  8.1 Levels of Abstraction

For non-recurrent deep neural networks, there exists a functional limit to the
depth of abstraction obtainable, which is proportional to the depth of the
network. Abstraction here refers to the depth of composed rules that a model
may learn to apply on low-level stimulus, such as pixels or tokens. For feed-
forward models like an MLP or transformers, this corresponds to the depth in
terms of stacked nonlinear layers. As an example, to handle program-like
systematic inputs, depth-recurrent and memory-augmented architecture were
utilized in works such as Dehghani . (2019); Graves . (2014). For traditional
Transformer models, like RoBERTaLARGE, we can say that the limit is
proportional to the number of layers in the network (e.g. $O(|L|)$; $|L|=24$).

#### Hiding Within Scale.

In practice, this is seldom an obvious problem because large networks would
have enough width-wise parameter capacity to directly encode “intuition”, i.e.
shortcuts to knowledge and abstraction. They may do so by tying the low-level
representation of some inputs to intermediate signals for the high-level
concepts they tend to manifest. As an artificial example: A small, shallow
model for classifying sentiment may tie a token such as scary to become a
strong signal for negative sentiment, say in a movie review setting. Sensible
in the domain of kids movies; but may in fact signify a positive sentiment
instead when observed in reviews for horror movies. The key point is that it
is an early shortcut, not whether it is accurate.

#### Intuition as Shortcuts From Raw Input.

Thus, earlier layers can encode information at a higher level than $|L|$-steps
of abstraction would suggest. This would be needed, and expected, for models
of a reasonable depth and width to be able to satisfy the feature extraction
needs of the _pretraining_ task at the last layer. Should a _downstream_ task
require a different signal from what the _pretraining_ task exposed, it may
have to _shift_ a large subset of the weights to relearn or resurface that
information from the shallower levels.

#### Proposition.

The methods we’ve discussed may enable downstream tasks to query the model for
hidden but useful information. For the downstream task to make use of such
features, it would likely need to transform them further. By applying the
$2$-layer MLPs $f^{V}_{n}$ on these intermediate features (in DWAtt), and a
linear transform in Concat, the task can extract a more useful representation
from each level/layer.

###  8.2 Modeling Capacity

The two models presented and highlighted—DWAtt and Concat—are aggregate views
of the features of all intermediate Transformer layers
$\\{{\mathbf{z}}_{n}\mid n\in|L|\\}$ (see Section 4). DWAtt’s added module
requires access to the last layer’s ${\mathbf{z}}_{L}$ to form the query,
while Concat does not. Neither method makes any use of external signals such
as, for example, a task embedding vector.

These three points together present an underlying property of the modeling
capacity of DWAtt and Concat: _The depth-wise layer mixing arrives at a model
that is, at most, as expressive as the underlying sequence-modeler_. See
Figure 5 where layer fusion under FE approaches but does not exceed FT, while
all methods are similar under FT.

#### Modeling Dimensions.

DWAtt and Concat operate _depth-wise_ over a sequence-modeling model. By that
very nature, it may not be the best option when the task at hand requires
increased or improved _spatial_ abstraction—the ability to learn connections
on a spatial axis (e.g. between tokens in sequences in a text Transformer).
Then, adding extra Transformer layers or full finetuning may be better
options.

##  9 Design Details

Model | Param | Value
---|---|---
AllLarge | Transformer Layers | 24
All | Learning Rate | 1e-5
DWAtt |  $\gamma^{Q}$, query bottleneck | 0.5
DWAtt |  $\gamma^{V}$, values bottleneck | 0.5
DWAtt |  ${{{d}_{\text{\,pos}}}}$, keys latent | 24
Table 4: Architecture Parameters

###  9.1 Training

inline,color=red]Update for NER. All trainings use the AdamW (Loshchilov
Hutter, 2017) optimizer with a linear decay learning rate (LR) schedule.
Training on WikiText-2 uses a batch size of $8$ samples, and a max LR of
$5\text{\times}{10}^{-5}$. Training on CoNLL-03 and WikiAnn uses a batch size
of $16$ and a max LR of $5\text{\times}{10}^{-5}$.

###  9.2 Layer Index Embedding

${\mathbf{k}}^{\mathsf{\vphantom{fg}pos}}_{n}\\!\in\mathbb{R}^{{{{d}_{\text{\,pos}}}}}\\!\\!\sim\mathcal{U}(0,1)$
are static positional embedding vectors of the layer index $n$.
${\mathbf{W}}^{K}$ is a learned affine transform.

${\mathbf{k}}_{n}=\operatorname{\vphantom{fg}\mathsf{PE}}(n)={\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}{{\mathbf{W}}^{K}}}{\mathbf{k}}^{\mathsf{\vphantom{fg}pos}}_{n}$
(10)

###  9.3 MLP Modules

For each role $\in\\{Q,V\\}$ that uses an MLP, one is defined for a layer $n$
as:

$\displaystyle f_{n}({\mathbf{x}})$
$\displaystyle={\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}{{\mathbf{W}}_{n}}}\cdot\operatorname*{\vphantom{fg}\mathsf{LN}}\left(\operatorname{\vphantom{fg}\mathsf{gelu}}\left({\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}{{\mathbf{U}}_{n}}}~{}{\mathbf{x}}\right)\right)$
(11)

Where ${\mathbf{U}}_{n},{\mathbf{W}}_{n}$ are the down and up projections of
the bottleneck
(${{d}_{\text{\,z}}}\mapsto\gamma{{d}_{\text{\,z}}}\mapsto{{d}_{\text{\,z}}}$),
respectively. This means, for example, that for value transforms one set of
weights is assigned to each layer $n$ with nothing shared.
$\operatorname{\vphantom{fg}\mathsf{gelu}}$ is the activation function from
Hendrycks Gimpel (2016).

##  10 Conclusion & Future Work

inline,color=yellow]@bkhmsi - Review. We present DWAtt—a new method of reusing
the latent representations of a deep neural network. We analyze DWAtt and a
similar, simpler method—Concat—from multiple aspects of performance and
scaling on NER and MLM tasks. Results suggest similar layer fusion methods can
be a strong tool for downstream adaptation. Performance gains from $1\%$ to
$6\%$ and as high as $30\%$ can be seen in various experiments for different
few-shot sizes and training times, under Finetuning or Feature Extraction, and
on base models of different depths. DWAtt and Concat have shown improved
performance even in Feature Extraction training against full Finetuning. We
believe this effect may extend to other tasks besides sequence labeling and to
other sequence modeling architectures besides the Transformer—and propose such
analysis for future work. We believe _add-on_ -style additions to pretrained
models, such as adapters and depth-wise mixing, to be a fertile ground for
research into low-cost adaptation of large models that has not been saturated
yet.

## References

  * AlKhamissi . (2021) AlKhamissi_Gabr_ElNokrashy_Essam_2021AlKhamissi, B., Gabr, M., ElNokrashy, M. Essam, K.  2021Apr.  Adapting MARBERT for Improved Arabic Dialect Identification: Submission to the NADI 2021 Shared Task Adapting marbert for improved arabic dialect identification: Submission to the nadi 2021 shared task.  Proceedings of the Sixth Arabic Natural Language Processing Workshop Proceedings of the sixth arabic natural language processing workshop ( 260–264).  Association for Computational Linguistics.  https://aclanthology.org/2021.wanlp-1.29
  * Ba . (2016) Ba2016LayerNormBa, JL., Kiros, JR. Hinton, GE.  2016\.  Layer Normalization. Layer normalization.  arXiv.  https://arxiv.org/abs/1607.06450 10.48550/ARXIV.1607.06450
  * Bapna . (2018) bapna2018trainingBapna, A., Chen, MX., Firat, O., Cao, Y. Wu, Y.  2018\.  Training Deeper Neural Machine Translation Models with Transparent Attention Training deeper neural machine translation models with transparent attention.  CoRRabs/1808.07561.  http://arxiv.org/abs/1808.07561
  * Ben-Zaken . (2022) BenZaken2022BitFitSPBen-Zaken, E., Ravfogel, S. Goldberg, Y.  2022\.  BitFit: Simple Parameter-efficient Fine-tuning for Transformer-based Masked Language-models Bitfit: Simple parameter-efficient fine-tuning for transformer-based masked language-models.  ACL. Acl.
  * Brown . (2020) gpt3Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, JD., Dhariwal, P.Amodei, D.  2020\.  Language Models are Few-Shot Learners Language models are few-shot learners.  H. Larochelle, M. Ranzato, R. Hadsell, MF. Balcan H. Lin (), Advances in Neural Information Processing Systems Advances in neural information processing systems ( 33, 1877–1901).  Curran Associates, Inc.  https://proceedings.neurips.cc/paper/2020/file/1457c0d6bfcb4967418bfb8ac142f64a-Paper.pdf
  * Clevert . (2015) eluClevert, DA., Unterthiner, T. Hochreiter, S.  2015\.  Fast and Accurate Deep Network Learning by Exponential Linear Units (ELUs). Fast and accurate deep network learning by exponential linear units (elus).  arXiv.  https://arxiv.org/abs/1511.07289 10.48550/ARXIV.1511.07289
  * Conneau . (2019) xlm-robertaConneau, A., Khandelwal, K., Goyal, N., Chaudhary, V., Wenzek, G., Guzmán, F.Stoyanov, V.  2019\.  Unsupervised Cross-lingual Representation Learning at Scale Unsupervised cross-lingual representation learning at scale.  CoRRabs/1911.02116.  http://arxiv.org/abs/1911.02116
  * Dehghani . (2019) universal-transformerDehghani, M., Gouws, S., Vinyals, O., Uszkoreit, J. Kaiser, L.  2019\.  Universal Transformers Universal transformers.  ArXivabs/1807.03819.
  * Devlin . (2019) devlin-etal-2019-bertDevlin, J., Chang, MW., Lee, K. Toutanova, K.  201906\.  BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding BERT: Pre-training of deep bidirectional transformers for language understanding.  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) 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) ( 4171–4186).  Minneapolis, MinnesotaAssociation for Computational Linguistics.  https://aclanthology.org/N19-1423 10.18653/v1/N19-1423
  * Dou . (2019) dou2019dynamicDou, Z., Tu, Z., Wang, X., Wang, L., Shi, S. Zhang, T.  2019\.  Dynamic Layer Aggregation for Neural Machine Translation with Routing-by-Agreement Dynamic layer aggregation for neural machine translation with routing-by-agreement.  CoRRabs/1902.05770.  http://arxiv.org/abs/1902.05770
  * Fan . (2020) Fan2020AccessingHRFan, A., Lavril, T., Grave, E., Joulin, A. Sukhbaatar, S.  2020\.  Accessing Higher-level Representations in Sequential Transformers with Feedback Memory Accessing higher-level representations in sequential transformers with feedback memory.  ArXivabs/2002.09402.
  * Graves . (2014) neural-turing-machineGraves, A., Wayne, G. Danihelka, I.  2014\.  Neural Turing Machines Neural turing machines.  ArXivabs/1410.5401.
  * He . (2016) resnetHe, K., Zhang, X., Ren, S. Sun, J.  2016\.  Deep Residual Learning for Image Recognition Deep residual learning for image recognition.  2016 IEEE Conference on Computer Vision and Pattern Recognition (CVPR)770-778.
  * Hendrycks Gimpel (2016) geluHendrycks, D. Gimpel, K.  2016\.  Bridging Nonlinearities and Stochastic Regularizers with Gaussian Error Linear Units Bridging nonlinearities and stochastic regularizers with gaussian error linear units.  ArXivabs/1606.08415.
  * Houlsby . (2019) adaptersHoulsby, N., Giurgiu, A., Jastrzebski, S., Morrone, B., de Laroussilhe, Q., Gesmundo, A.Gelly, S.  2019\.  Parameter-Efficient Transfer Learning for NLP Parameter-efficient transfer learning for nlp.  ICML. Icml.
  * Huang . (2016) huang2017densenetHuang, G., Liu, Z. Weinberger, KQ.  2016\.  Densely Connected Convolutional Networks Densely connected convolutional networks.  CoRRabs/1608.06993.  http://arxiv.org/abs/1608.06993
  * Lewis . (2020) lewis-etal-2020-bartLewis, M., Liu, Y., Goyal, N., Ghazvininejad, M., Mohamed, A., Levy, O.Zettlemoyer, L.  202007\.  BART: Denoising Sequence-to-Sequence Pre-training for Natural Language Generation, Translation, and Comprehension BART: Denoising sequence-to-sequence pre-training for natural language generation, translation, and comprehension.  Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics Proceedings of the 58th annual meeting of the association for computational linguistics ( 7871–7880).  OnlineAssociation for Computational Linguistics.  https://aclanthology.org/2020.acl-main.703 10.18653/v1/2020.acl-main.703
  * Li . (2018) Li2018LayerLi, J., Liu, C. Gong, Y.  2018Aug.  Layer Trajectory LSTM Layer trajectory lstm.  arXivabs/1808.09522.  http://arxiv.org/abs/1808.09522
  * F. Liu . (2020) liu2020layerLiu, F., Ren, X., Zhao, G. Sun, X.  2020\.  Layer-Wise Cross-View Decoding for Sequence-to-Sequence Learning Layer-wise cross-view decoding for sequence-to-sequence learning.  CoRRabs/2005.08081.  https://arxiv.org/abs/2005.08081
  * X. Liu . (2021) liu2021understandingLiu, X., Wang, L., Wong, DF., Ding, L., Chao, LS. Tu, Z.  2021\.  Understanding and Improving Encoder Layer Fusion in Sequence-to-Sequence Learning Understanding and improving encoder layer fusion in sequence-to-sequence learning.  International Conference on Learning Representations. International conference on learning representations.  https://openreview.net/forum?id=n1HD8M6WGn
  * Y. Liu . (2019) Liu2019RoBERTaLiu, Y., Ott, M., Goyal, N., Du, J., Joshi, M., Chen, D.Stoyanov, V.  2019\.  RoBERTa: A Robustly Optimized BERT Pretraining Approach Roberta: A robustly optimized bert pretraining approach.  ArXivabs/1907.11692.
  * Loshchilov Hutter (2017) adamwLoshchilov, I. Hutter, F.  2017\.  Decoupled Weight Decay Regularization. Decoupled weight decay regularization.  arXiv.  https://arxiv.org/abs/1711.05101 10.48550/ARXIV.1711.05101
  * Merity . (2017) Merity2017PointerSMMerity, S., Xiong, C., Bradbury, J. Socher, R.  2017\.  Pointer Sentinel Mixture Models Pointer sentinel mixture models.  ArXivabs/1609.07843.
  * Nakayama (2018) seqevalNakayama, H.  2018\.  seqeval: A Python framework for sequence labeling evaluation. seqeval: A python framework for sequence labeling evaluation.  https://github.com/chakki-works/seqeval Software available from https://github.com/chakki-works/seqeval
  * Nøkland (2016) DFANøkland, A.  2016\.  Direct Feedback Alignment Provides Learning in Deep Neural Networks. Direct feedback alignment provides learning in deep neural networks.  arXiv.  https://arxiv.org/abs/1609.01596 10.48550/ARXIV.1609.01596
  * Rahimi . (2019) rahimi-etal-2019-wikiannRahimi, A., Li, Y. Cohn, T.  201907\.  Massively Multilingual Transfer for NER Massively multilingual transfer for NER.  Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics Proceedings of the 57th annual meeting of the association for computational linguistics ( 151–164).  Florence, ItalyAssociation for Computational Linguistics.  https://www.aclweb.org/anthology/P19-1015
  * Shen . (2018) shen2018denseShen, Y., Tan, X., He, D., Qin, T. Liu, T.  2018\.  Dense Information Flow for Neural Machine Translation Dense information flow for neural machine translation.  CoRRabs/1806.00722.  http://arxiv.org/abs/1806.00722
  * Tjong Kim Sang De Meulder (2003) tjong-kim-sang-de-meulder-2003-conllTjong Kim Sang, EF. De Meulder, F.  2003\.  Introduction to the CoNLL-2003 Shared Task: Language-Independent Named Entity Recognition Introduction to the CoNLL-2003 shared task: Language-independent named entity recognition.  Proceedings of the Seventh Conference on Natural Language Learning at HLT-NAACL 2003 Proceedings of the seventh conference on natural language learning at HLT-NAACL 2003 ( 142–147).  https://www.aclweb.org/anthology/W03-0419
  * Vaswani . (2017) vaswani17attentionVaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, AN.Polosukhin, I.  2017\.  Attention is All You Need Attention is all you need.  Proceedings of the 31st International Conference on Neural Information Processing Systems Proceedings of the 31st international conference on neural information processing systems ( 6000–6010).  Red Hook, NY, USACurran Associates Inc.
  * Wallat . (2021) wallat_bertnesia_2021Wallat, J., Singh, J. Anand, A.  202109\.  BERTnesia: Investigating the capture and forgetting of knowledge in BERT BERTnesia: Investigating the capture and forgetting of knowledge in BERT.  arXiv:2106.02902 [cs].  [2021-12-03]http://arxiv.org/abs/2106.02902 arXiv: 2106.02902
  * A. Wang . (2019) wang2019superglueWang, A., Pruksachatkun, Y., Nangia, N., Singh, A., Michael, J., Hill, F.Bowman, SR.  2019\.  SuperGLUE: A Stickier Benchmark for General-Purpose Language Understanding Systems Superglue: A stickier benchmark for general-purpose language understanding systems.  CoRRabs/1905.00537.  http://arxiv.org/abs/1905.00537
  * A. Wang . (2018) wang2018glueWang, A., Singh, A., Michael, J., Hill, F., Levy, O. Bowman, SR.  2018\.  GLUE: A Multi-Task Benchmark and Analysis Platform for Natural Language Understanding GLUE: A multi-task benchmark and analysis platform for natural language understanding.  CoRRabs/1804.07461.  http://arxiv.org/abs/1804.07461
  * Q. Wang . (2019) wang2019learningWang, Q., Li, B., Xiao, T., Zhu, J., Li, C., Wong, DF. Chao, LS.  2019\.  Learning Deep Transformer Models for Machine Translation Learning deep transformer models for machine translation.  CoRRabs/1906.01787.  http://arxiv.org/abs/1906.01787
  * Q. Wang . (2018) wang2018multiWang, Q., Li, F., Xiao, T., Li, Y., Li, Y. Zhu, J.  2018\.  Multi-layer Representation Fusion for Neural Machine Translation Multi-layer representation fusion for neural machine translation.  Proceedings of the 27th International Conference on Computational Linguistics Proceedings of the 27th international conference on computational linguistics ( 3015–3026).  [2021-06-25]https://www.aclweb.org/anthology/C18-1255.pdf
  * Waskom (2021) seabornWaskom, ML.  2021\.  seaborn: statistical data visualization seaborn: statistical data visualization.  Journal of Open Source Software6603021.  https://doi.org/10.21105/joss.03021 10.21105/joss.03021
  * Wolf . (2019) Wolf2019HuggingFaceWolf, T., Debut, L., Sanh, V., Chaumond, J., Delangue, C., Moi, A.Brew, J.  2019\.  HuggingFace’s Transformers: State-of-the-art Natural Language Processing Huggingface’s transformers: State-of-the-art natural language processing.  ArXivabs/1910.03771.  http://arxiv.org/abs/1910.03771