Source-linked AI summary
DenseFormer: Enhancing Information Flow in Transformers via Depth Weighted Averaging
Matteo Pagliardini, Amirkeivan Mohtashami, Francois Fleuret, Martin Jaggi
TL;DR
Large Transformers deliver strong capabilities but impose substantial computational, memory, and data demands. DenseFormer adds depth-weighted averaging between blocks so later blocks directly reuse earlier representations. It matches deeper Transformers while improving efficiency and data use, with learned weights revealing structured information flow.
Problem
Larger Transformers increase computational cost and memory demands, while longer training requires more data and deeper models face diminishing returns.
Method
DenseFormer adds a learned Depth-Weighted-Average module after each Transformer block to combine current, previous, and embedded input representations.
Results
DenseFormer matches much deeper Transformers while being smaller, faster, and less memory-intensive at inference, and achieves lower perplexity at equal inference speed, 17.84 versus 17.94.
Takeaways & Limitations
DenseFormer improves the Transformer perplexity-speed trade-off across varied settings, with stable learned-weight patterns indicating structured reuse of distant-layer activations.
Takeaways & Limitations
More efficient implementations, better sparsity patterns, and methods for sharding DWA across multiple nodes remain future work.
Abstract
from arXiv · showhide
The transformer architecture by Vaswani et al. (2017) is now ubiquitous across application domains, from natural language processing to speech processing and image understanding. We propose DenseFormer, a simple modification to the standard architecture that improves the perplexity of the model without increasing its size -- adding a few thousand parameters for large-scale models in the 100B parameters range. Our approach relies on an additional averaging step after each transformer block, which computes a weighted average of current and past representations -- we refer to this operation as Depth-Weighted-Average (DWA). The learned DWA weights exhibit coherent patterns of information flow, revealing the strong and structured reuse of activations from distant layers. Experiments demonstrate that DenseFormer is more data efficient, reaching the same perplexity of much deeper transformer models, and that for the same perplexity, these new models outperform transformer baselines in terms of memory efficiency and inference time.
1 Introduction
DenseFormer addresses the computational and data demands of increasingly large Transformers by averaging representations from previous blocks. It matches deeper Transformers while improving size, speed, memory efficiency, and data efficiency, with learned weights revealing structured information reuse.
- Motivation: The architecture targets rising computational and memory costs, data scarcity, and diminishing returns from increasing Transformer depth.Large models carry increased computational and memory costs, while longer training requires larger datasets and deeper models may yield diminishing gains.
- Method: DenseFormer averages outputs from all previous blocks before each subsequent block, extending standard skip connections with learned depth-weighted access.The added DWA module follows every Transformer block and uses learned weights over current and earlier representations.
- Results: DenseFormer matches much deeper Transformers while being smaller, faster, and less memory-intensive at inference, without requiring more data.The introduction presents this as the main overall performance and efficiency result.
- Insights: Learned DWA weights show stable patterns across depths and random seeds, suggesting structured reuse of early features through inter-block connectivity.The authors hypothesize that direct reuse reduces the need to propagate early features through many layers and helps resolve ambiguity from ordinary skip connections.
- Contributions: DenseFormer’s contributions include empirical speed-performance comparisons across settings and additional insights supporting its benefits.The listed experiments vary datasets, batch sizes, and sequence lengths, while the implementation is publicly available.
2 Related Work
Related work addresses diminishing returns and information flow through residual, highway, dense, and depth-wise mechanisms, while also pursuing efficient Transformer implementations. DenseFormer operates between blocks and is intended to complement internal block modifications and hardware-aware optimizations.
- Information flow: Prior methods improve information flow from earlier layers through residual connections, Highway Networks, DenseNets, and related architectural mechanisms.These approaches were proposed to address diminishing returns and stabilize or improve deep-network training.
- Depth-wise mechanisms: Depth-wise Attention mixes outputs across Transformer blocks, resembling DenseFormer’s weighted averaging but applying the mechanism before the final projection layer.The supplied passage distinguishes the prior operation from DenseFormer’s proposal, whose weights are learned differently.
- Efficient Transformers: Transformer research also modifies activations, normalization, execution order, attention computation, or redundant operations to reduce computational challenges.These proposals primarily affect the internal structure of Transformer blocks.
- Hardware efficiency: Hardware-aware implementations provide memory and computational gains, motivating DenseFormer’s focus on keeping its added time overhead negligible during training and inference.DenseFormer is presented as compatible with existing efficiency proposals because its DWA modules operate between blocks.
3 Method
DenseFormer adds a Depth-Weighted-Average module after each Transformer block, mixing current, earlier, and embedded representations through learned weights. Dilated and periodic variants sparsify this connectivity to reduce overhead while preserving performance.
- DenseFormer: DenseFormer adds a DWA module after each Transformer block to compute a weighted average of the current block output, earlier outputs, and the embedded input.The learned α weights are the method’s only additional parameters, and the module is initialized to recover the standard Transformer.
- DenseFormer: Each block processes the previous DWA output, then its representation is mixed with X0 and earlier block outputs to produce the next signal.The architecture defines Y0 as X0, computes Xi through block Bi, and computes Yi through DWAi over the available representations.
- 3.1 Impact on Resources: DWA adds negligible model-size and memory overhead because it uses only depth-dependent scalar weights and representations already retained by standard Transformers.At depth i, DWA has i + 1 weights; typical models therefore add at most about 10^3 parameters, while stored activations are already needed for training or inference.
- 3.2 Dilated DenseFormer: Dilation sparsifies DWA connectivity by averaging only every k-th representation, reducing DWA computation by 1/k without noticeable degradation for small k such as 2 or 4.The efficient implementation also avoids unnecessary data movement, and the dilated variant selects representations satisfying j ≡ i (mod k).
- 3.3 Periodic DenseFormer: A 48-block DenseFormer matches the perplexity of a 72-block Transformer, while a 48-block 4x5-DenseFormer matches that result at 1.4× faster inference.These variants improve the speed–performance trade-off, and constrained-connectivity baselines perform worse than DenseFormer.
- 3.3 Periodic DenseFormer: A DWA period p inserts averaging only after every p blocks, reducing computational overhead by 1/p; a kxp-DenseFormer therefore uses 1/kp of regular DenseFormer overhead.Small periods greater than 1 can improve speed without noticeable performance degradation, although the ideal period depends on the dilation factor.
4 Results
DenseFormer consistently improves the speed–perplexity trade-off over standard Transformers across model sizes, datasets, sequence lengths, and training budgets. Its gains persist when matching depth, inference time, or training time, while dilation and DWA periods reduce overhead with little performance loss.
- Main comparisons: DenseFormer outperforms same-depth Transformers in perplexity and matches the perplexity of substantially deeper models while remaining faster and smaller.A 48-block DenseFormer matches a 72-block Transformer, whose parameter count and memory footprint are 45% larger.
- Main comparisons: 17.84 versus 17.94 perplexity shows that a 48-block DenseFormer beats a same-inference-speed 64-block Transformer.Using dilation 4 or DWA period 5 can widen this performance gap further.
- Additional experiments: DenseFormer’s gains are not reproduced by merely adding learned scaling factors to skip connections, supporting the importance of direct access to all previous layers.The experiments include OpenWebText2, PG-19, and longer 512-token sequences, where the advantage persists.
- Efficiency trade-offs: 4x1- and 4x5-DenseFormers push the speed–performance Pareto frontier forward, with small dilation and period values preserving comparable performance while boosting speed.The reported trade-off holds for both inference and training efficiency.
- Training efficiency: Under the same training-time budget, 4x5-DenseFormer reaches lower perplexity faster than a Transformer trained for more iterations.The comparison uses 40k versus 41.5k steps, with the Transformer receiving more data points yet still being outperformed.
5 Analyzing the Information Flow
The learned DWA weights reveal stable, structured reuse of representations across depths, including early embeddings, nearby blocks, and broad aggregation near the output. Small weights remain important, and restricted sparsity patterns do not recover DenseFormer’s full benefit.
- Learned weight patterns: DWA weights exhibit similar patterns in 48- and 72-block models and remain similar under dilation 4, indicating stable information-flow structure.The heatmaps are averaged over three seeds and show consistent patterns across runs.
- Learned weight patterns: High weights concentrate on the diagonal, nearby previous blocks, initial embeddings, and an aggregation region spanning many earlier layers near the end.Embedding weights are positive in earlier layers and negative in later layers.
- Weight importance: Removing more than 15% of the smallest DWA weights sharply increases perplexity, suggesting that small inter-block weights still contribute materially to prediction.The pruning is performed after training by sparsifying weights according to magnitude at inference.
- Alternative sparsity patterns: Restricted patterns such as Last K and Connect to Last fail to achieve DenseFormer’s perplexity improvement, indicating that broad inter-block connectivity matters.The alternatives limit access to recent blocks or place a single DWA after the final layer.
- Embedding correlations: Cosine similarity with input embeddings is high in early layers and decreases later, which the authors hypothesize reflects a shift toward constructing the next-token representation.The interpretation is explicitly presented as a hypothesis about the later processing stage.
6 Future Work & Conclusion
DenseFormer improves the perplexity/speed trade-off over Transformers across varied settings, while learned weight patterns provide insight into its information flow. Future work targets more efficient implementations, sparsity, and distributed sharding.
- 6 Future Work & Conclusion: DenseFormer establishes a superior perplexity/speed trade-off over Transformers across varied experimental settings.The experiments cover different datasets, batch sizes, and sequence lengths.
- 6 Future Work & Conclusion: Dilation and DWA periodicity increase inference speed without significantly hurting performance.
- 6 Future Work & Conclusion: Learned DWA weights reveal patterns that persist across different depths.
- 6 Future Work & Conclusion: Future work should develop more efficient implementations, better sparsity patterns, and methods to shard DWA across multiple nodes.These directions are motivated by the need to support large-scale distributed training.
- 6 Future Work & Conclusion: Naive implementations are slow when dilation and DWA periodicity are not used, motivating optimized tensor-based implementations.
A.1 Naive Pytorch implementation
The naive implementation stores prior representations and computes each DWA output by weighted averaging, with variants supporting dilation and periodicity. An optimized package replaces loops with tensor operations and exposes a simpler integration path.
- A.1 Naive Pytorch implementation: The naive implementation stores each block output and feeds previous representations to DWA modules after every block.
- A.1 Naive Pytorch implementation: DWA computes a weighted average by multiplying stored representations by learned α values and summing them.
- A.1 Naive Pytorch implementation: The kxp variant supports dilation and DWA frequency while reusing the same DWA module as the 1x1 variant.
- A.1 Naive Pytorch implementation: The optimized implementation removes Python loops by accumulating prior representations in tensors and updating the accumulator after each block.
- A.1 Naive Pytorch implementation: The denseformer package orchestrates DWA logic using the number of blocks, dilation factor, and DWA period, enabling Transformer conversion in three steps.
B.1 Information Flow
DenseFormer’s DWA weights form stable, depth-consistent patterns that emerge quickly during training and maintain changing similarity to the initial embeddings across depth.
- B.1 Information Flow: DWA weights exhibit consistent patterns across depths and match the patterns learned by 1x1-DenseFormer.
- B.1 Information Flow: For 72-block models, cosine similarity to initial embeddings remains high early but decreases toward later layers.The figure passage interprets the later-layer decrease as potentially preparing the next-token output.
- B.1 Information Flow: The DWA weight pattern is learned relatively quickly, appearing similar to the final pattern within the first 5000 iterations.
B.2 Analysis of Dilation and DWA Period
Dilation and DWA periodicity can accelerate inference, but larger settings may degrade perplexity and interact non-monotonically with model depth and each other.
- B.2 Analysis of Dilation and DWA Period: Small dilation coefficients increase inference speed without significantly deteriorating perplexity.The figure reports that larger dilation factors cause a noticeable performance drop, with degradation occurring after k = 4.
- B.2 Analysis of Dilation and DWA Period: 12-Dilated DenseFormer still outperforms the Transformer baseline despite the perplexity degradation from larger dilation.
- B.2 Analysis of Dilation and DWA Period: Increasing the DWA period further improves inference speed, while its perplexity profile is non-monotonic.The non-monotonic profile exposes interplay between dilation, period, and model depth.
- B.2 Analysis of Dilation and DWA Period: A 4x20-DenseFormer still outperforms the Transformer baseline despite increasing the DWA period.
B.3 Delaying the Training of DWA Weights
DenseFormer can delay learning its DWA weights while initially behaving like a standard Transformer, with the weights rapidly approaching their final values early in training.
- DWA weights rapidly converge to their final values within the first 5,000 training iterations.
- At initialization, DenseFormer is equivalent to a Transformer because the DWA weights preserve the standard information flow.
- The experiment trains the model as a Transformer for N iterations before tuning DWA weights through the remaining 40,000 iterations.
B.4 Rank Analysis
Rank analysis finds no significant difference between matrices learned by DenseFormers and standard Transformers, including when singular values are averaged across blocks.
- There is no significant difference in matrix rank between DenseFormer and Transformer architectures.
- Across 48-block models, averaged singular-value profiles show no significant differences between Transformers and DenseFormers.The averages exclude the embedding matrix and combine results across 3 seeds.
B.5 Experiments with a batch size of 128
With batch size 128, DenseFormer retains its performance and efficiency advantages, while learned DWA patterns remain consistent and alternative sparse connectivity patterns are less effective.
- Experiments with a batch size of 128: The batch-size-128 experiments reproduce the study using a smaller training batch than the batch size 400 used in the main comparison.
- Speed and performance trade-off: A 48-block DenseFormer performs on par with a 90-block Transformer in perplexity while using fewer blocks.The perplexity gap is larger than with batch size 400, and the comparison concerns inference speed and perplexity.
- Speed and performance trade-off: DenseFormer outperforms standard architectures of the same depth and those with the same inference speed on OpenWebText2.A deeper standard model can sometimes match a shallower DenseFormer, but the DenseFormer remains faster at inference.
- Results with other sparse patterns: Alternative sparse DWA patterns, including Last K and Connect to Last, do not achieve the perplexity boost of the original DenseFormer.The comparison covers 48-block architectures and includes dilation-based sparsity patterns.
- Visualizing the DWA weights: DWA weight patterns learned with batch size 128 are consistent with patterns obtained using a larger batch size.