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

TL;DR

DWAtt addresses the problem that intermediate representations in deep pretrained transformers are underutilized during adaptation. It introduces Depth-Wise Attention (DWAtt), an add-on that re-surfaces signals from all encoder layers and aggregates them with a layer-aware attention mechanism, contrasting with a simple layer concatenation baseline. Across few-shot NER (CoNLL-03, WikiAnn) and MLM (WikiText-2), DWAtt and Concat improve data- and step-efficiency relative to a deeper baseline, with DWAtt outperforming Concat as data scale increases; for CoNLL-03, layer fusion yields F1 gains in the range $3.68 ightarrow 9.73\%$ across different few-shot sizes. The results support layer fusion as a practical, data-efficient strategy for adapting large pretrained models, with potential applicability to other tasks and architectures.

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.

Figures (9)

• 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.
• 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.
• 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.
• 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.
• Figure 5: F1-Score on the WikiAnn English devset on the RoBERTaLARGE base model. The methods start behaving similarly at higher training epochs. DWAtt leads Concat leads R26 at medium shot scenarios.