EPAS: Efficient Training with Progressive Activation Sharing

TL;DR

EPAS addresses the high computational cost of training and inferring large Transformer models by exploiting redundancy in deep-layer activations. It introduces a switchable activation-sharing decoder and a deterministic, progressive schedule that expands a sharing region from the model's deep layers toward the shallow ones, reducing FLOPs without sacrificing accuracy. Empirically, EPAS achieves up to 11% faster training throughput and up to 29% faster inference throughput on LLaMA-family models, with minimal loss in final performance. Additionally, continual pretraining with EPAS converts pretrained models into efficient activation-sharing variants, yielding about a 10% improvement in LM accuracy under sharing, demonstrating the method's practical impact for efficient, budget-flexible large-language models.

Abstract

We present a novel method for Efficient training with Progressive Activation Sharing (EPAS). This method bridges progressive training paradigm with the phenomenon of redundant QK (or KV ) activations across deeper layers of transformers. EPAS gradually grows a sharing region during training by switching decoder layers to activation sharing mode. This results in throughput increase due to reduced compute. To utilize deeper layer redundancy, the sharing region starts from the deep end of the model and grows towards the shallow end. The EPAS trained models allow for variable region lengths of activation sharing for different compute budgets during inference. Empirical evaluations with QK activation sharing in LLaMA models ranging from 125M to 7B parameters show up to an 11.1% improvement in training throughput and up to a 29% improvement in inference throughput while maintaining similar loss curve to the baseline models. Furthermore, applying EPAS in continual pretraining to transform TinyLLaMA into an attention-sharing model yields up to a 10% improvement in average accuracy over state-of-the-art methods, emphasizing the significance of progressive training in cross layer activation sharing models.

Figures (5)

• Figure 1: Left: Overall solution from efficient training to inference using EPAS. Right: TinyLLaMA tinyllama model FLOPs reduction and train/inference throughput improvement expanding $QK$ activation sharing to 25% and 50% of the layers.
• Figure 2: The Switchable Activation Sharing Decoder with an example of attention sharing ($Q,K$). This decoder layer extends conventional transformer decoder layer by adding a conditional switching branch to reuse $Q_{i-1},K_{i-1}$ from previous layer instead of computing in current layer(left branch inside dashed box). When not using activation sharing mode, the computation follows the right branch as like convention decoder layer.
• Figure 3: An example of EPAS training approach. Beginning with all $L$ layers in compute-mode (e.g., 5 layers here), a region of $B$ layers transition into sharing-mode ( e.g., 1 layer here) at every $I$ step intervals. The progressive growth continues till maximum $S$ layers (e.g., 3 layers here) are in sharing mode. The trained model can be used either all layer in compute mode or up to $S$ layers in activation sharing mode.
• Figure 4: Smoothed loss versus time while scaling model sizes across LLaMA models with $125M$, $1.1B$, and $3B$ parameters. $QK$ sharing models trained with EPAS also exhibit slightly faster convergence during training in addition to have higher throughput during inference.
• Figure 5: Schematic illustration of layer grouping: multiple small blocks versus a single large block. The colors indicate compute- and sharing-mode. The number of layers shown is for illustrative purposes; both approaches can accommodate a variable number of layers.