Routing in Sparsely-gated Language Models responds to Context (2024)

Abstract: Language Models (LMs) recently incorporate mixture-of-experts layers consisting of a router and a collection of experts to scale up their parameter count given a fixed computational budget. Building on previous efforts indicating that token-expert assignments are predominantly influenced by token identities and positions, we trace routing decisions of similarity-annotated text pairs to evaluate the context sensitivity of learned token-expert assignments. We observe that routing in encoder layers mainly depends on (semantic) associations, but contextual cues provide an additional layer of refinement. Conversely, routing in decoder layers is more variable and markedly less sensitive to context.

## introduction language models (lms) have demonstrated exceptional capabilities in capturing linguistic nuances @xcite and generating coherent text @xcite . however, the dense nature of their architectures, where each token is processed by the total number of parameters, inherently limits their scalability, which is considered the predominant driver for their advanced expressiveness @xcite . sparsely-gated mixture-of-experts (moe) models as developed by @xcite and more recently integrated into the transformer architecture @xcite by @xcite and @xcite , emerged as a promising technique to scale up the parameter count of denselyconnected language models @xcite . beyond language models, this design paradigm was successfully applied to vision models @xcite and vision-language models @xcite , showcasing its versatility and effectiveness across various tasks. unlike applying the same parameters to every token as in dense transformers, the guiding design principle of sparse transformers is to selectively activate a subset of parameters for each token @xcite . specifically, mixture-of-experts layers operate by incorporating routers and making them learn to dynamically direct tokens to specific parameters, referred to as experts @xcite . this sparsity routing addresses the scaling issues of dense transformers while maintaining a constant number of computational operations. since routing is central to the mixture-of-experts paradigm, most ongoing research is dedicated to identifying and relieving various challenges associated with unstable gates @xcite and representation collapse @xcite @xcite . other research examined routing patterns @xcite @xcite to assess how effectively a sparse transformer can leverage its diverse set of experts. by tracing routing decisions across expert layers, @xcite discovered that expert assignments are less uniform among encoder layers than decoder layers and that meaningful specialization manifests primarily in syntactic properties rather than high-level semantics. @xcite further corroborated that routing is predominantly based on token identities and positions, regardless of context. this finding was termed context-independent expert specialization and justified by two observations: (1) tokens are routed to only a few fixed experts, and (2) consecutive token positions prefer similar experts. ## background mixture-of-experts (moe) has a long history in machine learning, dating back to the principle of adaptive mixtures of local experts @xcite . @xcite recently introduced sparselygated layers by extending the mixture-of-experts paradigm with techniques for conditional computation @xcite . by taking advantage of conditional computation, mixture-of-experts layers enable to scale up the number of trainable parameters while maintaining computational costs. building on transformer models @xcite , sparse mixture-of-experts layers can be interleaved with dense layers @xcite or upcycled from dense layers @xcite . sparse layers typically consists of a router and a fixed number of experts that are structurally identical to standard feed-forward neural networks. the router is responsible for assigning inputs to experts. each input is projected from its hidden state to the set of experts by multiplication with the router weights, which are learned jointly with the other network parameters. to produce a gradient for the router, the output of the computation is weighted by the corresponding probability of the assignment, since this probability is differentiable. this experts-as-a-layer approach dynamically activates a fixed subset of experts, ensuring that the number of floating-point operations remain constant, regardless of the total number of experts. to receive sufficient gradients for learning the router weights, @xcite conjectured that sparse mixture-of-experts layers require top-2 routing. as such, most implementations of sparse layers rely on two-way routing @xcite . however, this assumption is challenged by stable modifications for top-1 @xcite and adaptive topk routing @xcite , which allows variable expert assignment based on token complexity. to promote a balanced distribution of workload, @xcite defined a fixed expert capacity, which limits the number of tokens each expert can be assigned. the expert capacity is typically specified in the form of a hyperparameter, which acts as a multiplier factor for the expected number of tokens that would be assigned to each expert under a perfect uniform distribution. if the number of tokens assigned to an expert is not enough to fill its capacity, its set of tokens is padded to fill the remaining slots. if the number of tokens assigned to an expert overflows its capacity, the extra tokens are dropped. @xcite addressed the token dropout issue by reformulating the computation in terms of block-sparse operations that efficiently handle the dynamism present in sparse layers. since routing determines the token-expert assignments and thus dictates how effectively a model can leverage its set of experts, it is of central importance for the mixture-of-experts paradigm. there are two common classes of assignment algorithms for sparse layers: token choice in which tokens are dispatched to top-ranked experts and expert choice in which experts select the top-ranked tokens. token choice. the most common routing strategy is token choice @xcite @xcite , in which routing decisions are made by greedily selecting the topscoring experts for each token after projecting their hidden states to the number of experts. however, the greedy nature of this routing strategy suffers from notorious load imbalance issues that may cause the routers to collapse because experts that are assigned zero tokens no longer receive gradient updates @xcite . to encourage routers to make balanced token-expert assignments, additional adjustments such as noisy gating @xcite and imposing an auxiliary load balancing loss @xcite are required. @xcite developed a soft routing strategy with full differentiability that fills the capacity of experts using a weighted average of tokens. this provides a balanced and dropless mechanism for token-expert assignment. compared to the learning-to-route paradigm for routers @xcite , an alternative strategy is to reformulate the routing algorithm as a linear assignment problem that maximizes token-expert affinity @xcite or to eliminate the necessity for routers: stochastic routing @xcite leverages a consistency regularized loss for stochastic assignment, whereas deterministic hashing @xcite employs a parameter-free assignment algorithm that routes tokens by hashing. expert choice. rather than directing tokens to top-scoring experts, expert choice as proposed by @xcite has experts independently selecting top-scoring tokens, which guarantees perfect load balancing and allows for flexible allocation. ## methodology to illuminate the dynamics of routing with respect to context, we need to detail a sparsely-gated language model and the measurement to assess the degree of sensitivity within the sparse layers. we employ the switch @xcite transformer model, a sparsely-gated variant of the t5 @xcite ) sequence-to-sequence model, trained on a span corruption objective. this objective involves recovering variable-length contiguous segments masked in text, promoting a deeper understanding of contextual information compared to autoregressive models with dense layers @xcite and sparse layers @xcite . the architecture of the switch transformer consists of an encoder and a decoder, each comprising six sparse layers that alternate between dense and sparse configurations. each sparse layer contains a variable number of total experts in {8, 16, 32, 64, 128}, with a single active expert, where its assignment is managed through token choice routing combined with a load balancing loss. the choice of the switch transformer model is driven by its variable configurations of experts and its simple routing strategy. by tracing token-expert assignments in the sparselygated layers foot_0 , we can examine the sensitivity of the routing to similarity and the surrounding context. ## findings to examine how consistently sparsely-gated transformers route words based on context, we calculate the similarity between the distributions of experts for word pairs and correlate them with human judgments. we interpret strong correlation coefficients as context sensitivity. unless otherwise noted, we average the routing similarity across sparse layers. ## conclusion given the claims surrounding the factors influencing routing decisions in sparsely-gated mixture-ofexperts language models @xcite , we provide valuable insights into the influence of similarity and context. while similarity, encapsulated by token identities, form a stable basis for routing decisions, contextual cues provide an additional layer of refinement. however, the varying impact of context on the encoder and decoder reveals different sensitivities within the model components. the encoder demonstrates a strong ability to assign words in similar contexts consistently, revealing a high sensitivity to contextual cues, especially for configurations with many experts per sparse layer. the response of the decoder to context is poorer and more variable. this variability indicates instabilities in the utilization of context with respect to the number of experts. since our study demonstrates that context plays a significant role in routing, we hope that our approach sparks research on other linguistic properties and their influence on routing decisions, e.g., the influence of (affixal) negation @xcite or the consistency of routing for multi-word expressions @xcite .