UMoE: Unifying Attention and FFN with Shared Experts
Authors: Yuanhang Yang
Chaozheng Wang
Jing Li
Venue: NeurIPS 2025 spotlight ()
Keywords: Mixture-of-Experts
MoE
attention
Primary area: deep_learning
TL;DR: A novel MoE architecture that extends mixture-of-experts to both attention and feed-forward layers with unified expert designs and attention-FFN parameter sharing.

Abstract: Sparse Mixture of Experts (MoE) architectures have emerged as a promising approach for scaling Transformer models. While initial works primarily incorporated MoE into feed-forward network (FFN) layers, recent studies have explored extending the MoE paradigm to attention layers to enhance model performance. However, existing attention-based MoE layers require specialized implementations and demonstrate suboptimal performance compared to their FFN-based counterparts. In this paper, we aim to unify MoE designs in attention and FFN layers by introducing a novel reformulation of the attention mechanism, that reveals an underlying FFN-like structure within attention modules. Our proposed architecture, UMoE, achieves superior performance through attention-based MoE layers while enabling efficient parameter sharing between FFN and attention components.

UMoE: Unifying Attention and FFN with Shared Experts
Yuanhang Yang', Chaozheng Wang?, Jing Li
Institute of Science Tokyo, Tokyo, Japan
The Chinese University of Hong Kong, Hong Kong, China Hong Kong Polytechmic University, Hong Kong, China
yang-y ea2c0m. isst.is...... jc. czwang23@cse cuhk, edu. hk jing-amelia amelia. 1i@polyy......mm edu hk
Abstract
Sparse Mixture of Experts (MoE) architectures have emerged as promising ap- proach for scaling Transformer models. While initial works primarily incorporated MoE into feed-forward network (FFN) layers, recent studies have explored ex- tending the MoE paradigm attention layers enhance model performance. However, existing attention baased MoE layers require specialized implementations and demonstrate suboptimal performance compared their FFN- IFN-based counter- parts. In this paper, we aim to unify MoE designs in attention and FFN layers by introducing novel reformulation of the attention mechanism, that reveals an underlying FFN-like structure within attention modules. Our proposed architecture, UMoE, achieves superior performance through attention based MoE layers while enabling efficient parameter sharing between FFN and attention components.
Introduction
However, this scaling advantage comes with substantial computational costs, making continued scaling in- (MoE) architectures have emerged as a promising
Scaling plays crucial role in advancing the capabilities of large language models 22 3
creasingly impractical. Sparse Mixture-of- solution by selectively activating only a subset of model parameters termed experts- for each input 4 5 6 [7] This approach effectively decouples model size from computational cost, enabling efficient scaling with minimal overhead.
Recent work has demonstrated the effectiveness of MoE in Transformer architectures B particularly when applied to feed-forward neural network (FFN) layers. Building on this success, several studies have explored extending MoE to attention layers [12 13] [4], indicating potential for performance gains through attention scaling. Despite the potential, we find that existing MoE attention layers demonstrate suboptimal performance compared to FFN-MoE approaches, when provided with similar computational and parametric budgets. This performance gap challenges the practical utility of attention- Preetion-MoE architectures, parameters allocated to scaling attention layers might be more effectively utilized for scaling FFNs instead.
10
We identify two distinctions between iitention-MoE and FFN-MoE implementations that likely account for the observed performance differential: (1) the different expert design between attention and FFN layers, and (2) attention-MoE's necessity to compromise the expressiveness of vanilla attention mechanisms to accommodate sparse computation [12]. Motivated by these observations, we investigate a compelling question: can reformulate attention to reveal an underlying structure compatible with the same expert design as FFN layers, without compromising the expressive power
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Residual Stream
a) Vanilla Attention
FFN MoE
Add
Expert
Expert
Expert
KV Cache
Mix
Router
Attention MoE
Pre- Pre-Mixing Attention
Add
Fixed Expert
Expert
Export
Expert
Taken
Token
Mix
Cache
Mix
Mix
Mix
Router
Figure 1: Illustration of a UMoE layer, which incorporates MoE into both FFN and attention modules with shared experts. The primary dis- tinction between attention-Me MoE and FFN-MoE lies in an additional token mixing operation.
Figure 2: Two formulations of the multi-head attention mechanism. (a) Vanilla attention inter- leaves mixing operations with value and output projections. (b) Pre-mixing attention performs token mixing prior to projections.
of the attention mechanism? This a challenging question due to the inherent complexity attention mechanisms, including multiple projections and softmax calculations, which fundamentally differ from the straightforward two-matrix multiplication pattern of FFNs.
To bridge this structural gap. we reformulate the attention mechanism to reveal its underlying FFN-like structure. Our reformulation decomposes attention into two sequential operations: token mixing and token-wise expert processing. The token-wise expert processing, consisting of two consecutive matrix multiplications, can be implemented as an FFN with a small intermediate size. This implementation
naturally aligns with recent advances in fine- ane-grained FFN expert design c 15 3 enabling unified expert architectures and parameter sharing across both attention and FFN layers.
Based on this insight, we introduce UMoE, a unified MoE architecture that abstracts Transformer layers into three fundamental components: experts, token mixing operations, and roulers, shown in Fig. I The experts. implemented as standard two-layer FFNs, serve as the primary components for token processing and knowledge storage. The token mixing operations facilitate contextual information exchange through weighted summation of tokens. Routers are employed to dynamically dispatch tokens the most relevant experts enable sparse computation. UMoE, the distinction between the FFN and attention layers lies solely in the expert inputs: FFN layers process tokens independently, while attention layers process tokens simultaneously through weighted summation. This unified design not only simplifies the architecture but also enables parameter-efffcient scaling through expert sharing berween attention and FFN components.
To evaluate the effectiveness of UMoE, we conduct extensive experiments across various model sizes and tasks, including pre-training and zero-shot evaluations. With the reformulated attention mechanism, the attention-basee MoE layers of UMoE match or exceed the performance of previous FFN-based MoE layers. Moreover, by sharing parameters across attention and FFN modules, UMoE achieves superior performance in fully MoE architectures while maintaining the same parameter count. We also present a detailed routing analysis of UMoE, revealing expert specialization patterns across modules, with higher-ranked experts demonstrating interpretable attention patterns. Our code available at https: //git/h// cow/ yal.k/nnnkkkuaa
Related Work
Sparse Cixture-of-Experte (MoE). Sparse Mixture-of-Experts (MoE) models have gained in- creasing attention for their ability scale model capacity while maintaining computational effi-

---

The core component of these models is the sparsely activated MoE sub-layer, which selectively activates different parameter subsets for different inputs. recent Transformer- based implementations, MoE architectures primarily replace feed-forward network (FFN) layers with MoE sub-layers. Each MoE. layer consists of collection of experts, denoted as [B, where each expert E, is implemented as an FFN. Tokens are routed to subset of experts through routing mechanism, with the top-k k router [ being the most prevalent approach. Despite advances in routing mechanisms [66] 17 18 19 [1], the top-k router remains widely adopted due to its simplicity and robust performance Ω For given token æ E Rd, where d the hidden dimension, and a trainable weight matrix W, R Nxd the top-k k router computes the probability distribution experts as:
ciency
p = smax W.))
The set of top-k experts T is then selected based on p. where [7| k. Each expert processes the token independently and the final output of the MoF layer is computed as the weighted combination of these k experts' outputs:
y= (.E,
(2)
where each expert is implemented as an FFN with two matrices and non-linear activation function.
MoE for Attention. Several recent approaches have explored extending the MoE paradigm to attention layers in Transformers I2 with a primary focus on expert design. Because attention layers lack the consecutive matrix multiplication pattern found in FFNs, these approaches necessitate expert designs that differ from FFN-MoE models. The offtturr-o--Attnttiio (MoA) 2 propose to conceptualize individual attention heads as experts, scaling attention layers by increasing the number of attention heads. However, introducing sparsity into attention layers presents significant challenge: query vectors computed by specific expert (or head) require corresponding key and value vectors from the same expert, necessitating identical expert activation across all tokens. To address this constraint, MoA implements distinct query and output projections per head while maintaining shared key and value projections across attention heads.
SwitchHead [13] presents an alternative approach to implementing the MoE paradigm in attention layers. Rather than treating entire attention heads as experts, SwitchHead designates individual projection matrices within heads as experts. A straightforward implementation maintains four separate MoE sub-layers per head for query, key, value, and output projections. While scaling all projections yields performance improvements, empirical results show that value and output projections benefit most significantly from scaling.
In contrast to these approaches, UMoE unifies attention-MoE and FFN-MoE through a novel refor- mulation of the multi-head attention mechanism, enabling the shared expert design and parameters across both attention and FFN layers.
Other Related Work. Several studies have explored connections between MoE and attention from different perspectives. MoH (20) proposes using MoE for pruning attention heads in LLMs by continuing pre-training with a routing function. During inference, certain output projections (W.), viewed as experts, are selectively skipped based on routing decisions. Taking different approach, MH-MoE 21 incorporates concepts from multi-heas attention enhance -Faee MoE models. Instead of routing original input tokens to experts, MH-MoB decomposes each token into multiple low-dimensional sub-tokens, which are then processed in parallel by diverse sets of experts.
Method
The attention mechanism is the core of Transformers (22) processing token hidden states to capture contextual relationships. However, its structure differs from FFN layers, which complicates the unifi- cation of MoE designs across both modules. In this section, we present two alternative formulations of attention, pre-mixing and post-mixing, that reveal an inherent FFN-like structure within attention layers. Based on these formulations, we introduce novel MoE architecture, UMoE.
3

---

def UMoELayer (x, X); x: [1, d], X: [n, d]
Attention MoE indices, probs = TopXRouter (x) # Assign token x to Experta
residual_x= co X W_k
q_shared = W_q
for i, in zip(indices, probs):
q q-shared
[_ali] w_b[i]
# and (the hidden states are shared across experts. y Attention (Q=q, K-K, V-X) residual_x += (x)erta[illy) residual_x
FPN MoE
indices, probs TopXRouter (x) # Assign token x to Experts residual_x c. for 1, in zip(indices, probs):
residual_x += p Experta(i](x) return residual_x
Figure 3: Implementation details of UMoE layer. The input consists of a sequence X containing n token hidden states and x representing the final hidden state. For simplicity, this implementation focuses computing the output for the last token.
3.1 Formulations of Attention
Preliminaries. Consider a sequence of token hidden states X E Raxd where n the sequence length and d is the hidden dimension. In multi-head attention, each token attends all other tokens in the sequence through query, key, and value projections. For single token x (e.g., the last token in the sequence for simplicity), its attention output is computed as:
q W
K
XWR, k: V XW,..
qK
a = softmax
0 a
Vd
where W 9. W E dx and W. E _dxd. are learnable matrices, respectively, and R" is the attention weight. To enhance representation capacity, this process repeated h times in parallel, and the outputs combined:
y [oi; 02; : W..
where W, E Rddd,dd projects the concatenated outputs back to the original dimension d.
Pre-Mixing Formulation. While multi-head attention is typically expressed using concatenation, it can be equivalently expressed as a sum of per-head outputs, which helps reveal its connection to FFN layers. By decomposing W. into small matrices Wi RRd xd along the feature dimension, we can express the output as:
h
y
O W
Σ (.W:)
(6)
h Σ (W
(7)
This reformulation provides two distinct interpretations, as shown in Fig. 2
Eq. 6 The conventional view where value vectors are first aggregated then projected back into the hidden space with an output projection.

---

Eq. A new interpretation where token hidden states are first aggregated into contextualized representations, i.e., weighted averages of all tokens, before being processed by the value (W") and output (W") projections. We term this formulation pre-mixing attention.
While both interpretations yield same outputs, the pre- pre-mixing formulation enables the grouping of Wo and W,. This grouping reveals that pre-mixing attention exhibits a two-layer structure analogous to FFN modules, which can be implemented linear FFN with no activation function.
Post-Mixing Formulation. Alternatively, we can rearrange the computation as:
y a. (XW:
(8)
In this formulation, token hidden states are transformed by two successive projections independently for each token, before being aggregated using the attention weights.
3.2 UMoE
By grouping W. and W.. both pre-mixing and post-mixing attention can be naturally interpreted as MoE architecture, aligning with established FFN-MoE practices. Using pre-mixing attention as an example, let the expert E(x)= == xW. W. The multi-head attention can then be reformulated as:
Σ X))
(9)
By increasing the number of experts and introducing routing mechanism, such as top-k router, we derive MoE architecture, denoted as UMoE-Att. The output of a UMoE-Att layer is:
where T is the set of activated experts. y = P,E,((a,)),
(10)
iET
Referring to Eq. we observe that the primary distinction between FFN-MoE layers and UMoE-Att layers lies in their expert inputs: FFN experts operate on individual token hidden states æ, while attention experts process weighted combinations all token hidden states. This reveals a relationship: FFN-MoE layers can be interpreted specialized case of pre-mixing attention layers where the attention matrix is constrained to an identity matrix, limiting each token to self-attention only.
Fully MoE Architecture. Both the experts in UMoB-Att and the FFN layers of Transformer consist of two consecutive matrices. While attention layer experts utilize a relatively small intermediate size (d,), FFN layers typically employ larger dimensions. Recent advances in FFN-MoE models suggest the efficacy of using FFN layers with reduced intermediate sizes as experts (6. 15 []. This insight enables the direct adoption of experts in attention layers for FFN layers, resulting in a fully MoE architecture, denoted as UMoE. Fig. [ fillustrates the architecture of a UMoE layer. where the MoE paradigm is applied to both FFN and attention layers using shared expert set. Notably, to facilitate parameter sharing, experts are implemented as two-layer FFNs with an intermediate size of d, and incorporate non-linear activation function between matrix multiplications.
Pre-mixing Implementation. The token mixing operation pre-mixing attention is a weighted summation over token hidden states, which can be implemented as vanilla attention, accepting Q K, matrices as input and producing an output matrix. Each token generates distinct query vectors for different experts, while values (hidden states) and their associated keys are shared across experts. To generate expert-dependent queries for input tokens, each expert requires an additional query projection matrix, leading notable increase in parameters. To mitigate the parameter count disparity with existing MoE models, where experts typically comprise two matrices, we employ low-rank matrices (23] for query projection within UMoE experts. For given token æ. the query for expert i is computed as:
9: W æ
(11)
where the first term is shared across all experts, while the second term is expert-specific with unique parameters, Wi E Rddr and Wi, E Rr for each expert. Fig. presents the pseudo-code of a
UMoE layer.
5

---

Putting It All Together. As illustrated in Fig. UMoE integrates three key components: (1) experts implemented as fine- aine-grained FFNs with dual low-rank query projection matrices, (2) pre-mixing attention mechanism utilizing shared keys and values across experts, and (3) the top-k router for expert selection. It is noteworthy that while MoA [12] also shares keys and values across experts, the values of MoA are the results after applying value linear transformation to the input token hidden states. In contrast, the values of UMoE directly refer to the input token hidden states.
3.3 Discussion
Vanilla Attention vs Pre-mixing Attention. Vanilla attention and pre-mixing attention differ in terms of KV cache requirements and computational complexity. During inference, vanilla attention requires caching multiple keys and values per token, whereas pre-mixing attention requires only one key and hidden state per token. While grouped-query attention (GQA) is commonly adopted to reduce the KV cache size in attention layers, it cannot be combined with pre-mixing attention, since the latter already maintains only single key-value pair per token. Instead, multi-head latent attention (MLA) 8 promising alternative to GQA, can be applied by introducing a nown-projection to the hidden states before the token-mixinn mixing operation. Regarding the computation, while vanilla attention performs weighted summation over low-dimensional value vectors, pre-mixing attention operates on input token hidden states, introducing modest increase in computational complexity. This modest increase, however, becomes increasingly negligible as models scale to larger dimensions, effectively amortizing the additional computational overhead. A detailed comparative analysis is presented in Additionally, the abstract formulation of UMoE opens avenues for future research to explore more computationally efficient token mixing alternatives, such as linear attention mechanisms
Table 8 I.II.
24 25
Pre-mixing Attention vs Post-mixing Attention. As il- lustrated in Fig. post-mixing attention processes individ- ual tokens through experts prior to mixing. The architec- tural distinction between pre-mixing and post-mixing vari- ants represents different perspectives on token paremeeerreettr interactions in attention layers. Recent interpretability studies have drawn parallels between the two-matrix mul- tiplication pattern of FFNs and associative memory mod- ules, where value neurons, i.e. columns of the second ma- trix in FFNs, are retrieved by inputs 26 27 28). Within this framework, pre-mixing attention leverages token mix-
Mi
Cached Dutputs
Inland
Figure 4: Post-Mixing Attention.
ing generate contextualized inputs for precise retrieval. In contrast, post-mixing attention can be conceptualized as an ensemble of independent retrievals executed by preceding tokens. Our prelimi- nary experiments (A.2) demonstrates significant performance advantage of pre-mixing attention over its post-mixing counterpart. This observation suggests that generating contextualized inputs for token-paree Aaenmmaareetter interactions more effectively aligns with the principles of attention mechanisms.
Experiments
4.1 Setup
Datasets. We conduct language modeling pretraining two datasets: FineWeb-Edu 100B 29 and Wikitext- 30). FineWee-Edu has shown superior data efficiency when evaluated on knowledge- intensive benchmarks. Wikitext-103 consisting of approximately 100M tokens, is a smaller corpus that has been widely adopted in previous studies [3]. We apply the LLaMA tokenizer [11 with a 32K vocabulary size both datasets. The zero-shot performance of models trained FebeWebEEdu is evaluated using the na-eevlluaiion-hareeee
framework [32]
Baselines. We compare UMoE against three categories of baselines: dense models, FFN-MoB models with fine-grained experts g 3), and attention-MoE models (specifically MoA [12 and
SwitchHead [33). Attention-MoE models are configured with identical expert parametery We
'Swtthheead represents exception, it requires the number of experts to be divisible by the number of attention heads.

---

Table 1: Comparison of Dense and Sparse Mixture-of-Expeer (MoE) Models for Language Modeling. In MoE models, x denotes experts per layer with size in #Tottal' columns, while B in #Activee columns indicates the number of experts activated per token. Gray entries in '#ottal' columns indicate shared parameters between attention and FFN modules in UMoE models. UMoB-Att refers to UMoE variants with MoE only applied to attention modules.
Attention
FFN
PPL ())
Model
Params
MACs
#Total
#Active
#Total
#Active
Fineweb
Wikitext
Base Models
Dense
134 M
768
768
3072
3072
25.79
30.41
525
Fine-grained FFN-MoE
535 M
768
768
192 128
192
21.19
27.94
530G
MoA
525
192x
x92x4
3072
22.28
27.57
486G
SwitchHead
533 M
192 116
192
3072
3072
22.91
29.47
542
UMoE-Att
547
192
x92x4
27.45
611
UMoE
540 M
x92
x92
192x
192 16
20.44
26.67
616G
Large Models
Dense
1.1 B
2048
2048
5632
5632
17.53
25.46
4.59 T
Fine-grained FFN-MoE
3.8B
2048
2048
512x
512x
16.09
25.47
4.61 T
MoA
3.6B
512x57
512x4
5632
5632
16.72
25.14
3.99T
SwitchHead
3.7B
612x60
512x4
5632
5632
16.48
27.24
4.62T
UMoE-Att
3,8B
512x57
512x4
5632
5632
16.03
25.53
4.73T
UMoE
3.6B
512x64
512x4
512x
512x
15.95
25.44
4.75T
implement all MoE models with fixed expert per layer, following recommendations for optimal model performance 8 33). (MoE adopts the pre-mixing attention mechanism, which consistently outperforms post-mixing variants.
Experimental Setup. UMoE is implemented decoder-only Transformer with rotary position embedding (34), following deepseek-Mon [6]. The load balancing loss proposed by Switch Trans- former [5] is adopted to encourage balanced load across experts. The experts in experiments are implemented as two-layer MLPs to ensure fair comparison with the baselines
We evaluate model configurations: The base models comprise 12 layers with hidden size of 768, while the large models consist of 24 layers with hidden size of 2048. These configurations yield dense models with 134M and 1.1B parameters, respectively. MoE variants replace all attention or FFN layers with MoE layers. Due to computational constraints, unless otherwise specified, models pretrained on 50B tokens from FineWeb-Edu with batch size of 1024. For Wikitext- 103, following Csordás et al. 133, models are trained for 100k steps, though models typically overfit within 20k steps. Detailed hyperparameters are provided in A.4
4.2 Comparison with Baselines
Results. From Table I we observe that UMoE shows consistent superiority across different model sizes and datasets. In the base model regime, UMoE achieves the best performance. Notably, the attention-only variant of UMoE exhibits substantial improvements approaches. Even without parameter sharing, UMoE-Att establishes itself as a compelling alternative traditional FFN-based MoF models. The parameter sharing mechanism between attention and FFN modules further enhances the effectiveness without increasing the total parameter count. Despite equalizing the number of activated experts across all baselines, discrepancies due to different attention layer implementations, as measured by MACs (multipicaation accumulation operation). We additionally perform a MAC-matched comparison by increasing the number of activated experts of baseline models. Table 2 shows that even under comparable computational constraints, UMoE method achieves the lowest perplexity.
previous attention- based
observe subtle computational
UMME is designed as a flexible framework agnostic to the expert implementation. Using more advanced gated MLP variants, such as SwiGLU, may lead to improved performance, as the gating mechanism can enhance the expressiveness of both attention and beed-forward layers.

---

Table 2: MAC-matched comparison for base models by increasing the number of activated experts of baseline models.
Table 3: Parameter sharing strategies. / in- dicates shared components between modules while X indicates separate components.
Model
MACs
Active Params
PPL (
Component
UMoE
Fineweb
Wikitext
Fixed Experts
FFN-MoE
617G
768 + x2
20.80
27.39
Router
X
MoA
6
192x6 3072
22.00
27.63
SwischHead
Params
540
M
UMtE-Ait
611G
192x4 3072
20.81
27.45
UMoE
616G
192x4
192x 16
20.44
26.67
PPL.
22.82
23.11
23.05
23.02
Table 4: Zero-shot accuracy on downstream tasks.
The best score is marked in bold.
Model
Params
HellaSwag
PIQA
ARC-E
ARC-C
RACE
Lambada
MMLU
Wino
Avg.
Base Models
Dense
134
33.58
62.35
46.09
24.74
27.75
19.97
24.8
49.8
36.
MoA
525
37.82
65.58
26.19
28.83
22.33
38.49
SwitchHHeud
533 M
37.19
66.12
50.55
26.59
28.14
21.73
25.2
50.9
38.30
FFN-MoE
535 M
39.69
66.
52.95
26.71
29.76
23.46
25.3
52.1
39.55
UMoE (AU))
547 M
40.72
67.36
51.77
27.82
29.76
23.66
25.9
52.5
39.94
UMoB
540
41.28
66.65
51.86
29.01
28.71
23.77
26.6
40.06
Large Models
Dense
1.1
48.45
60.22
58.85
32.17
33.11
31.75
27.4
53.8
44.35
MoA
3.6B
50.61
70.28
61.47
33.25
32.38
33.15
25.6
54.7
45.55
SwitchHead
B.7
51.90
70.83
62.34
33.27
33.66
55.7
46.27
FFN-MoE
3.8
52.74
71.52
64.23
35.67
33.30
34.00
29.2
56.3
47.12
UMoR (Att)
3.8
53.20
71.44
63.30
34.39
32.82
34.78
29.3
57.4
47.08
UMoE
3.6
53.17
72.47
64.23
35.75
32.44
35.32
30.4
56.9
47.58
In larger-scale models, UMoE maintains its competitive advantage. While MoA shows marginally better performance on Wikitext-106, this result may not fully reflect model capabilities given the rela- tively small size of Wikitext- 103 (100M tokens) compared to the model scale. Following established practice [16], we further report validation perplexity. Fig. shoows that UMoE demonstrates faster convergence and lower validation perplexity compared baselines, indicating enhanced modeling capabilities. This superior performance translates to downstream tasks, with Table showing UMoE consistently achieving the highest average zero-shot accuracy across diverse tasks.
Efficiency. Following Zhang et al. 12 Jin et al. 20 employ MACS 5 an efficiency metric, as remains independent of hardware implementations. As shown in Table I the pre-mixing attention introduces modest computational overhead, resulting in approximately 1.17x slowdown for base models. However, this slowdown becomes increasingly negligible as models scale up: in large models, UMoE introduces only 1.03x slowdown compared to the dense baseline. This favorable scaling behavior arises from the different growth rates in computational complexity: expert processing scales quadratically with hidden dimension, while token aggregation in attention layers scales linearly.
4.3 Ablations
We conducted ablation experiments using base models trained on FineWeb-Edu with 20B tokens.
Parameter Sharing Analysis. We investigated various sharing strategies across FFN and attention layers for fixed experts (6] 333 and routers. As shown in Table 3 all configurations achieved comparable perplexity. Our default configuration, which employs separate fixed experts and routers across FFN and attention layers, yielded the optimal perplexity.
Expert Allocation. As suggested in Section 3.2 FFN-MoE layers can be interpreted as specialized case of pre-mixing attention layers with an identity matrix as attention matrix. This interpretation raises question: does UMoE perform better when allocating more experts to attention layers rather than FFN layers? According to Table we observe an trend when gradually shifting expert
MACss is measured using the DeepSpeed Flops Profiler.

---

Table 5: Impact of expert allocation between Attention and FFN layers (total experts 20).
Table 6: Effect of Activation Functions in Ex- pert Modules. Vidiicates experts with activa- tion functions while x indicates experts with- out activation functions.
# Expert
Model
PPL
Attention
FFN
Model
Act. Function
PPL
16
22.82
V
22.82
8
12
22.63
UMoE
X
24.43
UMoB
12
22.44
16
22.50
23.37
UMoE (Att)
20
0
21.75
x
23.99
Table 7: Top tokens for selected experts in the last attention and FFN layer of UMoE.
Expert ID
Top Tokens in Attention Layer
Top Tokens in FFN Layer
Each
This
Every
Each
This
This
This
Every
Each
Another
10
Film video lab. film
Video
Tag
Font. ISBN twitter.
DNS
46
The The
the
the
_Our
_a
his _my. Your Your
relatively _extremelyy very
Very
allocation from FFN to attention modules while maintaining total activated expert size of 20. The model achieves its best perplexity when all experts are allocated to attention layers. This finding provides empirical evidence supporting a specialized form of attention, with the attention mechanism exhibiting greater expressiveness. However, increasing the number of attention experts introduces substantial computational overhead due to token mixing operations. Future research could explore efficient attention alternatives within the attention MoE framework.
theoretical interpretation that FFN layers function as
Activation Function. Table presents our investigation into the impact of activation functions UMoE. The results demon- strate that incorporating activation functions between matrix multiplications within experts consistently improves model per- formance, reinforcing the crucial role of non-linearity in deep learning architectures. Notably, while the removal of activation functions reduces the experts to pure linear transformations both FFN and attention modules, UMoE remains trainable. We attribute this robustness to the preserved non-linearity from token mixing operations and layer normalization. Neverthe- less, the consistent performance degradation underscores the importance of activation functions in model expressiveness, particularly in the context of shared expert architectures.
MoA
Parplesity
FFN-MoE
UNCE
1050
0030
1950
1003
12008 14005 10000
Training Steps
Dense
Training Loss
UNGE
4.4 Expert Specialization
266
Table presents the routing patterns in the final layer UMoE, where experts are shared between attention and FFN modules while maintaining distinct routers. Notably, certain token cat- egories consistently route to the same experts across both mod- ules, as evidenced by experts and 46. Expert 3 consistently processes determiners, while expert 46 specializes in demon- strative pronouns.
0200
10000
Training
Figure S: Best valid PPL (top) and training loss (bottom) on Wikitext.
The analysis also reveals divergent specialization patterns that highlight the complexity of shared expert architectures. notable example expert 64, which exhibits distinct specializations: processing consecutive punctuation marks in the attention layer while handling degree adverbs in the FFN layer. This phenomenon suggests that shared experts can develop multiple specializations, potentially leading more efficient parameter utilization. However,

---

it also raises important questions about potential knowledge conflicts within individual experts, indicating promising directions for future research in routing mechanism design for shared expert architectures.
We also provide analysis the attention maps of UMoE in A.5 which confirms that higher-ranked experts show more focused attention distributions on relevant tokens compared to lower-ranked ones.
5
Conclusion
The paper proposes UMoE, a novel architecture that unifies MoE designs for attention and FFN layers. The key insight is reformulation of the attention mechanism that allows the value and output projections to be grouped into FFN-like experts. This unification enables parameter sharing across attention and FFN layers, resulting in fully MoE architecture that improves performance without introducing additional parameters. The paper presents extensive experiments demonstrating UMoE superiority over existing MoE architectures in terms of perplexity on language modeling datasets and accuracy on zero-shot tasks.
As for future work, we are looking at replacing the token mixing mechanism with more efficient alternatives to enable scaling up the number of activated experts in attention layers. In addition, we are also interested in investigating architectures that unify attention and FFN into a single layer, given our finding that FFN layers function as specialized case of attention with reduced expressiveness. Exploring parameter sharing across different Transformer layers 33 36 is also a promising direction.
References
[1] Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B. Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. CoRR, fas/2001.083611 2020. URL /tt::://axxv..e/ www.
08361
[2] Elias Frantar, Carlos Riquelme Ruiz, Neil Houlsby, Dan Alistarh, and Utku Evci. Scaling laws for sparsely-connectee foundation models. In The Twelfth International Conference Learning Representattioss, 2024.
[3] Jan Ludziejewski, Jakub Krajewski, Kamil Adamczewski, Maciej Pióro, Michat Krutul, Szymon Antoniak, Kamil Ciebiera, Krystian Król, Tomasz Odrzygóćdž, Piotr Sankowski, Marek Cygan, and Sebastian Jaszezur. Scaling laws for fine-grained mixture of experts. Proceedings of the 41st International Conference on Machine Learning, volume 235, pages 33270-33288, PMLR, 21-27 Jul 2024.
[4] Noam Shazeer, *Azalia Mirhoseini, *Krzysztof Maziarz, Andy Davis, Quoe Le, Geoffree Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gatee mixture-of- experts layer. In International Conference on Learning Representations, 2017.
[5] William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: scaling to trillion parameter models with simple and efficient sparsity. J. Mach. Learn. Res., 23(1), January 2022.
ISSN 1532-4435.
[6] Damai Dai. Chengqi Deng, Chenggang Zhao, R.x. Xu, Huazuo Gao, Deli Chen, Jiashi Li. Wangding Zeng, Xingkai Yu, Y. Wu, Zhenda Xie, Y.k. Li, Panpan Huang, Fuli Luo, Chong Ruan, Zhifang Sui, and Wenfeng Liang. DeepSeekMoF: Towards ultimate expert specializa- tion in mixture-of-experts language models. In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 1280- 1297, Bangkok, Thailand, August 2024. Association for Computational Linguistics. doi: URL https:///////// htattp:://aclaatoolgg.. .erpggg..rrg/// org/2024// acl long. 70/
[7] Fuzhao Xue, Zian Zheng, Yao Fu, Jinjie Ni, Zangwei Zheng, Wangchunshu Zhou, and Yang You. Openmoe: an early effort on open mixture-of of-eeeerss language models. In Proceedings of the 41st International Conference on Machine Learning, ICML'24. JMLR.org, 2024.
[8] DeepSeek-Al, Aixin Liu, Bei Feng, Bing Xue, Bingxuan Wang, Bochao Wu, Chengda Lu, Chenggang Zhao, Chengqi Deng, Chenyu Zhang, Chong Ruan, Damai Dai, Daya Guo, Dejian Yang, Deli Chen, Dongjie Ji, Erhang Li, Fangyun Lin, Fucong Dai, Fuli Luo, Guangbo Hao, Guanting Chen, Guowei Li, H. Zhang, Han Bao, Hanwei Xu, Haocheng Wang, Haowei Zhang, Honghui Ding, Huajian Xin, Huazuo Gao, Hui Li, Hui Qu, J. L. Cai, Jian Liang. Jianzhong Guo,
10

---

Jiaqi Ni, Jiashi Li, Jiawei Wang, Jin Chen, Jingchang Chen, Jingyang Yuan, Junjie Qiu, Junlong Li, Junxiao Song, Kai Dong. Kai Hu, Kaige Gao, Kang Guan, Kexin Huang, Kuai Yu, Lean Wang, Lecong Zhang, Lei Xu, Leyi Xia, Liang Zhao, Litong Wang. Liyue Zhang. Meng Li, Miaojun Wang, Mingchuan Zhang, Minghua Zhang, Minghui Tang. Mingming Li, Ning Tian, Panpan Huang, Peiyi Wang, Peng Zhang, Qiancheng Wang, Qihao Zhu, Qinyu Chen, Qiushi Du, R. J. Chen, R. L. Jin, Ruiqi Ge, Ruisong Zhang, Ruizhe Pan, Runji Wang. Runxin Xu, Ruoyu Zhang, Ruyi Chen, S. S. Li, Shanghao Lu, Shangyan Zhou, Shanhuang Chen, Shaoqing Wu, Shengfeng Ye, Shengfeng Ye, Shirong Ma. Shiyu Wang, Shuang Zhou, Shuiping Yu, Shunfeng Zhou, Shuting Pan, T. Wang, Tao Yun, Tian Pei, Tianyu Sun, W. L. Xiao, Wangding Zeng, Wanjia Zhao, Wei An, Wen Liu, Wenfeng Liang, Wenjun Gao, Wenqin Yu, Wentao Zhang, X. Q. Li, Xiangyue Jin, Xianzu Wang. Xiao Bi, Xiaodong Liu, Xiaohan Wang. Xiaojin Shen, Xiaokang Chen, Xiaokang Zhang, Xiaosha Chen, Xiaotao Nie, Xiaowen Sun, Xiaoxiang Wang, Xin Cheng. Xin Liu, Xin Xie, Xingchao Liu, Xingkai Yu, Xinnan Song, Xinxia Shan, Xinyi Zhou, Xinyu Yang. Xinyuan Li, Xuecheng Su, Xuheng Lin, Y. K. Li, Y. Q. Wang, Y. X. Wei, Y. X. Zhu, Yang Zhang, Yanhong Xu, Yanhong Xu, Yanping Huang, Yao Li, Yao Zhao, Yaofeng Sun, Yaohui Li, Yaohui Wang, Yi Yu, Yi Zheng, Yichao Zhang. Yifan Shi, Yiliang Xiong, Ying He, Ying Tang. Yishi Piao, Yisong Wang, Yixuan Tan, Yiyang Ma. Yiyuan Liu, Yongqiang Guo, Yu Wu, Yuan Ou, Yuchen Zhu, Yuduan Wang, Yue Gong, Yuheng Zou, Yujia He, Yukun Zha, Yunfan Xiong, Yunxian Ma, Yuting Yan. Yuxiang Luo, Yuxiang You, Yuxuan Liu, Yuyang Zhou, Z. F. Wu. Z. Z. Ren, Zehui Ren. Zhangli Sha, Zhe Fu, Zhean Xu, Zhen Huang, Zhen Zhang, Zhenda Xie, Zhengyan Zhang. Zhewen Hao, Zhibin Gou, Zhicheng Ma, Zhigang Yan, Zhihong Shao, Zhipeng Xu, Zhiyu Wu, Zhongyu Zhang, Zhuoshu Li, Zihui Gu, Zijia Zhu, Zijun Liu, Zilin Li, Ziwei Xie, Ziyang Song, Ziyi Gao, and Zizheng Pan. Deepseek-v3 technical report, 2024.
[9] Niklas Muennighoff, Luca Soldaini, Dirk Groeneveld, Kyle Lo, Jacob Morrison, Sewon Min, Weijia Shi, Pete Walsh, Oyvind Tafjord, Nathan Lambert, Yuling Gu, Shane Arora, Akshita Bhagia, Dustin Schwenk, David Wadden, Alexander Wettig, Binyuan Hui, Tim Dettmers, Douwe Kiela, Ali Farhadi, Noah A. Smith, Pang Wei Koh, Amanpreet Singh, and Hannaneh Hajishirzi. Olmoe: Open mixture-of-experts language models. CoRR, abs/2409.0200 002660, 2024.
[10] Albert Q. Jiang, Alexandre Sablayrolles, Antoine Roux, Arthur Mensch, Blanche Savary, Chris Bamford, Devendra Singh Chaplot, Diego de Las Casas, Emma Bou Hanna, Florian Bressand, Gianna Lengyel. Guillaume Bour, Guillaume Lample, Lélio Renard Lavaud, Lucile Saulnier, Marie-Anne Lachaux, Pierre Stock, Sandeep Subramanian, Sophia Yang. Szymon Antoniak, Teven Le Scao, Théophile Gervet, Thibaut Lavril, Thomas Wang, Timothée Lacroix, and William EI Sayed. Mixtral of experts. //doi.org//f 10. 000.448550/rr/ arxiv. 2401. 0008884088
CoRR, abs/2401, 040088 2024. URL https
[11] Barret Zoph, Irwan Bello, Sameer Kumar, Nan Du, Yanping Huang, Jeff Dean, Noam Shazeer, and William Fedus. St-moe: Designing stable and transferable sparse expert models. arXiv preprint arXhv:2202.089000 2022.
[12] Xiaofeng Zhang, Yikang Shen, Zeyu Huang, Jie Zhou, Wenge Rong, and Zhang Xiong. Mixture of attention heads: Selecting attention heads per token. In Proceedings of the 2022 Conference on Empirical Methods Natural Language Processing, EMNLP 2022, Abu Dhabi, United Arab Emirates, December 7-11, 2022, pages 4150 4162. Association for Computational Linguistics, 2022. doi: 10. 11t.:/22222...PPPAAAI 2022. emn. www.emmmnlmmmmmmm.mmmmmmmmmmm..m 278
18653/v1/
[13] Róbert Csordás, Piotr Piekos, Kazuki Irie, and Jürgen Schmidhuber. Switchhead: Accelerating transformers with mixture-ote mixture-of-experts attention. In The Thirty- Annual Conference on Neural Information Processing Systems, 2024. URL httpp:://ppeeneeveew..ee wwwp::/wwweefooo/memaam 80S$169GAZ
[14] Xiaoye Qu, Daize Dong, Xuyang Hu, Tong Zhu, Weigao Sun, and Yu Cheng. Llama-moe v2: Exploring sparsity of llama from perspective of mixture-of-er -feeeerrss with post-training. CORR, abs/2411.15708. 2024.
[15] Yuanhang Yang, Shiyi Qi, Wenchao Gu, Chaozheng Wang, Cuiyun Gao, and Zenglin Xu. XMoE: Sparse models with fine-grained and adaptive expert selection. In Findings of the Association for Computational Linguistics: ACI. 2024, pages 11664-11674, Association for Computational Linguistics, 2024. doi: 10. http://...nnnnnsaadllo.o.

---

[16] Stephen Roller, Sainbayar Sukhbuatar, Arthur Szlam, and Jason E Weston. Hash layers for large sparse models. In A. Beygelzimer, Y. Dauphin, Liang. and J. Wortman Vaughan, editors. Advances in Neural Information Processing Systems, 2021. URL forum?id--MgDDW55UUMM
net/
[17] Mike Lewis, Shruti Bhosale, Tim Dettmers, Naman Goyal, and Luke Zettlemoyer. BASE layers: Simplifying training of large, sparse models. In Proceedings of the 38th International Conference on Machine Learning, ICMI. 2021, 18-24 July 2021, Virtual Event, volume 139 of Proceedings of Machine Learning Research, pages 6265-6274, PMLR, 2021. htth
//proceedings.oo mlr. wwwpppeeessveee wttps///111//eees2aaa.eees
[18] Yangi Zhou, Tao Lei, Hanxiao Liu, Nan Du, Yanping Huang. Vincent Y Zhao, Andrew M. Dai, Zhifeng Chen, Quoc V Le, and James Laudon. Mixture-of-experts with expert choice routing. In Advances in Neural Information Processing Systems, 2022. URL
net/forumiid jdJo dJ
[19] Quzhe Huang, Zhenwei An, Nan Zhuang, Mingxu Tao, Chen Zhang, Yang Jin, Kun Xu, Kun Xu, Liwei Chen, Songfang Huang, and Yansong Feng. Harder task needs more experts: Dynamic routing in MoE models. In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics ( (Volume 1: Long Papers), pages 12883- 12895, Bangkok, Thailand, August 2024. Association for Computational Linguistics. URL https: attpss//aclannhhooog. org 2024 acl-long.............. 6967
[20] Peng Jin, Bo Zhu, Li Yuan, and Shuicheng Yan. Moh: Multi-head attention as mixture-of- head attention. CoRR, abs/2410. 1842, 2024. doi: 10.48550/ARXIV.44100 1842. URL
https://ddi..rge: htrg//..8850//// arxiv. 2410. 11842
[21] Xun Wu, Shaohan Huang, Wenhui Wang, Shuming Ma, Li Dong, and Furu Wei. Multi-head mixture-of-experts. In The Tiihty Annuul Conference Neural Information Processing www./eeeetnuiieee dttp///rmmmdddddd-//iii
Systems, 2024. URL http:tttttpppereerrreee
[22] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez. E ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Information Processing Systems. Curran Associates, Inc., 2017.
[23] Edward Hu, yelong shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. LoRA: Low-rank adaptation of large language models. In International Conference on Learning Representations, 2022. |tLpptts://oppeeeerii
[24] Songlin Yang, Bailin Wang, Yu Zhang, Yikang Shen, and Yoon Kim. Parallelizing linear transformers with the delta rule over sequence length. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024. /ttRatttss///peeeviiiie wtwp:eewwwwiie..ee/fotueee ner/forum? id=y8Rm4VNRPH
[25] Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. CoRR, abs/23 12.07552, 2023. https:///// /ttp::///oooooo///0/// 48550/arXiv
wttp://..........///rr...... 2312. 00752
[26] Mor Geva, Roei Schuster, Jonathan Berant, and Omer Levy. Transformer feed-forward layers are key-value memories. In Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, Association for Computational Linguistics, 2021. URL org/2021. www.l-main.nn.....eem.mmmmmmmamm 446/
[27] Zeyu Liu, Tim Dettmers, Xi Lin, Veselin Stoyanov, and Xian Li. Towards a unified view of sparse feed-forward network in pretraining large language model. In Proceedings of the 2023 Conference on Empirical Methods Natural Language Processing, pages 15038-15061, Singapore, December 2023. Association for Computational Linguistics. doi: 10. 10.1188/1/2222.. emnlp-main.933 URL (ttl:stts///acaattttat
org 2023. emnlp-manieee 930/
[28] Zeping Yu and Sophia Ananiadou. Neuron-level knowledge attribution in large language models.
In Proceedings of the 2024 Conference on Empirical Methods in Natural Language Processing, pages 3267-3280. Association for Computational Linguistics, 2024. doi: 10. (0.18633/1/2202. emnip-main. 1911 URL hww.hhttps:/// /tt.ttttss///ccaaaataaa http///aaaaattooggyyo org//024 emnlp-main.......m 191/
[29] Guilherme Penedo, Hynek Kydlícek, Loubna Ben Allal, Anton Lozhkov, Margaret Mitchell, Colin Raffel, Leandro von Werra, and Thomas Wolf. The fineweb datasets: Decanting the web for the finest text data at scale. CORR, abs/2406. 7557, 2024. httpss:tts:: //d::i///oo/ org/10
48550/ aaiiv 2406 17557
12

---

[30] Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models. In Sth International Conference on Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Conference Track Proceedings. OpenReview.net, 2017.
[31] Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, Dan Bikel, Lukas Blecher. Cristian Canton-Ferre Ferrer, Moya Chen, Guillem Cucurull. David Esiobu, Jude Fernan- des. Jeremy Fu, Wenyin Fu, Brian Fuller, Cynthia Gao, Vedanuj Goswami, Naman Goyal, Anthony Hartshorn, Saghar Hosseini, Rui Hou, Hakan Inan, Marcin Kardas, Viktor Kerkez, Ma- dian Khabsa, Isabel Kloumann, Artem Korenev, Punit Singh Koura, Marie-Anne Lachaux, Thibaut Lavril, Jenya Lee, Diana Liskovich, Yinghai Lu, Yuning Mao, Xavier Martinet, Todor Mihaylov, Pushkar Mishra, Igor Molybog, Yixin Nie, Andrew Poulton, Jeremy Reizen- stein, Rashi Rungta, Kalyan Saladi, Alan Schelten, Ruan Silva, Eric Michael Smith, Ranjan Subramanian, Xiaoqing Ellen Tan, Binh Tang, Ross Taylor, Adina Williams, Jian Xiang Kuan, Puxin Xu, Zheng Yan, Iliyan Zarov, Yuchen Zhang, Angela Fan, Melanie Kambadur, Sharan Narang, Aurélien Rodriguez, Robert Stojnic, Sergey Edunov, and Thomas Scialom. Llama 2: Open foundation and fine-tuned chat models. CoRR, abs/2307.092888, 2023. URL attps://doi.org
2307 09288
[32]
Leo Gao, Jonathan Tow, Baber Abbasi, Stella Biderman, Sid Black, Anthony DiPofi, Charles Foster, Laurence Golding, Jeffrey Hsu, Alain Le Noac"h, Haonan Li, Kyle McDonell, Niklas Muennighoff, Chris Ociepa, Jason Phang, Laria Reynolds, Hailey Schoelkopf, Aviya Skowron, Lintang Sutawika, Eric Tang, Anish Thite, Ben Wang, Kevin Wang, and Andy Zou. A framework for few-shot language model evaluation, 07 2024. URL hhttps://zente org/records/ 12608602
[33] Samyam Rajbhandari, Conglong Li, Zhewei Yao, Minjia Zhang, Reza Yazdani Aminabadi, Ammar Ahmad Awan, Jeff Rasley, and Yuxiong He. Deepspeed-moce moo:: Advancing mixture-of- experts inference and training to power next-generation Al scale. In International Conference on Machine Learning. ICML 2022, 17-23 July 2022, Baltimore, Maryland, USA, volume 162 of Proceedings of Machine Learning Research, pages 18332-18346. PMLR, 2022.
[34] Jianlin Su, Murtadha H. M. Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding. Neurocomputing, 568: 127063, 2024. /ttpttts:///do.ooo org/ 1016/j .ewcoom.emmmmmmmm.. 2023. 127063.
[35] Mostafa Dehghani, Stephan Gouws, Oriol Vinyals, Jakob Uszkoreit, and Lukasz Kaiser. Univer- sal transformers. In International Conference on Learning Representations, ICLR, New Orleans, LA, USA, May 6-9, 2019. OpenReview.net, 2019. URL id-HyzdRiR9Y7
[36] Róbert Csordás, Kazuki Irie, Jürgen Schmidhuber, Christopher Potts, and Christopher D. Manning. Moeut: Mixture-of-experts universal transformers. In Advances in Neural Information Processing Systems, NeurIPS, 2024, URL http::/pppp/:rr .nis..nipp............. www.:niipppppppppe cwtpcccrppppeceeee filessffiee///eee filpsecc/pper 2024/
/aper/ html
13

---

Table 8: Computational complexity analysis of vanilla attention and pre- pre-mixing attention mechanisms. Entries in gray denote identical complexity terms between the two mechanisms. Here, N denotes sequence length, d represents hidden dimension, h indicates the number of attention heads (activated experts), da and d, are key and value dimensions respectively, and r the rank of query projection matrices in the pre-mixing attention.
Operation
Vanilla
Pre-mixing
Output Projection
(N - d. h)
O d, d h)
Value Projection
N dr d h)
O d h)
Key Projection
O (N dx d
O(N dg d)
Query Projection
O (N dx d.
O(N dg d + N (dk + d) .r h-))
QK Multiplication
(N2 d2 h)
O dy h)
Weighted Sum
O d,, h)
O -d)
A Technical Appendices and Supplementary Material
A.1 Complexity
Table presents detailed computational complexity comparison between vanilla attention and pre- mixing attention mechanisms. The key distinctions lies in two operations: key projection and weighted sum computation. In key projection, pre-mixing attention achieves lower complexity. However, this efficiency is partially offset in the weighted sum operation, where the complexity increases due to the full dimensional mixing. is worth mentioning that the computational complexity of weighted sum grows linearly with the hidden dimension, while the FFN computation grows quadratically. The computational overhead weighted sum becomes less significant as model size increases.
def PostMixingMof (X) x: [n. d]
### Attention MoE
# Independent Expert Processing
all_token,keya []
all
for thie_token in X: indices, probs TopKRouter (this_token)
for 1, p zindines,s, probs):
y = p Experts (tthis token) all htt.tookenootppt..paapapap append
[_k (thissstthtk all_token_keys append
# Mixing
(X W_q) unsuqeeze(1) # [n, 1, d_g] all_token res.een.ey..rrehapeeeee (n, k, d_q)
all_token cotoken_outuuss..o. reshape (n, k, d_v)
attn_outpur Attention(Q, K, v) # [n. 1, d_v]
FFN MoE
Implementation of UMoE layer based on post-mixing attention. X is sequence of token hidden states. In the attention layer, tokens are processed by their top-k experts independently. The output embeddings of all tokens are aggregated according the attention weights.
Figure 6:
14

---

A.2 Post-mixing vs Pre-mixing
Fig. presents the pseudo-code of UMoE layer based on post-mixing attention. We conducted preliminary experiments Wikitext- 103 and FineWeb-Edd datasets. As illustrated in Fig. MoE models incorporating pre-mixing attention demonstrate substantially superior performance compared to their post-mixing counterparts.
As discussed in Section 3.1 these two reformulations the attention mechanism are mathematically equivalent, owing to the absence of non-linear transformations within the matrix chain multiplication in attention heads. However, when grouping the value and output projections and implementing them as FFN with non-linear activation function, these formulations yield distinct outputs in the MoE layers. Fig. Billustratss MoE models implemented based on these two reformulations. Both approaches can be interpreted as extensions of conventional FFN-based MoE models. Specifically, the pre-mixing approach enhances FFN-based MoE models by contextualizing the inputs, while post-mixing attention enables MoE layers to incorporate other tokens' outputs in generating the final output. Future research could explore the synergistic combination of post- and pre-mixing approaches fully leverage contextual information.
Dense
Post-mixing
Post-mixing
Pre-mixing
Pre-mine Pre-mixing
Training Loss
Training Loss
15000 17500 20000
1000
2000
3000
4000
5000
6000
7000
8000
Training Steps
Training Steps
Figure 7: Loss curves of UMoE with attention MoE layers implemented post-mixing and pre- mixing attention, respectively. Models are trained on Wikitext-10s (left) and FineWeb-Eae (right).
Post-mixing
Taken
Mix
Cached Outputs
Fixed Expert
Token-0
Token-t
Expert
Export
Expert
Router
Pre-mixing
Fixed Expert
Expert
Expert
Expert
Cached Inputs
Token-t
Token Mix
Token
Router
Figure 8: Two implementations of UMoE based on pre-mixing and post-mixing attention, respectively.
15

---

Table 9: Comparison of training throughput and inference latency across models. Inference latency denotes the pre- reeliiiing time (1024 tokens) measured on a single H100 GPU. All results are averaged over multiple runs.
Model Type
Model
Model Size
Pre-filling Latency (8) Training Throughput (tokens/s)
GPT
134M
0.1376
126,508
FFN-MoR
535M
0.1974
74,415
MoA
525M
0.1843
90,394
Base Model
SwitchHead
533M
0.2253
69,374
UMoE-Att
547M
0,20066
76,538
UMoE
540M
0.2197
71,461
GPT
1.1B
0.2950
11,634
FFN-MoE
3.8B
0.4080
9,799
MoA
3.6B
0,3624
12,321
Large Model
SwitchHead
3.7B
0.4328
8,882
UMoE-Att
3.8B
0.4120
11,377
UMoE
3.6B
0.4392
9,194
A.3
Wall-clock Efficiency Analysis
Table 9 presents a quantitative comparison of end-to-end training throughput and inference latency across various model architectures. All experiments were conducted on single NVIDIA H100 GPU. Inference latency corresponds to the pre-filling time for sequences of 1024 tokens, while training throughput is measured in tokens per second.
Across both scales, all Mixture-of--xperts (MoE) variants exhibit comparable inference latency but remain slower than the dense bascline, despite their similar theoretical MAC counts. This discrepancy primarily arises from additional computational overheads introduced by expert routing and sparse expert execution. Among MoE models, SwitchHead and UMoE show slightly higher latency due to their architectural designs-Switchie -eaitcHeaa applies two MoE layers per attention head, whereas UMoE employs two MoE layers within each Transformer block.
These results underscore a persistent challenge in MoE-based architectures: while their theoretical efficiency is well established, current hardware and software infrastructures (e.g., routing kernels, expert parallelization) are not yet fully optimized realize these potential speedups. Consequently, our main paper reports MAC-based efficiency as the primary metric. We anticipate that advances in GPU kernel design and distributed parallelization will further mitigate the observed wall-clock inefficiencies in future implementations.
A.4 Hyperparameters
Table [0 and II give the parameters used for based and large models, respectively. It takes roughly week for pretraining base models on FineWeb-EEd datasets. MoA and SwitchHead utilize identical parameters as UMoB-Att, excluding the low-rank query projections. Table [12 details the hyperpa- rameters used during training. For the Wikitext-100 103 dataset, we adopt the same hyperparameter configuration as Csordás et al. [3]
A.5 Attention Analysis
We analyze the attention patterns in UMoE by visualizing expert-specific attention maps. While UMoE (Large) contains 64 experts per layer across layers, we focus on the top experts (ranked by router scores) to maintain tractability. Each expert utilizes its own query projection matrix, allowing us to compute attention maps regardless of activation status.
To investigate attention behavior, we examine two inputs:
"Context: William Shakespeare wrote the famous play Romeo and Juliet the late 16th century. Question: Who wrote Romeo and Juliet? The Answer is"
16

---

Table 10: Hyperparameters of Base Models, MoA and SwitchHead use the same hyperparameters as UMoE-Att.
Hyperparameter
Dense
FFN-MoB
UMoE-Att
UMoE
Context Length
1024
1024
1024
768
Number of Layers
12
12
12
12
Hidden Size
768
768
768
768
Attention Heads
4
4
4
4
FFN Size
3072
192
16
3072
192
16
Query (Key) Dimension
512
512
512
512
Value Dimension
192
192
Query Lora Rank
16
16
Number of MoE layers
12
12
12
Expert Size
192
192
192
Experts per MoE Layer
128
116
128
FFN Experts per Token
Attention Experts per Token
4
4
Table Hyperparameters of Large Models. MoA and SwitchHead use the same hyperparameters as UMoE-Att.
Hyperparameter
Dense
FFN-MoE
UMoE-Att
UMoE
Context Length
1024
1024
1024
1024
Number of Layers
24
24
Hidden Size
2048
2048
2048
2048
Attention Heads
4
FFN Size
5632
512 x
5632
512 x
Query (Key) Dimension
512
512
512
512
Value Dimension
512
512
Query Lora Rank
36
36
Number MoE layers
24
Expert Size
512
512
512
Experts per MoE Layer
57
FFN Experts per Token
11
11
Attention Experts per Token
4
"Context: Tokyo the capital city of Japan and has a population of over 37 million people in its metropolitan area. Question: What is the capital of Japan? The Answer is"
UMoE successfully predicted the correct answers for both inputs, even after removing the context. For the final token in each input, collected attention maps from the top 8 experts across all layers.
Fig. 9 presents attention maps for the final token prediction, revealing distinct patterns between higher and lower-ranked experts. Higher-ranked experts demonstrate more focused attention distributions that align with task requirements. For instance, in the Shakespeare question, Expert 0 and Expert _3 show pronounced attention weights task-critical tokens "who wrote". Similarly, for the Tokyo question, the top expert exhibits sophisticated attention patterns by focusing on key contextual elements like "Japan" and "capital". These observations suggest that the routing mechanism effec- tively identifies experts capable of extracting task-relevant information through specialized attention patterns.
To create a comprehensive visualization, we aggregated attention maps of layers by summing them, as shown in Fig. 10 The values on the left (e.g., EO: 3.83) represent the sum of router scores
17

---

Table 12: Training Hyperparameters on FineWeb-Edu and Wikitext- 103.
Hyperparameter
Fin-Web-Edu
Wikitext
Global Batch Size
1024
64
Learning Rate
4e-4
2.5e-4
Training Steps
50000
100000
LR Scheduler
cosine
cosine
Warmup Ratio
0.05
0.05
GPU
H100
H100
Expen_e
05
Doerr: 23: 0.006
Exper, 3: 0,047
Dioee 4if
Dosee_a 0,030
Dove 52:
Contee
(a)
100
Expart 0.35
Expart 3.000
Exper 23 0.00
Exper 0.066
Expar, 3.22
Exper 36 001
005
Toon
(b)
Figure 9: Representative attention maps. The heatmaps show the attention weights of the last token produced by top 8 experts, ranked by their router scores. (a) Attention patterns for the Shakespeare question, where higher-ranked experts (c.g., Expert_ Expert_322 demonstrate focused attention on question-relevant tokens. (b) Attention patterns for the Tokyo question, showing similar task-specific attention concentration among top experts.
received by experts at each rank position across all layers. The accumulated patterns reveal that higher-ranked experts (particularly EO and El) maintain more targeted attention distributions focused on auestion-relevan tokens, while lower-ranked experts tend to focus heavily on the initial token,
Notably, we observe minimal attention paid to answer tokens present in the context. This phenomenon aligns with the conceptualization of experts (two consecutive matrices) as key-value memory modules. where input serves as query. In other words, the output of attention layers is the composition of values, i.e., columns of the second matrix, the activated experts, rather than token hidden states. This suggests that the token mixing should focus building appropriate query for accurate compositions. Therefore, the last token should pay attention to the tokens relevant to the answer token, rather than the answer itself.
18

---

00:185
Token
(a)
10.3.41
23:0.00
24/0.4
25.0.5
26.0.5 6.54
0.45
(b)
Figure 10: Layer-wise accumulated attention weights across the model. The values on the left (e.g.. EO: 3.83) represent the sum of router scores for experts. Higher-ranked experts (E0, E1) consistently show more focused attention distributions compared to lower-ranked experts.
A.6 Limitations
One limitation of UMoE. is, as discussed in Section 4.2 the modest additional computational cost introduced by reformulating the attention mechanism in relatively small models. Additionally, the datasets used in this paper do not cover mathematics or code, given the scope of our work and limited computational resources. Considering recent research on MoE models' reasoning limitations, exploring how reasoning abilities scale with UMoE may provide valuable insights for future work.
19

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: In the abstract and introduction, claim that our proposed attention-MoE layers outperform previous attention-Moß methods and match or exceed the performance of FFN-MoE models.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
The paper provides a "Limitations" section A.6)..
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (c.g.. independence assumptions, noiseless settings, model cell-specfiiaatinn asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect the scope of the claims made, e.g. if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex, system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
20

---

Answer:
[NA]
Justification:
The paper does not include theoretical results.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proof's can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided not)?
Question:
Answer: [Yes]
Justification: This paper proposes novel model architecture. The details of experiments are provided in Section and A.4 Our code is also available for reproduction.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, descríbing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is new model (e-g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing verifying the results.
5. Open access to data and code
21

---

Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [Yes]
Justification:
The code is released. The datasets used by this paper are all publicly available.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.co cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https
//nips //ipss.cc/pill wttp:ccppppcccceeeeeeeeee..c.ett htapa//eguddeee/eeeeabiiissssss /Cc////oddiiisssseeeeeee ssioiolicyyyyyy for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Question:
Answer: [Yes]
Justification:
The details of experiments are provided Section and A.4
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [No]
Justification: Considering performing language modeling pretraining, it would be too computationally expensive for this given our limited computational budget.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
22

---

The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma enror bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are range (e.g. negative enror rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification: Section A.4 provides information on the time of training and Table provides the information of GPU.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).
9. Code of ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics htttps:/neerr/ httpa//ppp..ccpbb.a.c/cccaaaaaaa
Answer: [Yes]
Justification:
The research conducted in the paper conform with the NeurIPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. if there is special consid- eration due to laws or regulations their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: There is no societal impact of the work performed since the focus of this work is to propose a general architecture.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
23

---

Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), faireess considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied partieular applications, let alone deployments. However, if there is direct path any negative applications, the authors should point it out. For example, it is legitimate point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (c.g., gated release models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility MI).
11. Safeguards
Question:
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Answer: [NA]
Justification:
There is no model or data released.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: We cite the datasets and the code package used for model evaluation. The Fen-Web dataset is under ODC-By 1.0 license and the Wikitext-100 103 is released under CC BY-SA 3.0.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

---

If assets are released, the license, copyright information, and terms of use in the thcode. has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
package should be provided. For popular datasets, paperswi
batasets.papersstt
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors encouraged to reach out the asset's s creators.
13. New assets
Question: Are assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification:
We released our code for reproduction with document how to run it.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Bthics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
25

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology. scientific rigorousness, or originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification: LLMs are used only for writing.
Guidelines:
The answer NA means that the method development in this research does not involve LLMs as any important, original, or non-standard components.
Please refer to our LLM policy (https: /ttp::://neurps.... hatps////coooefeeeeec////////// for what should or should be described.
26


Reviews:
[{"review_id": "3NRVZ2T51X", "reviewer": "Reviewer_nJ81", "summary": "This paper presents UMoE, a novel sparse Mixture‑of‑Experts (MoE) Transformer architecture that unifies MoE mechanisms across both attention and feed‑forward network (FFN) layers. The authors reformulate attention to reveal an FFN‑like structure, enabling shared experts between attention and FFN modules. This design supports efficient parameter sharing, eliminating the need for specialized attention MoE implementations while maintaining high performance. Experiments on language modeling (e.g., WikiText and FineWeb‑Edu) show UMoE achieves comparable or superior perplexity.", "questions": "See Weaknesses.", "limitations": "Yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "**Strengths**:\n1. The paper introduces a novel perspective by unifying sparse MoE mechanisms in both attention and FFN layers, which simplifies model architecture.\n2. By showing that attention can be reformulated as a sequence of FFNs, UMoE enables expert sharing across attention and FFN.\n3. Experiments on language modeling (e.g., WikiText and FineWeb‑Edu) show UMoE achieves comparable or superior perplexity, with detailed routing analysis demonstrating expert specialization and interpretable attention behaviors.\n\n**Weaknesses**:\n1. Although the paper discusses theoretical FLOPs and parameter counts, it does not report end-to-end training/inference wall-clock time comparisons against baselines, which limits claims on practical efficiency.\n2. The experimental setup is limited in scope, making it difficult to assess the generalizability of the proposed method. In particular, the paper does not evaluate UMoE in more challenging or diverse scenarios such as large-scale pretraining tasks or long-context modeling, where the benefits of unified sparse routing might be more or less pronounced. Including such experiments would strengthen the claims about the method’s broad applicability.\n3. The inference performance of UMoE remains unclear and warrants further evaluation. Since language modeling relies on autoregressive generation, efficient inference typically requires KV caching. However, UMoE appears to involve caching intermediate activations (e.g., the X used in expert routing and weighted summation), which may increase both computational overhead and memory usage during inference. This overhead could be particularly problematic when compared to baselines employing optimized techniques such as Grouped Query Attention (GQA). Without benchmarking inference latency and memory footprint, it is difficult to assess UMoE’s practical viability in real-world deployment scenarios.\n4. The writing quality of the paper requires proofreading and refinement. \n5. The paper lacks sufficient explanation for the performance gains over other MoE-based attention methods. While empirical results are promising, the authors should provide deeper analysis—e.g., ablation studies, theoretical insights, or visualizations—to clarify why UMoE achieves better results.", "was_revised": true, "final_justification": "Thanks for the additional latency and throughput results. While the proposed UMoE achieves better perplexity at similar MACs, the increased training time and inference latency are still notable.\n\nI believe it would be more meaningful to compare model quality under equal training throughput or inference latency. This reflects real-world constraints more accurately. Since this limitation is common in current MoE work, I will keep my score unchanged.", "rebuttal": "Thanks a lot for your detailed review - we are glad you found our proposed UMoE novel, which eliminates the need for specialized attention MoE implementations while maintaining high performance.\n\nRegarding your Weaknesses:\n\n> **Weakness 1: end-to-end training/inference wall-clock time comparisons against baselines**\n\nThis is a very constructive suggestion! We implement the forward pass of experts in MoE layers using GPU kernels written in Triton. Below we report the training throughput (tokens/s) and inference latency (ms). Specifically, inference latency is the time required for pre-filling sequences with 1024 tokens. All measurements were conducted on a single H100.\n\n| **Base Model**  | **Model Size** | **Pre-filling latency (s)** | **Training throughput (token/s)** |\n| --------------- | -------------- | --------------------------- | --------------------------------- |\n| GPT             | 134 M          | **0.1376**                  | **126,508**                       |\n| FFN-MoE         | 535 M          | 0.1974                      | 74,415                            |\n| MoA             | 525 M          | 0.1843                      | 90,394                            |\n| SwitchHead      | 533 M          | 0.2253                      | 69,374                            |\n| UMoE-Att        | 547 M          | 0.2006                      | 76,538                            |\n| UMoE            | 540 M          | 0.2197                      | 71,461                            |\n| **Large Model** |                |                             |                                   |\n| GPT             | 1.1 B          | **0.2950**                  | 11,634                            |\n| FFN-MoE         | 3.8 B          | 0.4080                      | 9,799                             |\n| MoA             | 3.6 B          | 0.3624                      | **12,321**                        |\n| SwitchHead      | 3.7 B          | 0.4328                      | 8,882                             |\n| UMoE-Att        | 3.8 B          | 0.4120                      | 11,377                            |\n| UMoE            | 3.6 B          | 0.4392                      | 9,194                             |\n\nAs the results indicate, all MoE variants show quite similar latency while requiring significantly more time than the dense baseline despite having similar theoretical MACs.\n\nOur analysis reveals that both the routing mechanism and sparse expert computation within MoE layers introduce noticeable overhead. Specifically, SwitchHead and UMoE are slower than other baseline MoE models because SwitchHead implements two MoE layers per attention head while UMoE has two MoE layers per Transformer block.\n\nThese findings highlight a known challenge with MoE architectures: while theoretically efficient in computation, current hardware and software implementations don't fully realize these benefits in practice. Due to these limitations, we use MACs as the main efficiency metric in our work. We believe these practical speed gaps will narrow with future advancements in GPU kernels and parallelization strategies. We will add this discussion to our paper.\n\n> **Weakness 2: The experimental setup is limited in scope ...... In particular, the paper does not evaluate UMoE in more challenging or diverse scenarios such as large-scale pretraining tasks or long-context modeling**\n\nThank you for this suggestion! Our paper focuses on proposing a new attention-MoE structure, following the evaluation protocols of previous related work [1, 2] by conducting pretraining on general datasets. Your suggestions regarding larger-scale training and testing on long-context scenarios are valuable and could provide new insights, yet these are beyond the scope of our current study. Since these tasks are resource-intensive, we leave them for future work.\n\n> **Weakness 3: However, UMoE appears to involve caching intermediate activations (e.g., the X used in expert routing and weighted summation), which may increase both computational overhead and memory usage during inference.**\n\n> **This overhead could be particularly problematic when compared to baselines employing optimized techniques such as Grouped Query Attention (GQA)**. \n\n> **Without benchmarking inference latency and memory footprint, it is difficult to assess UMoE’s practical viability in real-world deployment scenarios.**\n\nThis is an insightful observation, especially regarding GQA! We addressed the inference latency question in W1. Here, we focus specifically on KV cache and memory footprint. The footprint is measured with ```torch.cuda.memory_allocated()```. We consider a scenario using a large model generating the next token based on 8192 tokens. All models have 4 attention heads.\n\n|             | footprint (GB) | Cache size (MB) | K Cache Shape  | V Cache Shape   |\n| ----------- | -------------- | --------------- | -------------- | --------------- |\n| FFN-MoE     | 9.75           | 1536.0 MB       | [4, 8192, 512] | [4, 8192, 512]  |\n| UMoE        | 8.21           | 960.0 MB        | [1, 8192, 512] | [1, 8192, 2048] |\n| FFN-MoE-GQA | 7.78           | 384.0 MB        | [1, 8192, 512] | [1, 8192, 512]  |\n| UMoE-MLA    | 7.74           | 384.0 MB        | [1, 8192, 512] | [1, 8192, 512]) |\n\nAs shown, **UMoE has a smaller KV cache size and memory footprint compared to vanilla attention** due to our pre-mixing attention design. We acknowledge that our proposed pre-mixing attention cannot utilize Grouped Query Attention (GQA) since we already have only one key and value, which cannot be further reduced. However, we can leverage Multi-head Latent Attention (MLA) [3], a promising alternative to GQA proposed by DeepSeek. MLA reduces memory footprint by decreasing KV cache dimensions.\n\nWe trained a base-size UMoE on Fineweb-Edu using 10B tokens with MLA, which resulted in only a slight degradation in performance (perplexity: 23.209 to 23.565), demonstrating MLA's effectiveness.\n\n> **Weakness 4: The writing quality of the paper requires proofreading and refinement.**\n\nWe appreciate this feedback and will thoroughly revise the manuscript, addressing writing issues and incorporating suggestions from all reviewers to improve clarity and presentation.\n\n> **Weakness 5: The paper lacks sufficient explanation for the performance gains.**\n\nThank you for this suggestion. The main distinction between MoE-based attention methods, including our UMoE, lies in expert design—specifically how to increase attention parameters most effectively. One possible explanation for our superior performance is that the two-matrix multiplication pattern in the attention layer resembles memory layers (as mentioned in Section 3.3). This suggests that matrices $W_v$ and $W_o$ should be scaled together, which previous work failed to accomplish. MoA didn't scale them together, and both MoA and SwitchHead break key-value correspondence during scaling. In the revised manuscript, we will expand Section 3.3 to provide a more comprehensive discussion.\n\n[1] Csordás, R., Piękos, P., Irie, K., and Schmidhuber, J. SwitchHead: Accelerating Transformers with Mixture-of-Experts Attention, NeurIPS 2024. https://openreview.net/forum?id=80SSl69GAz\n\n[2] Zhang, X., Shen, Y., Huang, Z., Zhou, J., Rong, W., and Xiong, Z. Mixture of Attention Heads: Selecting Attention Heads Per Token, EMNLP 2022. https://aclanthology.org/2022.emnlp-main.278/\n\n[3] DeepSeek-AI et al. DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model, 2024. https://arxiv.org/abs/2405.04434"}, {"review_id": "rD672w6rO7", "reviewer": "Reviewer_VfmG", "summary": "The authors propose a novel architecture that unifies MoE design to FFN and attention layers simultaneously, termed UMoE. This can potentially allow parameter sharing between FFN and attention experts. The authors build this architecture based on the observation that FFN layers are a specialized form of attention. They also provide LLM pre-training experiments to compare the performance of the different MoEs designs matched at a similar compute budget.", "questions": "1. See weaknesses W1, W2, W3\n2. In table 5, if attention is 0, does UMoE reduce to FFN-MoE? If not, please highlight the differences. Otherwise, this is a good insight to mention.", "limitations": "None", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n1. The subject addressed in the paper is worthy of investigation. Allowing for parameter sharing between attention and FFN layers with a a strong pre-training performance can be very interesting. \n2. Numerical experiments show that the author's proposed architecture can achieve a strong performance compared to SOTAs.\n3. The method is well-explained. The ablation study is interesting and insightful.\n\nWeaknesses:\n1. In line 230, the authors mention that the evaluation dataset can be small and not a reliable metric for perplexity evaluation, why don't they consider another dataset, like C4?\n2. Can the authors provide evaluation perplexity on other evaluation datasets to understand to what extent the performance of UMoE is generalizable?\n3. Can the authors provide results of Table 2 for the larger model as well?\n4. Further scaling to a ~3B dense model and other more hyperparameters for pre-training can improve the experimental results.", "was_revised": true, "final_justification": "My score would have been 4.5 but decided to recommend acceptance of the paper: 5.\nI think the proposed unified MoE architecture across attention and FFN layers is very interesting and novel --> opens new research directions for pre-training. This comes with good numerical performance. Additionally, throughput and memory footprint results are interesting.\nThis would represent a new baseline to compare against for pre-training tasks, hence my recommendation towards acceptance.\nThe results are interesting and potentially impactful but definitely not groundbreaking, the score shouldn't be higher.", "rebuttal": "We thank the reviewer for the time spent reviewing our work. We are grateful that the reviewer found our study insightful and interesting with strong performance! Below, we present the responses to the issues raised.\n\n> **Weakness 1: In line 230, the authors mention that the evaluation dataset can be small and not a reliable metric for perplexity evaluation, why don't they consider another dataset, like C4?**\n\nThank you for this question! Our work includes both training and evaluation on two datasets, FineWeb-Edu and WikiText-103. Our primary dataset is FineWeb-Edu, which is comparable in scale to C4. FineWeb-Edu is a large-scale, diverse dataset widely used in recent language modeling research. We believe the evaluation results on FineWeb-Edu provide the most reliable insights into our model's performance.\n\nWhile WikiText-103 is indeed smaller, we included it primarily for fair and direct comparison with the baseline models MoA and SwitchHead, which also reported results on this dataset. The limited size of WikiText-103 helps us understand model fitting speed (as shown in Figure 5). Evaluating on additional datasets like C4 would certainly provide further insights, which we consider valuable for future work.\n\n> **Weakness 2: Can the authors provide evaluation perplexity on other evaluation datasets to understand to what extent the performance of UMoE is generalizable?**\n\nWe appreciate this suggestion. To directly address the concern about generalizability, we conducted additional experiments on C4 to verify the reliability of our method. Due to computational constraints, we trained the models on 10B tokens with the following results:\n\n| Model    | Perplexity |\n| -------- | ---------- |\n| FFN-MoE  | 24.622     |\n| UMoE-Att | 24.116     |\n| UMoE     | **23.945** |\n\nAdditionally, we evaluated our models trained on FineWeb-Edu with C4's validation set:\n\n| Model      | Perplexity (↓) |\n| ---------- | -------------- |\n| FFN-MoE    | 28.22          |\n| MoA        | 29.35          |\n| SwitchHead | 28.60          |\n| UMoE-Att   | 28.13          |\n| UMoE       | **27.98**      |\n\nThese results confirm that our conclusions remain consistent across datasets: UMoE presents a strong alternative to previous FFN-MoE method.\n\n> **Weakness 3: Can the authors provide results of Table 2 for the larger model as well?**\n\nThank you for pointing this out. We apologize for this oversight; the table was prepared but inadvertently omitted from the appendix. The results for the larger model are as follows:\n\n| Model      | MACs   | Active Params      | Fineweb (↓) | Wikitext (↓) |\n| ---------- | ------ | ------------------ | ----------- | ------------ |\n| FFN-MoE    | 4.81 T | 2048 + 512 x 12    | 16.00       | 25.54        |\n| MoA        | 4.91 T | 512 x 6 + 5632     | 16.61       | 26.79        |\n| SwitchHead | 4.87 T | 640 x 4 + 5632     | 16.45       | 26.12        |\n| UMoE-Att   | 4.73 T | 512 x 4 + 5632     | 16.03       | 25.53        |\n| UMoE       | 4.75 T | 512 x 4 + 512 x 11 | **15.95**   | **25.44**    |\n\nThe results for the larger model are consistent with our other experiments, showing that UMoE consistently outperforms previous attention-MoE baselines even when they are allocated similar or greater computational budgets, while achieving performance on par with FFN-MoE.\n\n> **Weakness 4: Further scaling to a ~3B dense model and other more hyperparameters for pre-training can improve the experimental results.**\n\nWe agree with the reviewer that scaling our experiments to larger models (~3B parameters) and exploring a wider range of hyperparameters would be valuable for further demonstrating the practical effectiveness of UMoE. \n\nGiven our limited computational resources (requiring over two weeks to train models exceeding 1B parameters, with unstable resource availability), we focused on demonstrating the core concept within these constraints. Nevertheless, we believe that our findings with moderately-sized models are still significant and meaningful for the research community. UMoE presents an interesting and simple architecture that can inspire future work, and we hope this paper will encourage exploration of its scaling potential with more extensive resources.\n\n> **Question: In table 5, if attention is 0, does UMoE reduce to FFN-MoE? If not, please highlight the differences. Otherwise, this is a good insight to mention.**\n\nThank you for this insightful question!  To clarify, when the activated expert count (k) for attention is set to 0 in Table 5, UMoE does not precisely reduce to standard FFN-MoE.\n\nIn UMoE, we utilize a shared pool of experts for both attention and FFN layers, but with different numbers of activated experts (top-k) for each. If k=0 for attention, this would disable all attention mechanisms since no experts would be activated for attention computations. The model would become pure FFNs without any token interaction capabilities, functionally equivalent to a 1-gram language model predicting solely based on the current token position. Therefore, we did not train such a model configuration."}, {"review_id": "4LTCB8NR1V", "reviewer": "Reviewer_1ZC3", "summary": "This paper propose a new way to look at attention module of transformer architecture, and realize that part of the attention module could be viewed as a FFN expert. The author proposes to have a fully MoE-ed transformer layer -- UMoE. Through experiments, they show UMoE is able to perform better, given similar FLOPs.", "questions": "* Line 113, the notation is misleading --- notation doesn't agree with W_o; they should really be W_q^I\n* In Table 5, does the result imply we should just always skip FFN and only have attention; this is surprising because this seems to disagree with some existing work on importance of FFN.\n* Line 154, \"token hidden states\" is more precise; token embedding typically refers to the output of embedding matrix (even before position embedding, if any, is applied).\n* What does \"mixing\" in \"pre/post-mixing\" refers to? I thought a_i X is a mixing operation.", "limitations": "Might be useful to make a plot of PPL-MAC to show the Pareto-frontier.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Strength: \n* The attempt to move towards a fully MoE-ed layer is intriguing and interesting.\n* experimentation result is thorough; but might requires more description like, in main experiment, do author have MoE layer every 2 transformer layer; or in Table 5, is the expert allocation done on all layers.\n\nWeakness:\n* important explanation missing: from Eqn 11 and Figure 3, it seems W_k and W_q are fully activated to calculate the attention, which makes the saving not as appealing as it appears when the author talks about the attention is also being sparsely activated.\n* important inconsistency: Equation 11 disagree with figure 3; there doesn't seem to have a expert-specifc weight; also I have question on Equation 11: since no acitvation function is involved, how is this different from having a separate W matrix and do xW? Also, if I take Eqn 11 as it is, how could this be equivalent to a vanilla attention? \n* In Table 5, UMoE's improvement over FFN-MoE seems marginal given UMoE's MAC is higher; on base model, the performance gap is clearer, but UMoE is also using relatively much more ; this might imply the method doesn't scale well. Similar observation on downstream task in Table 4. Given that MAC is used for efficiency, if the author switches to FLOPs, the improvement seems even less appealing since 1 MAC = 2 FLOPs. \n* some important aspect of MoE is not addressed like how the author deals with load-balancing?", "was_revised": true, "final_justification": "My concern was addressed with additional experiments from the authors", "rebuttal": "We thank the reviewer for the detailed feedback and constructive suggestions. We are encouraged that our work was found intriguing and our experiments thorough. Below, we address each of the reviewer's points.\n\n# Experimental Details\n\n> do author have MoE layer every 2 transformer layer? in Table 5, is the expert allocation done on all layers?\n\nWe apologize for omitting these details. To clarify:\n\n1. Following previous work [1], in UMoE and all MoE baselines, every Transformer layer is replaced with one corresponding MoE layer, creating a fully MoE-based architecture.\n2. In Table 5, the same expert allocation strategy is applied to all layers.\n\n[1] Niklas Muennighoff, Luca Soldaini, Dirk Groeneveld et al. OLMoE: Open Mixture-of-Experts Language Models, 2025. URL https://openreview.net/forum?id=xXTkbTBmqq.\n\n# Weaknesses\n\n**W1 & W2: Inconsistency between Equation 11 and Figure 3**\n\n> important explanation missing: from Eqn 11 and Figure 3...\n\n> important inconsistency: Equation 11 disagree with figure 3.\n\nWe thank the reviewer for identifying this inconsistency. Figure 3 was simplified and did not accurately represent the low-rank query projections (Eqn 11). We will update the pseudo-code in Figure 3 to be fully consistent with our formulation:\n\n```python\n# Revision for Figure 3\n# S is a sequence of token hidden states, x is the last token hidden state of S\n# For simplicity, this implementation focuses on computing the output for the last token.\n\n# Attention MoE\nK = S @ W_k\nq_shared = x @ W_q\n\nindices, probs = TopKRouter(x, k) # k is head_num\nfor i, p in zip(indices, probs):\n    # Expert-specific query via low-rank adaptation\n    q_i = q_shared + x @ W_a[i] @ W_b[i]\n    # core attention\n    mixing_output = Attention(Q=q_i, K=K, V=S)\n    # Expert Computation\n    expert_output = p * Experts[i](mixing_output)\n\n# ... FFN MoE\n```\n\n> how is this different from having a separate W matrix and do xW?\n\nThanks for this question. The core idea behind Equation 11 is parameter efficiency. In the attention layers of UMoE, the router selects top-$k$ experts for each token. For each selected expert, a unique query vector should be generated using **expert-specific weights**. Having a separate $W_q^i \\in R^{d \\times d_q}$ for $i$-th expert as expert-specific weights is a solution. However, this would significantly increase parameter count. \n\nOur low-rank factorization, $W_q^i \\approx W_q+W_a^iW_b^i$, is a more efficient solution. It allows each expert to have a specialized query projection by adding only the small, expert-specific low-rank matrices ($W_a^i  \\in R^{d \\times r}, W_b^i  \\in R^{r \\times d_q}$) to a shared base matrix ($W_q$).\n\n> Also, if I take Eqn 11 as it is, how could this be equivalent to a vanilla attention?\n\nOur modification affects only the query projection. The core attention mechanism—calculating attention  scores and producing a weighted sum of value vectors—remains unchanged.\n\n> it seems W_k and W_q are fully activated to calculate the attention.\n\nYes, `W_k` and the shared `W_q` is utilized for all tokens. The computation in attention layers remains sparse. This sparsity comes from activating only a subset of experts for the remaining computations, not from the initial key/query projections.\n\n**W3: Computational Efficiency and Scaling**\n\n>  In Table 5, UMoE's improvement over FFN-MoE seems marginal given UMoE's MAC is higher; on base model, the performance gap is clearer, but UMoE is also using relatively much more; this might imply the method doesn't scale well. Similar observation on downstream task in Table 4. Given that MAC is used for efficiency, if the author switches to FLOPs, the improvement seems even less appealing since 1 MAC = 2 FLOPs.\n\nThis is an insightful and crucial point! Table 2 presents a MAC-matched comparison, where UMoE consistently outperforms baselines even when they are allocated similar or greater computational budgets. This suggests UMoE's advantage stems from a more effective formulation of computation rather than simply using more computation.\n\nWhile the performance gap between UMoE and FFN-MoE narrows at a larger scale, UMoE remains superior to prior Attention-MoE methods. We see this as a primary contribution: **UMoE establishes Attention-MoE as a viable and competitive paradigm for scaling Transformers, on par with state-of-the-art FFN-MoE models.** This is a significant step forward, as previous Attention-MoE methods have struggled to match the performance of their FFN counterparts.\n\nTo further investigate performance scaling, we ran a new experiment on our large model. We hypothesized that performance of UMoE could be improved by reallocating activated experts from the FFN to the attention module while keeping the total number of activated parameters constant. Due to time constraints, we trained for 10B tokens on FineWeb-edu.\n\n| **k for Attention Layer** | **k for FFN Layer** | **Perplexity (Validation)** |\n| ------------------------- | ------------------- | --------------------------- |\n| 4                         | 11                  | 19.76                       |\n| 6                         | 9                   | 19.68                       |\n| 8                         | 7                   | **19.55**                   |\n\n**W4: Load Balancing**\n\n> some important aspect of MoE is not addressed like how the author deals with load-balancing?\n\nThank you for this question! UMoE employs the standard auxiliary loss for load balancing introduced by the Switch Transformer. We regret omitting this important implementation detail and will add it to our revision.\n\n# Questions\n\n> Q1: Line 113, the notation is misleading --- notation doesn't agree with W_o; they should really be W_q^I\n\n> Q3: token hidden states\" is more precise.\n\n> Might be useful to make a plot of PPL-MAC to show the Pareto-frontier.\n\nWe appreciate these suggestions for improving clarity and presentation. We have revised the paper accordingly and will add a PPL-MAC plot.\n\n> Q2: does the result imply we should just always skip FFN and only have attention; this is surprising because this seems to disagree with some existing work on importance of FFN.\n\nThis is indeed an interesting finding we aimed to highlight — As described in Lines 141-143 and Figure 1, we can interpret a standard FFN-MoE layer as a special case of our attention formulations where the attention score is constrained to an identity matrix. On the contrary, the experts in attention layer operate on contextually-aware inputs, making them potentially more expressive. Our results in the new experiment above (W3) reinforce this, showing that shifting expert capacity to the attention module is beneficial. This is an intriguing finding that warrant further exploration.\n\n> Q4: What does \"mixing\" in \"pre/post-mixing\" refers to? I thought a_i X is a mixing operation.\n\nYes, $a_i X$ is indeed a mixing operation. In our terminology, \"mixing\" refers to the weighted summation based on attention scores. We will clarify this term in the paper to avoid ambiguity."}, {"review_id": "4OGJkTMEbT", "reviewer": "Reviewer_fL6Z", "summary": "This paper proposes a new method to combine FFN MoE and attention MoE. They reformulate the attention mechanism, propose a pre-mixing formulation and post-mixing formulation, and then apply them with MoE. The experimental results show that the pre-mixing formulation can have a better performance. The UMoE architecture achieves superior performance through attention-based MoE layers while enabling efficient parameter sharing between FFN and attention components.", "questions": "See weaknesses above.", "limitations": "See weaknesses above.", "rating": 4, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Pros:\n\n- The paper is technically sound.\n\n- The paper is easy to understand.\n\n- The experimental results show the effectiveness of the proposed method.\n\nCons:\n\n- The authors claim in line 35 that MoEs are applied on the straightforward two-matrix multiplication pattern of FFNs, which is not the case in most SOTA LLMs. They use SwiGLU instead of MLP.\n\n- The experiments are only conducted on small LLM models with 1.1B parameters as a baseline. Models with more parameters should be used to further verify the usefulness of the proposed method.\n\n- The inference time, memory usage and accuracies of vanilla attention pre-mixing attention and post-mixing attention without MoE should be compared with an ablation study.", "was_revised": true, "final_justification": "I will keep my score.", "rebuttal": "Thank you for your detailed review. We appreciate your recognition that our paper is technically sound and demonstrates the effectiveness of our proposed method.\n\n> **Weakness 1: The authors claim in line 35 that MoEs are applied on the straightforward two-matrix multiplication pattern of FFNs, which is not the case in most SOTA LLMs. They use SwiGLU instead of MLP.**\n\nThank you for this insightful observation. We acknowledge that current SOTA LLMs predominantly use SwiGLU instead of traditional MLPs. Our focus on MLPs in the paper was primarily to highlight the structural similarity between MLPs and the ($W_v, W_o$) matrices in attention layers, which enabled us to abstract a portion of attention as an expert naturally.\n\nAs illustrated in Figure 1, UMoE is designed as a framework that remains agnostic to the expert implementation. To demonstrate this flexibility, we conducted additional experiments with 10B tokens using SwiGLU as experts:\n\n| Model       | # Params | # FFN Intermediate Size | Perplexity |\n| ----------- | -------- | ----------------------- | ---------- |\n| UMoE-SwiGLU | 540 M    | 128                     | 23.565     |\n| UMoE-MLPs   | 540 M    | 192                     | **23.209** |\n\nThese results indicate comparable performance, with MLP experts performing slightly better. This may be because SwiGLU requires specific hyperparameter tuning or is less effective with smaller expert sizes.\n\n> **Weakness 2: The experiments are only conducted on small LLM models with 1.1B parameters as a baseline. Models with more parameters should be used to further verify the usefulness of the proposed method.**\n\nWe agree that scaling our experiments to larger models would further validate UMoE's effectiveness. Given our limited computational resources (requiring over two weeks to train models exceeding 1B parameters, with unstable resource availability), we focused on demonstrating the core concept within these constraints. Nevertheless, we believe that our findings with moderately-sized models are still significant and meaningful for the research community. UMoE presents an interesting and simple architecture that can inspire future work, and we hope this paper will encourage exploration of its scaling potential with more extensive resources.\n\n> **Weakness 3: The inference time, memory usage and accuracies of vanilla attention, pre-mixing attention and post-mixing attention without MoE should be compared with an ablation study.**\n\nThank you for this constructive suggestion. An ablation study comparing vanilla attention with our pre-mixing and post-mixing reformulations is essential to isolate the performance gains of UMoE from the attention mechanism itself. We have run this experiment on a single H100 GPU, and the results are presented below:\n\n| GPT                  | **Perplexity** | **Pre-filling latency (s)** | **Memory footprint** |\n| -------------------- | -------------- | --------------------------- | -------------------- |\n| w/ Vanilla attention | 25.381         | 0.2823                      | 4.68 G               |\n| w/ Pre-mixing        | 25.380         | 0.2998                      | 3.93 G               |\n| w/ Post-mixing       | 25.374         | 0.3190                      | 6.18 G               |\n\nAs shown, the perplexity is nearly identical across all three configurations. This is expected, as the reformulations only change the execution order of matrix multiplications within attention and should not alter the final output. \n\nRegarding efficiency, the pre-mixing attention variant shows a slightly lower latency  and memory footprint than post-mixing, making it a more efficient alternative. Furthermore, our experiments indicate that pre-mixing yields superior performance when combined with MoE architectures. Based on these findings, we selected pre-mixing as our primary configuration for UMoE."}, {"review_id": "7YDFEkm45r", "reviewer": "Reviewer_twFT", "summary": "This paper proposes a new MoE architecture, called UMoE. It unifies the MoE for both attention and FFN modules, and supports unified training of them using shared experts. Experiments results on efficiency and efficacy demonstrates the superiority of this method.", "questions": "1. Could the author present the training/inference speed (steps/s or tokens/s) for the method and the baseline?\n2. Is anything missing in the head of Table 3? There are no contents in the right columns of UMoE.", "limitations": "yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "### Strengths\n1. The UMoE framework unifies the MoE architecture for both modules of the transformer\n2. The increased expressiveness of Attention-MoE brings performance improvements for the MoE architecture\n\n\n### Weaknesses\n1. I observe that the performance improvements narrows when scaling the model sizes in Table 4. So I am worried about the scaling of this method .\n2. Many benchmarks here could be evaluated using few-shot. So a comparison with the baselines using few-shot prompting is welcomed.", "was_revised": true, "final_justification": "All my concerns have been addressed and I think this work is novel and solid.", "rebuttal": "We sincerely appreciate your thoughtful feedback and positive evaluation of our work.\n\n> **Weakness 1: I observe that the performance improvements narrows when scaling the model sizes in Table 4. So I am worried about the scaling of this method .**\n\nThank you for this insightful observation! We acknowledge that the performance gap between UMoE and FFN-MoE narrows at a larger scale. However, UMoE remains superior to prior Attention-MoE methods. This represents a key contribution: **UMoE establishes Attention-MoE as a viable and competitive paradigm for scaling Transformers, on par with state-of-the-art FFN-MoE models**. This is a significant step forward, as previous Attention-MoE methods have struggled to match the performance of their FFN counterparts.\n\nTo further investigate performance scaling, we ran a new experiment on our large model. We hypothesized that performance of UMoE could be improved by reallocating activated experts from the FFN to the attention module while keeping the total number of activated parameters constant. Due to time constraints, we trained for 10B tokens on FineWeb-Edu:\n\n| **k for Attention Layer** | **k for FFN Layer** | **Perplexity (Validation)** |\n| ------------------------- | ------------------- | --------------------------- |\n| 4                         | 11                  | 19.76                       |\n| 6                         | 9                   | 19.68                       |\n| 8                         | 7                   | **19.55**                   |\n\nThis suggests that with architecture-specific tuning, the performance of UMoE at scale can be further enhanced.\n\n> **Weakness 2: Many benchmarks here could be evaluated using few-shot. So a comparison with the baselines using few-shot prompting is welcomed.**\n\nThanks for this suggestion! As requested, we evaluated our models and the FFN-MoE baseline in a 3-shot setting using the `lm-evaluation-harness`. We found that other Attention-MoE baselines (MoA, SwitchHead) consistently underperformed FFN-MoE, so for clarity, we present a direct comparison with the strongest baseline below.\n\n| **Model**        | **Params** | **HellaSwag** | **PIQA**  | **ARC-E** | **ARC-C** | **RACE**  | **Lambada** | **MMLU** | **Wino** | **Avg.**  |\n| ---------------- | ---------- | ------------- | --------- | --------- | --------- | --------- | ----------- | -------- | -------- | --------- |\n| **Base Models**  |            |               |           |           |           |           |             |          |          |           |\n| FFN-MoE          | 535 M      | 38.48         | 67.03     | **64.78** | 31.06     | 29.48     | 22.02       | 32.1     | 52.9     | 42.23     |\n| UMoE (Att)       | 547 M      | 40.02         | **67.56** | 62.0      | **33.77** | **29.89** | 21.94       | 33.4     | 53.4     | **42.75** |\n| UMoE             | 540 M      | **41.38**     | 66.45     | 61.36     | 33.66     | 28.43     | **22.15**   | **34.0** | **53.7** | 42.64     |\n| **Large Models** |            |               |           |           |           |           |             |          |          |           |\n| FFN-MoE          | 3.8 B      | 53.38         | 74.40     | 70.08     | 41.99     | **37.99** | 35.11       | 30.4     | **57.9** | 50.16     |\n| UMoE (Att)       | 3.8 B      | 53.52         | 74.08     | 69.50     | **43.01** | 36.81     | 36.59       | 30.9     | 56.5     | 50.11     |\n| UMoE             | 3.6 B      | **54.72**     | **74.94** | **71.63** | 40.93     | 35.79     | **37.43**   | **32.8** | 57.4     | **50.70** |\n\nThese few-shot results are consistent with our zero-shot findings: UMoE (Att) is highly competitive with FFN-MoE, and UMoE model demonstrates the best overall performance.\n\n> **Q1: Could the author present the training/inference speed (steps/s or tokens/s) for the method and the baseline?**\n\nThis is a very constructive suggestion! We implement the forward pass of experts in MoE layers using GPU kernels written in Triton. Below are the training throughput (tokens/s) and inference latency measurements (for pre-filling 1024-token sequences) conducted on a single H100. We didn't measure token generation speed for inference since we found decoding can hardly fully utilize GPU.\n\n| **Base Model**  | **Model Size** | **Pre-filling latency (s)** | **Training throughput (token/s)** |\n| --------------- | -------------- | --------------------------- | --------------------------------- |\n| GPT             | 134 M          | **0.1376**                  | **126,508**                       |\n| FFN-MoE         | 535 M          | 0.1974                      | 74,415                            |\n| MoA             | 525 M          | 0.1843                      | 90,394                            |\n| SwitchHead      | 533 M          | 0.2253                      | 69,374                            |\n| UMoE-Att        | 547 M          | 0.2006                      | 76,538                            |\n| UMoE            | 540 M          | 0.2197                      | 71,461                            |\n| **Large Model** |                |                             |                                   |\n| GPT             | 1.1 B          | **0.2950**                  | 11,634                            |\n| FFN-MoE         | 3.8 B          | 0.4080                      | 9,799                             |\n| MoA             | 3.6 B          | 0.3624                      | **12,321**                        |\n| SwitchHead      | 3.7 B          | 0.4328                      | 8,882                             |\n| UMoE-Att        | 3.8 B          | 0.4120                      | 11,377                            |\n| UMoE            | 3.6 B          | 0.4392                      | 9,194                             |\n\nAll MoE variants exhibit similar latency but require significantly more computation time than the dense baseline despite having comparable theoretical MACs. Both the routing mechanism and sparse expert computation within MoE layers introduce substantial overhead. For instance, SwitchHead and UMoE are slower than other MoE models because SwitchHead implements two MoE layers per attention head, while UMoE has two MoE layers per block (one in attention, one in FFN).\n\nOur paper's focus on MACs as the primary metric for computational efficiency, as it reflects the theoretical advantages of the architecture, abstracting away from current hardware and software implementation limitations. We believe these practical speed gaps will narrow with future advancements in GPU kernels and parallelization strategies. We will add this discussion to our paper.\n\n> **Q2: Is anything missing in the head of Table 3? There are no contents in the right columns of UMoE?**\n\nThank you for noting this oversight! The right columns of Table 3 represent variants of our model with different parameter-sharing strategies. We apologize for the lack of clarity. We have revised the table header and caption to clearly explain each column."}]

Author rebuttal:

Meta-review: This paper argues that part of attention block of transformer architecture can be viewed as FFN expert. With this insight, authors propose a full mixture-of-expert transformer layer - UMoE. Experiments indicate that given similar compute, UMoE achieves superior performance and enables efficient parameter sharing between FFN and attention. 

All the reviewers commend novel insights and proposed method that enables expert sharing across attention and FFN. The paper is well written and easy to follow and the contributions are clear.  Empirical results are compelling, however authors are encouraged to include additional details wrt experiment setup and analysis of the experimental results, including additional ablations. 

Authors did a great job during the rebuttal phase, providing additional evaluations requested by the reviewers. We expect these new results and explanations to be incorporated into the final version of the paper. 

Additionally reviewers encourage further proofreading of the paper
Decision: Accept (spotlight)