L ogit R outer: a novel Attention variant for reducing Myopic Routing in Mixture of Experts (2025)

Abstract: Mixture of Experts (MoEs) have emerged as strong alternatives to traditional transformers, offering significant advantages in terms of training and inference efficiency. At the core of this architecture lies the router, responsible for selecting which experts are activated for each token. However, despite these advances, routing mechanisms continue to face stability challenges that the basic architecture fails to fully address. One such issue is Myopic Routing, where each token determines its route independently, without considering the routing decisions made for other tokens. To address this limitation, the LogitAttention mechanism is introduced—a variant of traditional attention—and, building upon it, the LogitRouter, a novel routing architecture that incorporates contextual information about the routing of other tokens. Due to budget constraints, a set of simple experiments is designed to obtain preliminary evidence of performance trends. These experiments are empirically validated on established benchmarks such as BoolQ, MMLU, and ARC. Finally, the work concludes with an in-depth discussion of architectural variants, applicability, limitations, and future directions, which aims to support continued research in this area.

## introduction large language models (llms) @xcite have represented one of the most significant breakthroughs in artificial intelligence in recent years. since the introduction of the transformer architecture @xcite , llms have continuously improved in both performance and capabilities, largely driven by their remarkable scaling laws @xcite . however, these advances have come at the cost of substantial energy and computational requirements @xcite , which has prompted the search for more efficient architectures capable of achieving results comparable to those of traditional transformers. in this context, mixture of experts (moe) @xcite @xcite has recently emerged as one of the most prominent architectures, serving as the foundation for many current state-of-the-art models @xcite @xcite . the main advantage of moe lies in its ability to reduce overall computational cost-during both training and inference-without causing a significant drop in final performance. this is achieved by replacing the traditional feed-forward network (ffn) sublayers, which typically hold the majority of a model's parameters, with a sparse moe layer. a sparse moe layer (or simply moe layer) is composed of multiple smaller ffn sublayers that share the same architecture (typically swiglu @xcite ) and are referred to as experts. these experts operate in parallel. to avoid activating all the parameters in the layer, a specialized component called the router is introduced. the router selectively activates and assigns weights to only a subset of experts, leaving the remaining ones unused. since the early works in this area @xcite @xcite , the conventional design of the router consists of a linear sub-layer, followed by a top-k selection function that activates only the k experts with the highest logit values. a softmax activation is then applied to ensure that the final weights follow a smooth probability distribution. this approach reduces the total computational load and, consequently, the energy and financial costs of running the network @xcite @xcite , at the expense of a minor (and often negligible) drop in performance. this setup naturally gives rise to a trade-off between efficiency and performance: scaling laws indicate that activating more parameters typically leads to improved outcomes, which means that mixture of experts (moe) models aim to find a balance point where further gains in performance do not justify the associated increase in cost @xcite . this balance is precisely the responsibility of the router: given a pool of experts, it decides which ones to activate and which to skip, seeking the optimal compromise between performance and efficiency. nevertheless, the standard router architecture still relies on a relatively basic formulation. while it has proven effective, it is far from definitive. as a result, a significant portion of current moe research focuses on designing better routers, as this component is not only central to the architecture but also the only element that represents a true innovation relative to standard transformers. literature gap. in this context, the problem known as myopic routing is identified. this issue arises when tokens are routed without considering the routing decisions of other tokens. to address this limitation, the logitattention mechanism is introduced, along with the logitrouter architecture-a novel type of router that mitigates this problem by learning to update routing decisions a priori using contextual information derived from the routing of other tokens. this architecture is evaluated through fine-tuning on the sft models of olmoe @xcite , and the resulting models are subsequently compared against the original ones across various benchmarks. contributions. the contributions of this work can be summarised as follows: the structure of this work is as follows: section 2 discusses the limitations of current moe architec-tures, the solutions proposed in the literature, and the motivation behind the present proposal. next, section 3 introduces logitattention, logitrouter, and their respective variants. section 4 presents the experimental setup, while the results and their analysis are detailed in section 5. following this, section 6 outlines potential directions and methods for extending and completing the research. finally, section 7 provides a concise summary of the overall contributions of the work. ## related work since gshard @xcite first applied mixture of experts (moe) to transformers, numerous routing architectures have been proposed @xcite @xcite , most of which have been evaluated primarily through empirical means. however, it was glam @xcite that first demonstrated moe's ability to outperform traditional transformers on more complex tasks. still, this promising architecture soon revealed several training instabilities. load balancing loss. one challenge inherited from early lstms (hochreiter and schmidhuber, 1997) is the load balancing loss @xcite , where some experts are over-utilized while others are underused, leading to suboptimal performance for the latter and ultimately for the model as a whole. gshard @xcite addressed this issue by introducing an auxiliary loss function-auxiliary load balancing loss-which remains widely used today. another effective approach was introduced in switch transformers @xcite , which challenged the then-common assumption that at least two experts must be activated to enable meaningful performance comparison. instead, they adopted a top-1 gating function while significantly increasing the number of experts, achieving better load distribution. another notable proposal is expert choice routing @xcite , which assigns a fixed number of tokens to each expert, rather than assigning a fixed number of experts to each token. representation collapse. one of the earliest empirically observed problems was the well-known representation collapse @xcite , where experts fail to utilise their full representational capacity. to mitigate this, the router z-loss auxiliary loss function was introduced @xcite , and it remains in use across various models. dynamic routing. beyond empirical observations, several theoretical issues have also been identified. a key concern is that all tokens typically activate the same number of experts, and thus the same number of parameters, contradicting the intuition that different tokens may require different levels of computational effort to reach optimal predictions. several strategies have been proposed to dynamically allocate the number of experts per token: routing fluctuation. dynamic routing, while promising, introduces new challenges @xcite . one such issue is routing fluctuation @xcite , where tokens repeatedly select different experts throughout training. as a result, many experts are updated at each step, but only a subset is used during inference, leading to poor optimization @xcite . to address this, (s)moe @xcite leverages token relationships captured through attention mechanisms to inform expert selection more effectively. however, unlike the aforementioned approaches that mainly focus on the gating mechanism or expert configurations, some recent proposals explore alternative sublayer architectures: ## experimental setup to ensure a degree of robustness in the results, two models were selected as the basis of the experiments: https://huggingface. co/allenai/olmoe-1b-7b-0924-sft and https://huggingface.co/allenai/ olmoe-1b-7b-0125-sft @xcite . these are, to date, the only competitive mixture of experts (moe) models small enough (around 7b parameters) to allow for low-budget training, while providing full access to their weights, training data, and evaluation details. for simplicity, we refer to them as olmoe-0924 and olmoe-0125. the reason for choosing their supervised fine-tuning (sft @xcite ) versions is that the sft datasets used in both models ( https: //huggingface.co/datasets/allenai/ tulu-v3.1-mix-preview-4096-olmoe and https://huggingface.co/datasets/allenai/ tulu-3-sft-olmo-2-mixture , respectively) are significantly smaller than their pretraining datasets. as a result, fine-tuning is faster and more cost-efficient. based on this, the following experiments were conducted: logitrouter models. first, starting from the olmoe-0924 and olmoe-0125 checkpoints, we replaced their routers with logitrouter modules. we then performed fine-tuning on both models using their original sft datasets, training only the router parameters while keeping all other parameters frozen. the fine-tuning procedure followed the hyperparameters described in the original paper @xcite , with the following exceptions: since logitrouter is a generalisation of the routing architecture used in olmoe (see 3.2), we initialized it as follows: w l := i n , w v := 0 n×n , and w r was initialized using the original router weights of the base models, ensuring that the initial behaviour remained identical. the matrices w q and w k were initialized using a standard truncated normal distribution @xcite . the resulting models are referred to as olmoe-0924-logit and olmoe-0125-logit, respectively. however, this experimental setup introduces two distinct sources of variation: (i) the new router architecture, and (ii) the new post-training strategy involving fine-tuning on the original sft datasets while freezing all non-router parameters. base router models. to isolate these factors, we conducted a second set of experiments applying only the post-training strategy, without modifying the routing architecture. in other words, using again olmoe-0924 and olmoe-0125 as base models, we repeated the fine-tuning process exactly as described above, but without replacing the original routers. to ensure a fair comparison under identical conditions, all hyperparameters were kept the same. these models are denoted as olmoe-0924-base and olmoe-0125-base. ## results the models were evaluated across three different benchmarks: arc (including arc-challenge and arc-easy), mmlu @xcite , and boolq @xcite , following the mcf strategy outlined in the olmes protocol @xcite . evaluations were conducted using varying numbers of few-shot examples (0, 1, 2, and 5), and the final performance metric for each model was defined as the maximum score achieved across these few-shot settings for olmoe-0125-base. two separate analyses were conducted: one to compare the performance of the original models with that of the base models, and another to compare the base models with the logit models. the first analysis aims to evaluate the impact of the previously defined post-training strategy on the final model, while the second assesses how changes in architecture affect performance. a summary of the results is presented in post-training strategy. the rows corresponding to the base models show how their performance differs from that of the original models. overall, it is evident that performance tends to degrade when these models undergo post-training using this strategy. in 5 out of 8 experiments, the post-trained models performed worse, and in 3 of those cases, the performance drop represented the largest absolute differences observed between the two versions. this suggests that there is no consistent improvement trend resulting from the post-training process. this outcome is expected, as the router becomes decoupled from the rest of the network under this setup-a configuration that is not commonly used in practice. several studies have shown that this decoupling may lead to network instability @xcite . moreover, it is important to emphasise that this particular experiment, although useful in this work for comparison purposes with the final models, does not constitute a genuine model optimisation procedure. rather, it corresponds to an arbitrary post-training step. the only potential "improvement" could be attributed to training the model for an additional epoch. however, if the original models did not undergo this extra epoch, it is likely because, after careful hyperparameter tuning and analysis, two epochs were deemed optimal. in fact, one can think of a well-trained and optimised llm as having reached a local optimum in the search space @xcite , meaning that any arbitrary modification is likely to degrade performance. therefore, the fact that performance actually improved over the original model in 3 out of 8 benchmarks suggests that there are still better optima to be found. this also ensures that the subsequent comparison with log-itrouter is fair and cannot be solely attributed to a drop in performance caused by this variable. logitrouter architecture. the rows corresponding to the logit models display their performance differences relative to the base models. this comparison is particularly relevant because it isolates the architectural effect of the router on performance. the results indicate a general trend toward improved performance, as only 2 out of the 8 experiments showed a performance drop-and in both cases, the decrease was minor. thus, we can reasonably conclude that switching the router architecture to a logitrouter produces a positive effect. in fact, in more than half of the evaluations (5 out of 8), the logit models even outperformed the original models. for instance, in arc-c, performance increased in both comparisons despite the initial performance drop observed in the base models. ## next steps given the constraints of a limited budget, most of the work remains to be completed. therefore, all the steps that need to be followed are outlined below, in order to assess the actual impact of log-itrouter and the other ideas previously introduced. ## conclusion this work introduced logitrouter, a novel routing architecture for mixture of experts, designed to incorporate contextual information into the token routing process. this approach addresses the newly identified myopic routing problem by enriching the routing decisions with broader contextual cues. despite limitations in budget the results indicate a promising trend. in particular, there were notable improvements on benchmarks such as arc-c compared to the baseline model, suggesting strong potential for this architecture in future developments. additionally, the paper provides a thorough discussion of each stage of the research process-ranging from the proposal and design of variants to the experiments and results-and outlines several promising directions for future work.