Source-linked AI summary
Funnel-Transformer: Filtering out Sequential Redundancy for Efficient Language Processing
Zihang Dai, Guokun Lai, Yiming Yang, Quoc V. Le
TL;DR
Transformer language processing is expensive partly because full-length token representations are maintained even for sequence-level tasks needing one vector. Funnel-Transformer compresses hidden sequences with an encoder, restores token representations with a decoder when needed, and outperforms standard Transformers on most evaluated tasks with comparable or fewer FLOPs.
Problem
Full-length hidden-state sequences create redundancy for sequence-level tasks while Transformer pretraining and finetuning require substantial computation and memory.
Method
Funnel-Transformer gradually compresses sequence length in an encoder, reinvests saved FLOPs in model capacity, and uses a decoder to recover token-level representations when required.
Results
With fewer or comparable FLOPs, Funnel-Transformer outperforms corresponding standard-Transformer baselines on most tasks across GLUE, text classification, and reading comprehension.
Takeaways & Limitations
Sequential-resolution compression can improve efficiency and model performance across varied sequence-level prediction tasks under the pretraining-finetuning paradigm.
Abstract
from arXiv · showhide
With the success of language pretraining, it is highly desirable to develop more efficient architectures of good scalability that can exploit the abundant unlabeled data at a lower cost. To improve the efficiency, we examine the much-overlooked redundancy in maintaining a full-length token-level presentation, especially for tasks that only require a single-vector presentation of the sequence. With this intuition, we propose Funnel-Transformer which gradually compresses the sequence of hidden states to a shorter one and hence reduces the computation cost. More importantly, by re-investing the saved FLOPs from length reduction in constructing a deeper or wider model, we further improve the model capacity. In addition, to perform token-level predictions as required by common pretraining objectives, Funnel-Transformer is able to recover a deep representation for each token from the reduced hidden sequence via a decoder. Empirically, with comparable or fewer FLOPs, Funnel-Transformer outperforms the standard Transformer on a wide variety of sequence-level prediction tasks, including text classification, language understanding, and reading comprehension. The code and pretrained checkpoints are available at https://github.com/laiguokun/Funnel-Transformer.
1 Introduction
Funnel-Transformer targets the high cost of Transformer pretraining and finetuning by addressing redundancy in full-length hidden-state sequences, especially for sequence-level tasks.
- Transformer models deliver strong NLP performance but require substantially more FLOPs and memory than traditional NLP models.
- The proposed approach seeks more efficient language processing while retaining the benefits of language pretraining.
- The paper focuses on redundancy from maintaining full-length hidden representations across every Transformer layer.
- Sequence-level tasks such as text classification and ranking often extract only one vector, so token-level detail may be unnecessary.
2 Method
Funnel-Transformer compresses hidden-state sequences through an encoder and restores token-level representations with a decoder when pretraining or other token-level prediction requires them.
- Transformer Architecture: A Transformer layer combines multi-head self-attention and a position-wise feed-forward network, with residual connections and layer normalization.
- Pretraining Objectives: Masked language modeling corrupts 15% of tokens and trains the model to reconstruct the originals from the corrupted sequence.
- Discussion: Token-level pretraining requires full-sequence hidden states, whereas classification and ranking generally need only a single sequence vector.
- Encoder: The encoder reduces sequence length between blocks using pooling while preserving the Transformer layer structure within blocks.
- Encoder: Pool-query-only attention uses pooled states for queries and residual signals while unpooled states provide keys and values, producing expressive compression.
- Decoder: The decoder upsamples the reduced representation, adds it to first-block token-level states, and applies two Transformer layers when token-level predictions are needed.
- Complexity and Capacity: A Transformer layer costs O(T^2D + TD^2), so halving sequence length yields a super-linear computation reduction and frees FLOPs for additional depth.
- Complexity and Capacity: Memory complexity is O(T^2 + TD), and the analysis focuses on computation complexity because memory is offset by a 1/D multiplier.
3 Related Work
The paper positions Funnel-Transformer among efficiency methods that remove redundant operations or representations, distinguishing it from related sequence-compression approaches.
- Funnel-Transformer is presented as achieving performance gains by compressing hidden-state sequence length under language pretraining.
- Its key distinctions include pool-query-only down-sampling, its up-sampling method, and relative attention parameterization.
4 Experiment
Experiments pretrain Funnel-Transformers at base and large scales, then evaluate them on sequence-level language understanding, classification, and reading-comprehension tasks under comparable computation.
- Experimental Settings: Base-scale experiments use 1M pretraining steps with batch size 256 on Wikipedia and Book Corpus, enabling comparisons with standard Transformers.
- Experimental Settings: Large-scale experiments use 500K steps with batch size 8K on five corpora and evaluate primarily sequence-level tasks.
- Base-scale Results: The study compares large, base, and small Transformer baselines with Funnel-Transformer layouts having similar or fewer FLOPs.
- Base-scale Results: Table 1 reports GLUE development performance and text-classification error rates, while FLOPs and parameter counts refer to encoder-only finetuning.
- Base-scale Results: With similar or fewer FLOPs, Funnel-Transformer outperforms standard Transformers on most tasks except STS-B, especially for smaller models.
- Base-scale Results: ELECTRA pretraining shows a similar overall trend, although the GLUE gain is slightly smaller.
- Running Time Comparison: The study also evaluates actual runtime and memory footprint because FLOPs do not fully determine execution speed.
- Large-scale Results: Large-scale Funnel-Transformers outperform corresponding baselines on most GLUE tasks with fewer or comparable FLOPs, supporting scalability.
Model CoLA SST-2 MRPC STS-B QQP MNLI QNLI RTE WNLI AVG
The ablation study compares pooling operations, positional parameterizations, special [cls] handling, and block layouts in F-TFM. Relative positional parameterization and the three-block layout are identified as important design choices.
- Pooling operation: Mean and max pooling perform similarly and significantly better than Top-Attn pooling for selecting “hub” states.
- Special designs: Pool-query-only and maintaining a separate non-pooled [cls] vector each improve the proposed model.
- Positional parameterization: Relative positional parameterization is key to F-TFM performance compared with learned absolute position embeddings.
- Block layout design: The three-block B6-6-6 layout outperforms the compared two-block B8-8 and four-block B5-5-5-5 layouts with similar FLOPs and parameter counts.
5 Conclusion & Discussion
The paper frames Funnel-Transformer as a way to compress sequential resolution, save computation, and reinvest saved FLOPs in model capacity and performance. It identifies improving compression, block layouts, FLOP reinvestment, and combining compression with distillation or quantization as future directions.
- Funnel-Transformer compresses sequential resolution to save computation and reinvests saved FLOPs in model capacity and performance.
- Future work includes better compression schemes, optimized block layouts, and improved strategies for reinvesting saved FLOPs.
- Combining Funnel-Transformer with knowledge distillation and quantization is proposed as a direction for improving practical impact.
A Implementation Optimization
To preserve the separate [cls] state without incurring irregular sequence lengths, the implementation truncates the final pooled token so lengths remain powers of two. This avoids accelerator inefficiency and associated slowdown.
- The implementation keeps [cls] intact during pooling to preserve the structure commonly used in pretraining.
- Truncating the final pooled token restores a power-of-two sequence length after retaining [cls].
- Maintaining power-of-two lengths avoids memory misalignment, wasted parallel computation, and the resulting slowdown.
A.2 Relative Positional Attention Implementation
The implementation avoids the prohibitive cost of explicitly materializing all pairwise relative-position encodings. It uses either gather-based reconstruction or TPU-friendly factorization to keep position-term computation manageable.
- Relative positional attention: Relative positional attention combines content and positional terms, with sinusoidal encodings representing relative distance i − j.The content term can be computed with complexity O(TD^2 + T^2D).
- Naive implementation: O(T^2D^2) computation and O(T^2D) memory make the naive position-encoding implementation the dominant cost.The naive approach stacks T^2 position-encoding pairs into a tensor.
- Standard Solution: Gather / Shift: The gather/shift solution stores 2T − 1 relative encodings and reconstructs the full tensor with a permutation tensor selecting one encoding per position pair.Because the attention computation is linear, gathering can be rearranged to reduce complexity to O(2TD^2 + 2T^2D).
- Optimization for TPU: On TPUs, gathering T^2 elements is dramatically slower than assumed, motivating a separate implementation that avoids this operation.The TPU alternative rewrites the positional term using factorized components.
A.3 Potential Model Extensions
Funnel-Transformer generalizes an encoder–decoder framework in which the encoder compresses its input sequence and the decoder extracts task-relevant representations. This supports sequence-level, token-level, and sequence-to-sequence tasks.
- General framework: The encoder compresses xenc into shorter hidden representations henc, while an optional decoder uses henc and xdec to solve the NLP task.The decoder input is problem-specific and may be absent.
- Sequence-level prediction: For sequence-level prediction, no decoder input is needed; the model extracts the [cls] representation from henc for the task-specific structure.This applies to tasks such as text classification.
- Token-level prediction: Token-level prediction uses a decoder to recover token representations from compressed henc, optionally supplying the original sequence or early token-level states.The additional low-level information can ease optimization.
- Sequence-to-sequence problems: For sequence-to-sequence problems, Funnel-Transformer uses source text as xenc and target text as xdec within its compressed encoder–decoder framework.The stated examples include machine translation, summarization, and dialog generation.
B Experiment Setting and Hyper-parameters
Experiments use BERT’s uncased word-piece preprocessing, base- and large-scale pretraining settings, and XLNet-derived finetuning procedures. Hyperparameters and evaluation procedures are specified separately for pretraining, GLUE, classification, RACE, and SQuAD.
- Preprocessing: The experiments use BERT’s uncased word-piece model with a vocabulary of about 30K and remove selected HTTP symbols for seven text-classification tasks.Other preprocessing follows the tokenizer’s default lowercasing and processing.
- Pretraining hyperparameters: Training can become unstable beyond 24 layers for B10-10-10H1024 at base scale, so its F-TFM learning rate is reduced to 8e-5.The authors report slower training and potentially slightly worse finetuning performance as side effects.
- Pretraining hyperparameters: The models use P-FFN inner states of 4 × hidden size, attention-head dimension 64, and hidden size/64 attention heads.These settings define the principal attention and feed-forward dimensions.
- Mask sampling: MLM uses word-span masking, whereas ELECTRA uses single-token sampling at base scale because word-span sampling significantly lowers performance there.The masking behavior differs by pretraining objective and scale.
- Finetuning: Finetuning largely inherits XLNet hyperparameters, searches learning rates 1e-5, 2e-5, and 3e-5 for GLUE and classification, and reports medians across multiple seeds.RACE and SQuAD use separate summarized settings, including layer-wise decay.
C Additional Experimental Results
Under large-scale training, Funnel-Transformer achieves comparable text-classification performance with fewer FLOPs than Transformer baselines. Several configurations improve results or provide substantial speedups while preserving performance closely.
- Large-scale text classification: With fewer FLOPs, F-TFM achieves comparable results to previous Transformer-based results across seven large-scale text-classification tasks.The comparison is reported in Table 10.
C.2 Training Cost Comparison
Funnel-Transformer reduces training cost across finetuning and pretraining while supporting deeper or otherwise comparable-capacity models. Across GPU and TPU evaluations, it achieves speed, memory, and efficiency advantages over standard Transformer baselines.
- Finetuning speed: B6-6-6H768 matches or exceeds the base Transformer’s speed despite being deeper and having more parameters.The comparison is made within the base-model group.
- Finetuning speed: 30%-50% speedup is achieved by B4-4-4 without losing too much performance relative to the base Transformer.B4-4-4 has the same depth and model parameters as the base model.
- Finetuning speed: B8-8-8 is significantly faster than the large Transformer, while B10-10-10 runs at almost the same speed.These comparisons concern the large-model group.
- Memory efficiency: All tested F-TFM variants use less memory than standard Transformer models.The lower memory footprint is reported for both base and large model groups.
- Pretraining speed: Funnel-Transformer substantially improves pretraining speed despite the additional decoder cost under the masked-language-modeling objective.The speed gain is slightly smaller than during finetuning, but the method remains more efficient in both stages.