Mixture-of-Modules: Reinventing Transformers as Dynamic Assemblies of Modules (2024)

Abstract: Is it always necessary to compute tokens from shallow to deep layers in Transformers? The continued success of vanilla Transformers and their variants suggests an undoubted “yes”. In this work, however, we attempt to break the depth-ordered convention by proposing a novel architecture dubbed mixture-of-modules (MoM), which is motivated by an intuition that any layer, regardless of its position, can be used to compute a token as long as it possesses the needed processing capabilities. The construction of MoM starts from a finite set of modules defined by multi-head attention and feed-forward networks, each distinguished by its unique parameterization. Two routers then iteratively select attention modules and feed-forward modules from the set to process a token. The selection dynamically expands the computation graph in the forward pass of the token, culminating in an assembly of modules. We show that MoM provides not only a unified framework for Transformers and their numerous variants but also a flexible and learnable approach for reducing redundancy in Transformer parameterization. We pre-train various MoMs using OpenWebText. Empirical results demonstrate that MoMs, of different sizes, consistently outperform vanilla transformers. More interestingly, after removing 50% of the multi-head attention modules and 25% of the feed-forward modules, an MoM model still holds comparable performance. Additionally, by properly adjusting the number of modules and compressing the model depth, one can have an MoM that achieves comparable performance to GPT-2 (774M) while saving 16% TFLOPs and 42% memory usage during forward computation.

## introduction transformer-based language models @xcite have demonstrated remarkable abilities across a wide range of challenging natural language tasks @xcite . in addition, the success of transformer in natural language processing (nlp) is also inspiring innovations in other fields such as computer vision @xcite and biomedicine @xcite . a transformer architecture typically consists of stacked layers that are identical in structure, whereby layers are organized in the order of depth, using the output of the previous layer as the input for the next. while this design convention has been widely accepted as a matter of course in the transformer era, we seek to challenge it by reconsidering whether the static and depth-ordered organization can fully unleash the potential of transformers, given their wellknown issues with over-parameterization @xcite and efficiency @xcite . before us, some rudimentary studies have touched on the question-they dissect transformer into modules such as attention heads and feedforward networks (ffns) and allow relatively flexible module call order. for example, mixtureof-experts (moe) @xcite ) sets up multiple ffns within the same layer and activates a specific subset during inference. earlyexiting @xcite @xcite and mixture-of-depths (mod) @xcite ) bypass certain layers when computing each token. on the one hand, these efforts indeed lead to improvements in terms of either efficacy or efficiency through the introduction of dynamic mechanisms into the vanilla structure of transformers, and thus corroborate our questioning regarding the established convention; on the other hand, they still follow the depth-ordered paradigm (i.e., tokens are passed from shallow layers to deep layers), leaving significant room for better architectures. in this work, we completely disrupt the traditional practice in the design of transformers by breaking down the depth-ordered organization. numerous studies have indicated that knowledge in transformers is often dispersed across multiple ffns in different layers @xcite @xcite , and many attention heads serve similar or identical functions, such as copying specific token information towards the end position of the input @xcite . encouraged by this evidence, we pose the question of whether the computation of a token can "move" freely across layers, that is the token can be computed by flowing to modules in deeper layers, sticking to modules of the same layer, or even going back to modules in previous layers. to answer the question, we propose a novel architecture dubbed mixture-of-modules (mom) in which the core idea is to define a neural network as dynamic assemblies of modules derived from vanilla transformer, as depicted in figure 1 . the basis of mom is a finite set of modules. each module is defined by multi-head attention (mha), a feed-forward network (ffn) (including add & norm), or a specialized module labeled "skip". each mha or ffn module is identical in structure and different in parameterization. skip enables skip operations for arbitrary tokens at arbitrary time steps. given a token, each time two routers select modules from the set and integrate the modules into the computation graph during the forward pass. hence, the whole computation graph of the token is formed as an assembly of modules, and the routers learn to optimize the organization of the modules in the assembly. we introduce a two-phase approach for training mom models. in the first phase, we pre-train a vanilla transformer on a large-scale corpus. then, in the second phase, we decompose the pre-trained transformer into modules as a warm-up of mom, randomly initialize the routers, and further update both the modules and the routers under the mechanism of dynamic assembly. mom has three major advantages over existing transformer-based architectures: (1) it provides a unified framework for various transformer variants, incorporating popular methods such as mixtureof-experts, early-exiting, and mixture-of-depths as special cases. the framework sheds light on architecture design for future works; (2) it brings unprecedented flexibility in forward computation. with the dynamic assembly mechanism, the conventional concept of "depth" now has a new interpretation: it refers to the number of routings performed per token. useful modules such as specific attention heads can be repeatedly utilized as long as they are capable of processing. hence, mom offers a dynamic and learnable approach to reducing redundant parameters in transformers; and (3) it offers efficient structures that achieve performance comparable to vanilla transformers but require significantly fewer flops and less memory in forward computation. we pre-train mom in three sizes-122m (small), 346m (medium), and 774m (large)-using open-webtext @xcite , and assess their performance with glue benchmark @xcite and xsum @xcite . empirical results indicate that (1) moms consistently outperform vanilla gpt-2 models on both text understanding and generation tasks; (2) parameters are quite redundant in vanilla transformers. mom can hold the comparable performance after removing 50% of the mha modules and 25% of the ffn modules; and (3) for those concerned with efficiency, mom-large can reduce tflops by 16% and memory usage by 42% in forward computation, while maintaining comparable performance to gpt-2 (774m), via properly increasing the number of modules and compressing the model depth. our contributions are three-fold: (1) proposal of mixture-of-modules to disrupt the depth-ordered convention in transformer construction, and reinvent transformers as dynamic assemblies of modules; (2) empirical verification of the efficacy of mixture-of-modules on glue benchmark and xsum; and (3) a series of new insights into the over-parameterization issue of vanilla transformers, and their implications for future architecture design. ## related works mom owns a dynamic mechanism of module selection and combination, and thus is related to conditional computation techniques @xcite @xcite . existing work on conditional computation can be categorized into two groups: dynamic depth and dynamic width. in these fields, terms such as gating and routing are used interchangeably, hereafter referred to as "routers" for clarity in presentation. the assembly lasts h rounds, with the current illustration showcasing progress in the third round. for each token, routers select the best k attention modules, denoted as m a k , and the best k feed-forward network modules, denoted as m f k , from a module set m (including "skip" modules). these selected modules collectively constitute assembled modules f a and f f , which are then appended to the existing computation graph. detailed notations are presented in §3. as a typical approach in dynamic depth, earlyexiting @xcite @xcite accelerates model inference through terminating forward computation at intermediate layers. the decision to exit often relies on confidence-based metrics @xcite @xcite or pre-determined strategies @xcite . with some degree of generalization, layer-skip @xcite @xcite represents a more adaptive variant of early-exiting, enabling certain layers to be skipped without terminating the entire forward computation. existing works mainly facilitate it by training a router @xcite or layer pruning @xcite . finally, if we view parameter copying as a particular way to increase network depth with controlled model size, then some parameter sharing methods @xcite , wherein certain modules or layers share parameters, also fall in the dynamic depth group. in terms of dynamic width, mixture-of-experts (moe, @xcite @xcite ) is a representative method. an moe model conceptualizes an ffn module as an "expert" for storing knowledge. comprising multiple such experts, an moe layer replaces the traditional ffn layer within transformers, aiming for superior performance in handling knowledgerelated tasks. during forward computation, a router network dynamically assigns each token to the top k experts out of a total of n experts, thereby increasing the maximum network width by k times. other dynamic width methods, such as coda @xcite and colt5 @xcite , use similar routing mechanism to select whether a token passes through a heavy or light pathway for not only each ffn layer but also each attention layer. mom breaks the depth-ordered paradigm followed by existing approaches when performing forward computation. it not only unifies a number of approaches described above but also offers a more flexible and learnable way to achieve conditional computation. ## methodology the idea of mixture-of-modules (mom) is inspired by the "society of mind" theory by marvin minsky @xcite , which explains the true intelligence as certain and very special ways of combinations of simple and modular units (in the book, they are termed "agents"). in §3.1, we first provide an overview of mom. then, we detail the assembly of modules and the routers in §3.2 and §3.3. after that, we present the training procedure of mom in §3.4. finally in §3.5, we show that mom unifies various techniques of dynamic computation allocation within transformers as special cases. ## conclusions in this work, we propose mixture-of-modules (mom), a novel architecture that reinvents transformers as a collection of individual modules and the dynamic assembly process conducted with these modules. this novel view offers us an opportunity to explore a wide range of different configurations of model architecture and unify a series of transformer variants. with exhaustive experiments, we not only validate the effectiveness of mom by both significant efficiency and performance gains but also reach new insights about transformers.