Mixture of Hidden-Dimensions: Not All Hidden-States’ Dimensions are Needed in Transformer
Authors: Yilong Chen
Junyuan Shang
Zhenyu Zhang
Jiawei Sheng
Tingwen Liu
Shuohuan Wang
Yu Sun
Hua Wu
Haifeng Wang
Venue: ICML 2025 poster (2025)
Keywords: Large Language Models
Transformer
Efficient Model
Conditional Activation
Hidden-dimension Sparsity
Primary area: deep_learning->large_language_models

Abstract: Transformer models encounter inefficiency when scaling hidden dimensions due to the uniform expansion of parameters. When delving into the sparsity of hidden dimensions, we observe that only a small subset of dimensions are highly activated, where some dimensions are commonly activated across tokens, and some others uniquely activated for individual tokens. To leverage this, we propose MoHD (Mixture of Hidden Dimensions), a sparse architecture that combines shared sub-dimensions for common features and dynamically routes specialized sub-dimensions per token. To address the potential information loss from sparsity, we introduce activation scaling and group fusion mechanisms. MoHD efficiently expands hidden dimensions with minimal computational increases, outperforming vanilla Transformers in both parameter efficiency and task performance across 10 NLP tasks. MoHD achieves 1.7% higher performance with 50% fewer activatied parameters and 3.7% higher performance with 3× total parameters expansion at constant activated parameters cost. MoHD offers a new perspective for scaling the model, showcasing the potential of hidden dimension sparsity.

Mixture of Hidden-Dimensions:
Not All Hidden-States' Dimensions are Needed in Transformer
Yilong Chen 12
Junyuan Shang?
Zhenyu Zhangè
Jiawei Sheng" Tingwen Liu Haifeng Wang
Shuohuan Wang
Yu Sun
Hua Wu²
Abstract
Transformer models encounter inefficiency when scaling hidden dimensions due to the uniform ex- pansion of parameters. When delving into the sparsity hidden dimensions, we observe that only small subset of dimensions are highly ac- tivated, where some dimensions are commonly activated across tokens, and some others uniquely activated for individual tokens. To leverage this, we propose MoHD (Mixture of Hidden Dimen- sions), a sparse architecture that combines shared sub-dimensions for common features and dynam- ically routes specialized sub-dimensions per to- ken. To address the potential information loss from sparsity, introduce activation scaling and group fusion mechanisms. MoHD effi- ciently expands hidden dimensions with minimal computational increases, outperforming vanilla Transformers in both parameter efficiency and task performance across NLP tasks. MoHD achieves 1.7% higher performance with 50% fewer activated parameters and 3.7% higher per- formance with 3x total parameters expansion constant activated parameters cost. MoHD offers new perspective for scaling the model, showcas- ing the potential of hidden dimension sparsity.
Activaioon Magnkkuds
10.0
Token 0
MoHD Parameters Efficiency
7.5
Token 2
5.0
47
2.5
Activate
46
X
0.0
500
1000
15000
Avg.
45
Shared Sub-Diis Specialized Sub--lima
44
Transformer MoHD
MoHD Curve
W
Router
1000
2000
MoHD
Embedding
# Activated Params (M)
Figure I. We observe that Transformer hidden states have both token-shared and token-specific activated dimensions. Based on this, we propose MoHD, which combines shared and specialized sab-dimensions for mixed activation. Compared to Transformers, MoHDD offers significantly higher parameter efficiency.
can enhance their capabilities with sufficient data. However, the large number of parameters in LLMs often results in high training and inference costs. To address this, flexible model architectures (Jiang et al., 2024b; Cai et al., 2024b;c) are sought to enable parameter scaling while maintaining com- putational efficiency. In Transformers, the parameters are determined by hidden and intermediate dimensions. Some studies (Qiu al., 2024; Liu et al., 2024) identify the spar- sity of to design adaptive networks, such MoE (Cai al., 2024c; Dai et al., 2024; Xue et al., 2024a), apply pruning (Xia et al., 2023; Chen et al., 2023; Ma et al., 2023b) and local activation mechanisms (Liu et al., 2023a) to boost efficiency.
1. Introduction
activations, and use them
Large Language Models (LLMs)) (Anthropic, 2023; Ope- nAI, 2023; Touvron et al., 2023a) have achieved impressive performance in various natural language processing tasks. Recent study (Kaplan et al., 2020) suggests that scaling mod- els by increasing parameters and computational resources
A Although elastic scaling of the intermediate dimension has been studied, scaling hidden dimensions with controllable computational costs remains underexplored (Figure 9). In general, the hidden dimension reflects the embedding size of all tokens, and expanding it can increase the model capacity to capture intricate patterns. However, existing Transform- ers (Vaswani et al., 2023) treat all token dimensions equally, resulting in significant computational and memory overhead as the hidden dimensions scale.
'Inttittee of Information Engineering, Chinese Academy of Sci- ences "School of Cyber Security, University of Chinese Academy of Sciences 'Baidu Inc. Correspondence to: Tingwen Liu <liut- ingwen@ @ aic.ac.en>, Project lead: Junyuan Shang <shangjun- yuan @ baiidu.com..
Proceedings of the = International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Mixture of Hlidden-Dimensee Not All Hidden-States" Dimensions are Needed in Transformer
Considering to limited understanding of hidden dimensions in L.Mss, we study the activation magnitudes and find significant sparsity, where 50% of the dimensions con- tribute 92.54% of the activation magnitudes (Figure 2 left). Among highly activated dimensions, we identify shared dimensions consistently activated across tokens, modeling common features, and specialized dimensions uniquely ac- tivated for individual tokens, capturing higher-level seman- ties (Figure 3). This inspires efficient design selectively activating shared and specialized sub-dimensions (Figure I). Furthermore, we observe a consistent Activation Flow Pat- tern across different blocks (Figure 2 middle, Figure 5), where Attention outputs vary while FFN outputs remain stable, guiding us to design unique and separate sparsity architectures to maintain activation flow integrity.
analyses on MoHD's routing mechanism and sparsifica- tion phenomenon. MoHD is the first method to introduce sparse mixture activation for expanding hidden dimensions, offering new approach to designing efficient architectures.
2. Observation
This section presents key findings for the design of MoHD. In Section 2.1, we observe long-tail sparsity in hidden di- mension activations. Section 2.2 examines activation flow in Transformers. In Section 2.3, we identify shared continuous and unique discrete high activation behaviors across tokens, guiding the design of hidden dimension sparsification.
2.1. Sparsity in Tokens' Hidden Dimension
In this paper, we propose MoHD (Mixture of Hidden- Dimensions), novel approach that greatly expands hid- den dimension capacity through sparse conditional acti- vation, while keeping the activated parameters nearly un- changed. Specifically, MoHD introduces two types of sub- dimensions in each layer of the model's Attention and FFN components: shared sub-dimensions that are always acti- vated capture common dimensional information across different tokens, and specialized aub-dimensions that are selectively activated to capture roken-specifi specialized dimensions. To ensure load balancing, we apply a balane- ing loss to the specialized sub-dimensions. An activation scaling mechanism and a grouped fusion mechanism are introduced to mitigate information loss from dimensional downsampling and maintain the activation flow. In this way, MoHD can be used to scale the model's hidden dimension without increasing the number of parameters, or to signif- icantly reduce the active hidden dimension to lower the computational cost.
We analyze activation magnitudes of 4096 hidden dimen- sions in LLaMA2-7B (Figure 2). The left panel reveals the long-tail sparsity, consistent with observations in Liu et al. (2024): in the 16th layer Attention input, the top 1000 dimensions account for 71.96% of total magnitude. The middle panel highlights functional divergence between Transformer's components: Attention shows higher, fluc- tuating activations, whereas FFN maintains lower, stable activations. This contrast underscores the need for differen- tiated activation designs. For details, see the Appendix F.
2.2. Activation Flow in Transformer
We also investigate activation magnitude variations within a single Transformer block, as shown in Figure 10. Consistent Activation Flow patterns were observed across blocks. The activation magnitude variations in each layer follow a simi- lar flow pattern: The Attention module compresses input activations to 58% through projections, demonstrating its ability to suppress irrelevant information. FFN selects and maintains the stability of the activation flow magnitude based on the Attention's output. Residual connections regulate activation changes by restoring compressed magni- tudes. This observation inspires us to maintain the activation flow reduce information loss. More in Appendix F.2.
To demonstrate the effectiveness of MoHD, we pretrain Vanilla Transformers with 355M, 495M, and 1.13B parame- ters based LLaMA's architecture (Touvron et al., 2023b), and their MoHD versions in both hidden dimension com- pression and expansion settings with scaling factors: 50%, 75%, 2x, 3x, and 4x. We evaluated these models on 10 NLP tasks, showing the advantages of the MoHD architec- ture. Results indicate that MoHD consistently outperforms Vanilla and Mixture of Experts Transformers with the same activated parameters across all model sizes. In the com- pression setting, MoHD reduces activated parameters by 50%, retaining 99% of original performance. In the expan- sion setting, MoHD keeps activated parameters constant while expanding hidden dimensions 4x, achieving up to an 8.37% relative performance improvement. Notably, MoHD- 355M outperformed LLaMA2-355M and even matched LLaMA2-1.131B performance, while reducing activated parameters to 28.9% of LLaMA's. To explore the impact of increasing hidden dimensions, we conducted detailed
2.3. Continuous High Activation
We investigate the temporal correlation of activation sparsity ratio by analyzing high activation values across consecu- tive tokens and identifying indices repeatedly activated by multiple tokens. A clear correlation in activations is oh- served over consecutive tokens. Figure 2 right shows the number of commonly activated dimensions across to 9 consecutive tokens, with the x-axis representing the thresh- old for high activation. Using the top 20% activation values as the threshold, 2672 dimensions are commonly activated across 2 tokens, and 673 dimensions remain active across 9 tokens. Figure 3 further illustrates correlated activation
2

---

Mixture of Hidden-Dimensee Not All Hidden-States" Dimensions are Needed in Transformer
Activation Migniutte Curve
6 1000
Sparsity
Sparsity
Hoh Actioboon Count
4000
de
Consecttie Tokoons Count
12.5
3000
Dimension on Dimens on 1000-1096
10.0
7.5
0.6
2000
0.5
1000
20. 59% 6.49%
0.97%
10000 Hidde Dimersion Index (Sorted)
3000
10
15
20
25
30
20
400
60
Layer
High Activatton Threshold (%)
Figure 2. Left: Activation magnitudes sorted in descending order with the percentage representing the cumulative activation sum. Middle: Sparsity of hidden dimension activations in Attention and FFN outputs across layers. Right: Number of shared activation dimensions varying activation magnitude thresholds, with curves showing the count for consecutive tokens ranges from 2 (blue) to 20 (red).
tivation sparsity ratio 5, MoHD dynamically activates subset S S [d] of parameters in the weight matrix W E Rd through hidden dimension sparsity. Inspired by observation in Section 2.1, We partition W into N sub- dimensions of size de such that Ndg d. structured as W [W1, W2, ... WN] where W; W[i 1)de ide). The routing gate Gate 6, N) selects top-K sub- dimensions via:
teken bola
2oRtg 1i0ge www AA Hidden tirmimatio Baa
(1 (1) , 0. otherwise,
Si, (opk((s;)) K),
8; Softmax;
Figure 3. Activation patterms of 4,096 hidden dimensions, clustered and reordered across five tokens the 16th layer of LLaMA2-7B.
(1) where 8; denotes poken-to-ssubsace affinity, of is the cen- troid of i -th sub-dimension at layer I, and K = [6N]. The routing weights enable dynamic amplification/suppres sion sub-dimension representations during optimization. Activated sub-dimension outputs {y, are concatenated produce the final d-dimensional output.
patterns over 5 tokens, where 4096 hidden dimensions are clustered based on their activation patterns. About 400 di- mensions are commonly activated across all tokens, while roughly 200 dimensions are uniquely activated in each to- ken. This indicates that each token's activations contain both shared sub-dimensions activated across tokens and token-specific sub-dimensions. Shared activations cap- ture similarities, while unique activations reflect differences. These observations inspire the shared-we specialzzd activation mechanism in MoHD's design.
- 9: (z,, = Wax
The gated concatenation operator i= Ii (im;(z,) dynami- cally selects sub-dimensions through sparse activation (spar- sity 8). Only = [6N] sub-dimensions receive non-zero weights gs, inducing sparsity in both the token embedding and weight matrix W. This produces output y/, Rdd while reducing activated parameters in Ws 8 || loo
3. Mixture of Hidden Dimensions (MoHD)
In this section, we propose the Mixture of Hidden Dimen- sions (MoHD) architecture to scale the hidden dimension without increasing activated parameters. The workflow is illustrated in Figure 4. Based on observations from Sec- tion 22.3, we introduce the Shared and Specialized Sub- Dimension Mixed Activation mechanism in Section 3.1. We also apply sparsified components to Attention and FFN blocks (Appendix A) and address information degradation with activation scaling and grouped fusion mechanisms in Section 3.2. In Section 3.4, we discuss a balance loss and present the implementation Section 3.5.
3.2. Activation Flow Maintenance
In Section 3.1, our sparse activation mechanism gener- ates dimension-sparsifee outputs ys to reduce computa- tion, but risks information degradation due to softmax- induced weight skewing and suppression of low-weight di- mensions during parallel concatenation. To preserve consis- tent Activation Flow (Section 2.2), we first introduce Sub- dimension Scaling with = Σ g ensuring stable acti- vation magnitudes across dimensions. A Grouped Fusion Layer then projects y, back to the original d-dimensional space using Monarch matrices (Dao et al., 2022; Chen et al., 2024a), structured as M Σij mij with a receptive field
3.1. Mixture of Sub-Dimensioon Activation
Let X E R denote embeddings of n tokens and E Rlxd denote single token embedding. Given ac-
3

---

Mixture of Whidden-Dimensions Not All Hidden-States" Dimensions are Needed in Transformer
Embedding
Embbddinn
Meell Feod Forword
Motid Muats-laad Adiiniion
Pualon
MoHO Pred Fonnard
Ho.HD Blocke
Ritshorm
Stop
Metoo
SLU
Mati-Head Artention
Emieeiiing
Roster
Viput
Emstediiing
Emneeddlnng
Enbedding
Figure An illustration of MoHD: A single MoHD Block mirrors the structure a LLaMA Block, consisting of two main components: MoHD Attention and MoHD FFN, each with pre-norm and residual connections. In both components, matrices are selectively activated hased on dimensions chosen by the Router. The Router weights shared dimensions and selects few sparsely activated ones for each token. The resulting outputs from these matrices weighted, concatenated, and mapped back the original dimensions.
MoHD Component
while another portion is selectively activated, likely encod- ing coken-specific features. Motivated by this, propose structured routing mechanism that decouples shared and specialized feature encoding. Specifically, we partition each layer's hidden space into two complementary subspaces (Fig. 4): Shared Sub-Dimensions (ip portion) maintain constant activation to capture cross-token common features. Specialized Sub-Dimensions (8 yp portion) employs dy- namic routing via Topk((ō N) selection, encoding fine- -re-ggrined contextual patterns.
Router Weights
Router
Scaling
Group Fusion
Residual Connection
Activation Flow
Impaired Activation Flow
Figure 5. Maintaining Activation Flow: Sparse activations can result information loss. The scaling factor adjusts the activations to restore the original magnitude, and the grouping and fusion mechanism restores activations their original dimensionality.
Softmax,
Ss:
st E 1
9x
81:
8t Topk (18) pN < < N), (6 N))),
0, otherwise.
(5)
reducing the mapping complexity O(d/r). Resid- ual Connections further preserve critical signal flow by bypassing suppressed dimensions. As shown in Figure 5, these components maintain ć-level parameter sparsity while mitigating information loss and computational overhead:
The routing mechanism ensures that Shared Sub- Dimensions are consistently activated for all tokens, consolidating common information, while Specialized Sub-Dimensions are encouraged to diversify. All sub- dimensions are weighted to allow selective amplification/- suppression of representations during optimization.
mi.d/r
My.
Ys:
md/rr1
M//r
3.4.
Load Balance
where the Monarch matrix M enables efficient grouping and transformation, reconstructing information across the orig- inal hidden dimensions while maintaining computational efficiency. In summary, the forwarding process for a single MoHD module can be formally represented as follows:
Research on conditional computation (Zhou et al., 2022; Jiang et al., 2024b; Dai et al., 2024) has shown that auto- matically learned routing strategies can often lead to load imbalance issues, where the model tends select only a few aab-dimensions, leaving others underutilized and insuf- ficiently trained. To distribute tokens more evenly among different sub-dimensions and smooth out the rouler score distribution, we incorporate Sub-Dimension Load Bal- ance Loss (Dai et al., 2024): Define 3 is a scaling factor and IK (argmax(g.) )=x) is an indicator function that returns the i-th sub-dimension has the highest gating score for the
y Ma | (x.,, g, X₂ Gate(x, N).
(4)
3.3. Mixed Activated lub-Dimensions
As discussed in Section 2.3, a portion of the hidden dimen- sions consistently activated, capturing shared features,

---

Mixture of Hidden- Hiidden-Dimennions: Not All Hidden-States" Dimensions are Needed in Transformer
s-th sequence position and 0 otherwise.
Easy(ARC-E) (Clark al., 2018b), and 10-shot HellaSwag (Hella.) (Zellers al., 2019), alongside 25-shot accuracy for ARC Challenge (ARC-C) (Clark et al., 2018a). In the assessments of continued QA and text understanding, we report 0-shot accuracy for LogiQA (Liu et al., 2020), 32-shot BoolQ (Clark et al., 2019), and O-shot LAMBADA (Lam.) (Paperno et al., 2016). All reported results are calculated with the mean and stderr of multiple experiments.
Eves Hf (argmax(u.)
LB = 8
(6)
N
-=19: 93
The term represents the normalized gating score
2==1 9)
for sub-dimension ensuring that the contributions of each sub-dimension are proportional to their selection frequency.
Baseline. Following the architecture of LLaMA2, con- structed models at three parameter scales: 355M, 495M, and I.13B, with hidden dimensions of 1024, 1536, and 2048, as shown in Table 7. For each parameter scale, we develop three variants: Vanilla Transformers (LLaMA architecture) and MoHD-baaed models. The flexible MoHD architec- ture allows for compressing activated parameters without changing the total parameter count expanding the model parameters while retaining the activation size. We exper- iment with five hidden dimension scaling factors-0.5xx 0. 0.75x, 2x, 3x, and 4x-to showcase MoHD's ability to reduce activation and enhance model capacity. All models are the same initialized and pre-trained on 50 billion tokens.
3.5. Implementation
MoHD sparsifies hidden dimensions in Transformer blocks by synchronizing sparsification across up- pp-projection and ddwn-prrojectin matrices, as well as the input x. We define Hidden dimension sparsity and formulate sparsely FFN and Attention mechanisms in Appendix A. Based on these mech- anisms, we construct MOHD BLOCK using MoHD MHA and MoHD FFN components (Details in Appendix. G). Residual connections are applied mitigate information loss, and LayerNorm layers are placed before both MHA and FFN inputs. The forward of the block is defined as:
BLOCK (x) FFNDOHD (MHAMoHD(()) (x) + + x.
4.2. Result
The final training objective combines the cross-entropy loss LCE for language modeling with the load balance loss Lg:
Capability in Compression. Table (blue rows) demon- strates the capabilities of MoHD on the 355M, 495M, and IB versions of LLaMA2 with 50% and 75% hidden dimen- sions activated. Results show that MoHD maintains even improves performance with partial activation. At 355M, MoHD with 50% % activation incurs only a 0.4% perfor- mance loss compared to the baseline, highlighting hidden dimension sparsity. Furthermore. MoHD with 75% acti- vation outperforms the fully activated baseline, achieving gains of 0.5%, 1%, and 1.8% for the 355M, 495M, and 1B models, respectively. Performance gains increase with model size: MoHD 50% achieves relative improvements of -0.4%, +0.3%, and +1.7% for the 355M, 495M, and IB models, showing potential for larger-scale applications. Although compressing activations to 50% slightly reduces Commonsense metric scores, MoHD retains strong lan- guage modeling capabilities under low activation settings.
L= L=LE + LB-
(8)
4. Experiments
4.1. Experimental Setup
Data. To pretrain MoHD models and baseline models, we employ the RedPajama (TogetherAI, 2023), which parallels the LLaMA training data across seven domains: Common- Crawl, C4, GitHub, Wikipedia, Books, ArXiv, and Stack- Exchange. This dataset comprises a validation set with 2 million tokens, a training set containing 50 billion tokens.
Training, Our experimental framework utilizes the Sheared- LLaMA codebase (Xia al., 2023) implemented on the 8
Composer package (Team, 2021), and is executed NVIDIA A100 GPUs (80GB). The models are trained with a sequence length of 4096, employing global batch size of 256. MoHD models are trained for 50000 steps (50B token budget). The learning rates were set at 3e-4 for all parameters. The baselines and all MoHD models follow the same training setup, starting from random initialization and training on the same amount of data.
Capability in Expansion. Table 1 (pink rows) shows MoHD's scalability with 3x, and 4x hidden dimension expansion. MoHD achieves performance comparable to models with equivalent parameters while using fewer activated parameters. For instance, MoHD 2x (355M) ex- ceeds the baseline by 2.2%, outperforming LLaMA2-495M and LLaMA2-1.13B. As the scale increases, the benefits of MoHD become more pronounced, yielding improve- ments of 2.2%, 0.7%, and 3% for the 355M, 495M, and 1. 1.13M parameter models, respectively. Tripling the hid- den dimension yields optimal results: MoHD x3-1.48B achieves a 2.2% improvement over the baseline and 1.5%
Evaluation. We employed the in-evaluation-harnes. (Gao et al., 2021) to evaluate our models. For common sense and reading comprehension tasks, we report 0-shot accuracy results for SciQ (Welbl et al., 2017), PIQA (Bisk et al., 2020), WinoGrande (WG) (Sakaguchi et al., 2020), ARC
5

---

Mixture of Hlideen-Dimensiin Not All Hidden-States" Dimensions are Needed in Transformer
Table Comprehensively evaluate the basic capabilliiess of models with different activated parameters. particular, MoHD 50%-355M represents model with 355M total parameters using MoHD to compress 50% hidden dimensions. Green and Red values indicate metrics that exceed fall below baseline, respectively. # Activate refers to all activated parameters excluding Embedding layers.
Commonsense Rending Comprehension
Continued
Knowledge
Model-Params
# Activate
Avg.
ScQ
PIQA
WG
ARC-E
ARC-C
Hella.
LegiQA
BoolQ
Lam.
MMLU
LLaMA2-3550
289M
74.0
65.2
50.5
44.7
20.1
31.1
19.5
59.7
36.6
25.2
42.7
MoHD 50%-355M
145M
74.0
65.6
50.4
43.9
19.7
30.7
20.6
54.7
37.7
25.6
42.3
MoHD 75%-355M
217M
75.3
65.6
50,9
44.7
20.9
31.1
22.3
55.8
38.9
26.2
43,2
MoHD x2 -710M
289M
76.6
67.5
49.8
47.7
23.0
33.4
20.7
60.5
43.3
26.5
44.9
MoHD x3-1.06B
289M
77.1
67.8
51.1
47.6
21.8
33.9
20.6
55.8
43.6
25.6
44.5
MoHD x4-1.42B
289M
77.6
67.9
49.1
47.0
23.3
33.9
22.1
57.5
44.3
24.7
44.7
LLaMA2-495M
396M
75.4
66.5
51.3
45.5
19.9
32.0
21.7
60.5
38.9
25.8
43.8
MoHD 50%-495M
198M
76.9
67.1
52.7
46.4
20.1
32.3
21.5
57.0
40.7
26.2
44.1
MoHD M5%-495M
297M
76.4
67.3
50.6
45.8
21.1
33.0
23.7
61.8
41.7
26.2
44.8
MoHID x2 -989M
396M
77.1
51.1
47.6
21.8
33.9
20.6
55.8
43.6
25.6
44.5
MoHD -3-BBBBB
396M
77.0
69.0
51.1
48.8
23.6
35.6
22.0
58.6
48.2
26.1
46.0
MoHD x4 -1.98B
396M
79.1
67.4
49.8
49.1
22.0
35.2
20.7
60.9
47.4
26.1
45.8
LLaMA2-1.133
IB
81.0
68.1
51.8
49.3
23.2
35.0
21.7
47.0
38.9
26.4
44.2
MoHD 50%-1.13B
503M
78.9
67.8
48.7
21.2
35.2
21.5
61.1
48.8
25.5
45.9
MoHD 75%-1.13B
755M
80.3
69.3
52.3
50.8
24.5
36.1
22.3
51.2
48.4
25.0
46.0
MoHD x2 22.77B
IB
81.2
70.9
54.1
53.0
24.6
38.3
22,4
50.5
52.1
25,5
47.2
MoHD x3 B344BB
IB
83,6
69.8
53.1
51.9
25.4
38.3
21.0
56.5
53.0
26.6
47.9
MoHD x4 -4.55B
IB
82.4
70.0
52.8
51.6
23.4
38.0
23.4
54.9
50.7
26.6
47.4
LLaMA2-2.7B
2.54B
82.5
70.8
56.3
54.4
27.8
39.3
23.5
44.4
37.7
25.3
46.2
47-
18-
MOHD x3-14488
Perplexity
16-
MOHD 75% 495M
MOHD 50N-495M
65
14-
Transformer
Transformer
MoHD
12-
43
MoHD Curve
MoHD Curve
500
10C0
1500
2000
2500
1000
2000
3000
4000
5000
10
20
30
40
Activaitdd Params (M)
ALL Params
Pretraining Budget (B)
Figure Left: Average score with activated parameters, with point size representing total parameters. Middle: Average score with total parameters, with point size representing activated parameters. Right: Perplexity curves for models pre-trained on S0B tokens.
over MoHD x2. MoHD effectively expands hidden dimen- sions and leverages sparsity to boost performance.
495M, MoHD x3-1.48B, MoHD 75%-495M, and MoHD 50%-495M. MoHD shows training stability and btter efficiency, with smoother and lower perplexity curves com- pared to LLaMA2-49MM. MoHD x3-1.48B shows the great- est improvements, while MoHD 75%-495M and MoHD 50%-495M maintain strong training performance, demon- strating robust learning even with partial parameter activa- tion. Overall, MoHD effectively expands hidden dimen- sions, improving learning capabilities during training.
Parameters Efficiency. Figure 6 (Left) shows that MOHD achieves a high parameter efficiency. On the 400M and IB scales, MoHD delivers 2.2% and 3% improvements, respectively, over the baseline with less than 50% activa- tion. Performance gains grow with activated parameters, underscoring MoHD's 's scalability at larger scales. Figure 6 Middle highlights MoHD's advantage with increasing to- tal parameters: at smaller scales, MoHD matches baseline performance with fewer activations, while at larger scales, MoHD achieves greater gains under the same total parame- ter count by leveraging hidden dimension redundancy.
Compare to MoE. As Sparsely-activated Transformer in a different dimension, we compare MoHD and MoE. DeepSeek MoE (Dai et al., 2024) models are trained from scratch at 355M as the baseline Even with 75% sparsity, MoHD outperform the baseline while MoE casue perfor-
Training Stability. Figure 6 (Right) visualizes evaluation perplexity during pretraining 50B tokens for LLaMA2-
"using Multi-head Attention for faimness. Settings in Table
6

---

Mixture of Hlideen-Dimension Not All Hidden-States" Dimensions are Needed in Transformer
Table Performance comparison of DeepScek-V2 V2 MoF (Dai et et.., 2024) and MoHD at 355M scale (50B training budgets).
Table 3. Eval Perplexity with ablation on MoHID (10B training budgets). "w.o." indicates the method was ablated.
Method
SeiQ
PIQA
Hella.
ARC-E
LogiQA
Lam.
Avg. f
Method
Perplexity 4
L.-35M
74.0
65.2
31.1
44.7
19.5
36.6
45.2
MoHD x2 -710M
10.25
MoE 75%
73.7
64.4
30.6
43.1
19.1
35.0
44.3 (-0.9)
w/o Mixed Activated Sub-Dimensiios
11.08 (+0.83)
MoHD 75%
75.3
65.5
31.1
44.7
22.3
38,9
46.3 (+1.1)
w/o Balance Loss
10,41 (+0.16)
MoE x2
75.1
65.2
31.9
45.2
20.7
38.0
46.0 (+0.8)
w/o Group Fusion Layer
10.47 (+0.22)
MoHD x2
76.6
67.5
33.4
47.7
20.7
43.3
48.2 (+3.0)
w/o Sub-Dimension Scaling
11.41 (+1.16)
MoF. x3
76.2
65.8
32.2
46.2
20.1
38.8
46.5 (+1.3)
LLaMA2-355M
11.61(+1.36)
MoHD x3
77.1
67.8
33,9
47.6
20.6
43.6
48.4 (+3.2)
mance loss. MoHD outperforms MoE by an average of 2% in both compression and expansion settings. By expanding the hidden dimension, MoHD enhances the representa- tion ability of each token, demonstrating higher parame- ter efficiency than MoE. More details in Appendix D.
Table 4. Eval Perplexity in the MoHID setting is performed for the Attention FFN of LLaMA2-355M (10B training budgets).
Method
# Activate
Perplexity +
LLaMA2-355M
289M
11.61
MOHHD-109%AATT--0000SS
289M
11.43 (-0.18)
4.3. Ablation Studies
195M
11,31 (-0.30)
To evaluate the importance of each method in Section 3 within MoHD, conducted detailed ablation experiments. In Table 3, we compare the ablation results of MoHD x2 with 710M parameters to the baseline with the same activa- tion under zero-shot pretraining on 10B tokens, based on Eval PPL.. The specific analysis is as follows:
MOHD-5U%ATTN, -
239M
12.25 (+0.64)
WOOHD-500%AATTT-5555557
145M
12.05 (40.44)
MOHD-100%ATTT--2555
147M
12.24 (+0.63)
MOHD-25%AATTN ATT
213M
14.31 (+2.70)
MOHD-25%%ATTN--555FFF
72M
13,20 (+1.59)
Mixed Activation Sub-Dimension Ablation. We ablated the Mixed Activation Sub- Sub-Dimension method by using fully specialized sub-dimensions without shared ones, leading to 0.83 increase in PPL, indicating negative impact performance. As discussed in Section 2.3, shared activation dimensions should be activated together in a mixed activa- tion mode. Figure 8a shows that this mode significantly outperforms fully sparse activation, highlighting its suitabil- ity for hidden dimensions patterns of the Transformer.
one for both, shown in Table 4 and Figure 11. All models were trained on 10B tokens. Sparsifying FFN showed more redundancy, leading to a minimal or even improved performance. We observe a -0.30 PPL reduction at 50% sparsity, likely due to the reduced redundant ac- tivations that mitigate overfitting. Sparsifying Attention resulted in a greater performance drop, with a +1.04 PPL increase at 50% sparsity, indicating its sensitivity to sparsification. Finally, joint sparsification of Attention and FFN achieved the best parameter efficiency, with the 50% ATTN-500%%
Balance Loss Ablation. The balanced loss effectively en- hances MoHD (0.16 improvement). It mitigates the risk of routing collapse, ensuring that most sub-dimensions are utilized more evenly. This increases the efficiency of sub- dimension utilization and improves the overall efficiency.
model reaching a PPI. of 1 12.05 using only 145M activated parameters, outperforming both 50%FFN and 50%ATTN configurations. More detailed analysis in Appendix H.1.
Flow Maintenance Ablation. Ablation experiments high- light that maintaining activation flow is crucial for MoHD. As shown in Table 3, removing Sub-dimension Scaling re- sults in a 1.16 performance drop, as the model loses critical information after sparsifying the hidden dimension. The Group Fusion Layer offers an additional 0.22 performance gain without significantly increasing parameters, improving dimension utilization while preserving information integrity.
Router Probability. To observe sub-dimension selection in MoHD, we visualize the attention and FFN router weight distributions at the 5th layer across five data domains in Figure 7. Each weight represents the average selection prob- ability across 4096 tokens. The sub-dimensions show spe- cialization across domains: for example, Attention Subdim is crucial for code-related data, with higher probabili- ties in GitHub and StackExchange. In contrast, Subdim 3 is more relevant to commonsense knowledge, with higher probabilities in Wiki, CC, and ArXiv. In the FFN router, Subdim 4 specializes in code tasks, while Subdim focuses on commonsense knowledge. This specialization validates MoHD's effectiveness allocating sub-dimensions based
4.4. Analysis
Decoupled MoHD Setting. To explore MoHD, we built three models: one sparsifying Attention, one for FFN, and
7

---

Mixture of Hidden-Dimennn Dimicn-o-Dmmsssn::::: Not All Hidden-States' Dimensions are Needed in Transformer
Proportion (%6
0.3
0.2
0.1
0.0
Wiki
CC
Github
StackExchange
Anxiv
Attn Sub Dim
Attn Sub-Dim
Attn Sub Dim
FFN Sub Dim
FFN Sub-Din Lub-Dim FFN Sub-Dim
Altn Sub-Din Dim
Altn Sub-Dim
FFN Sub-Dim
FFN Sub Rob--ii
Figure illustration Sparse Sub-dimension Routing Proportion in MoHD Attention Layer five Domains.
12.5
based MLPs (You et al., 2022; Li et al., 2023). This sparsity has been leveraged to improve LLM efficiency during in- ference. Liu al. (2023b) accelerated LLM inference by omitting zero-valued weight channels from GPU registers, while Song al. (2023) and Alizadeh et al. (2024) extend this to CPU offloading, reducing memory transfer overhead. Recent works Mirzadeh et al. (2023); Zhang et al. (2024); So et al. (2022); Song et al. (2024a;b); Wang et al. (2024b) have integrated activation sparsity into LLMs to boost effi- ciency, Lee et (2024) introduce CATS for training-free -ree sparsity in SwiGLU-based LLMs,, and Liu al. (2024) ex- tend training-free sparsity to large models. Building these studies, we focus on hidden dimension sparsity and continuous activation across tokens, leading to sparse acti- vation architecture that enhances parameter efficiency and hidden dimension scalability.
12.0
Transfomer
12.0
11.5
11.0
11.0
Eval
10.5
10.0
Transfomerr
10.0
3,5
0/4
1/4
2/4
3/4
4/4
102
128
Shared Sub.Dims/ Total Sub.Dims
Numbers Sub Dims
(a) Eval PPL. under different Shared Sub-Dim ratio settings.
(b) PPL. under different
Sub-Dim Numbers settings.
Figure Performance of MoHD x2-710M with varying sub- dimension allocation ratios and finer-grained sub-dimension set- tings. models are pre-trained from scratch on tokens.
data domains, enhancing parameter efficiency. Addi- tionally, sub-dimension probabilities generally stay within 0.2-0.3, ensuring all sub-dimensions are actively chosen.
Shared Activation v.s. Specialized Activation. Figure 8a shows the PPL values of MoHD x2-710M trained on 10B data with varying Shared Sub-dimensions and Specialized- dimensions proportions. The Baseline and MoHD models have identical total parameters. Specialized sub-dimensions effectively increases the hidden dimensions of the model, improving performance in the 0/4 setting. The best perfor- mance achieved with a 3/4 Shared activation ratio, proving that a number of sub-dimensions are still needed to model cross-token activation patterns.
5.2. Sparsely-activated Transformer
Sparse Transformer models, like Sparse Mixture-of-Exper (MoE) architectures, utilize input adaptivity to reduce com- putational overhead by activating only a subset of subnet- works, or "experts," " for each input token (Fedus et al., 2022: Riquelme et al., 2021; Zhou al., 2022; Jiang al., 2024a; Xue et al., 2024b; Gu et al., 2024; 2025). Recent develop- ments have introduced heterogeneous experts, integrating experts with varying capacities and specializations (Wu et al., 2024; He, 2024; Dean, 2021; Zhou et al., 2022; Dai et al., 2024). Some studies (Qiu al., 2024) have explored sparsely-activated architectures based on FFN intermediate activations. However, Transformer model has specifi- cally implemented sparse activation in hidden dimensions. Inspired by (Qiu et al., 2024; Dai et al., 2024), we ana- lyzed hidden dimensions and proposed a novel architecture, opening a new research direction. Extend related works in Appendix C.
Fine Gain Sub-dim Dimension. Figure shows the test PPL values of MoHD x2 -710M after pre-training 10B data with varying sub-dimensions. The best performance is achieved with 16 sub-dimensions (256 size). Increasing the sub-dimensions to 128 (32 size) and 256 (16 size) yields only marginal improvements. This indicates that more sub- dimensions do not enhance performance but increase com- putational costs. The results validate the effectiveness of the grouping fusion layer, showing that small number of parameters suffice to maintain efficient activation flow.
6. Conclusion
5. Related Work
In this paper, we presented MoHD, sparse conditional acti- vation architecture designed to address inefficiencies in scal- ing Transformer hidden dimensions. By integrating shared sub-dimensions for common token features and dynami- cally activating specialized sub-dimensions through a rout-
5.1. Activation Sparsity
Activation sparsity refers to the large proportion zero- valued hidden states models, naturally occurring in ReLU-
8

---

Mixture of Tlidden-Dimension Not All Hidden-States" Dimensions are Needed in Transformer
ing mechanism, MoHD achieves improved efficieney and flexibility while preserving activation flow through activa- tion scaling and group fusion mechanisms. Our evaluations demonstrate that MoHD outperforms standard Transform- ers across multiple NLP tasks, achieving superior parameter efficiency. These results underscore that MoHD provides a new direction for efficiently expanding model parameters.
perpositions of a sigmoidal function. IEEE Trans. Inf. Theory, 39:930 945, 1993. URL hwtph:tt////att /api. semanticscholas org/ CorpusID CorpusID:1538392 153883918.
Bartlett, P. L. and Mendelson, S. Rademacher and gaussian complexities Risk bounds and structural results. J. Mach.
Learn. Res., 3:463 482, 2003. URL https: /ppi.
semant semahticscoolar... org/ Corpus 463216.
Impact Statement
Bisk, Y. Zellers, R., Gao, J., Choi, Y., et al. Piqa: Reasoning about physical commonsense in natural language. In Pro- ceedings of the AAAI conference on artificial intelligence, volume 34, pp. 7432-7439, 2020.
In our study, we leverage open and online accessible data and techniques, mitigating privacy concerns. Our method emphasizes enhancing model parameter efficiency and re- ducing size to create powerful, compact, and openly ac- cessible models, thereby promoting the open dissemina- tion and democratization of NLP technologies. MoHD offer scalable solution for constraied environ- ments, reducing computational and memory costs without sacrificing capability. Its ability to match or surpass larger models highlights its potential to democratize access to peeeforeffrmeece language models, particularly in edge computing and low-resource settings. By redefining how hidden dimensions are utilized, this work paves the way for efficient large language model architectures, aligning with sustainability goals in Al development. Our work is commit- ted to advancing accessible and efficient NLP technologies, fostering more inclusive and automated future for Al.
Cai, R., Muralidharan, S., Heinrich, G., Yin, H., Wang, Z., Kautz, J., and Molchanov, P. Flextron: Many-in-one flexible large language model. ArXiv, abs/2406. 10260, 2024a. URL https: /ppi. semanticscholes org CorpusID:27056055:
Cai, R., Muralidharan, S., Heinrich, G., Yin, H., Wang, Z., Kautz, J., and Molchanov, P. Flextron: Many-in- one flexible large language model, 2024b. URL https: / //arxiv. //arxiv.org/ablll... 2400 10260.
Cai, W., Jiang, J., Wang, F., Tang. J., Kim, S., and Huang, J. A survey on mixture of experts, 2024c. URL https: //arxiv. org. /as/////////aa///////////aaa abs//b//////
Chen, T., Ding, T., Yadav, B., Zharkov, I., and Liang, L. LoRAShear: Efficient Large Language Model Structured Pruning and Knowledge Recovery, October 2023. URL https: / /aaxxvxvvvvv/////////a/////aavla hbtp://b///23111 18356v2.
Acknowledgments
We would like to thank Yinqi Yang, Yanxi Xie, Naibin Gu, Kun Huang and members of the IIE KDsec group for their valuable feedback and discussions. We are very grateful to Mengzhou Xia for providing the concise and effective SLearingLLLMMA experimental code and for her assistance during the reproduction process. Work done during Yilong Chen's internship Baidu Inc. This research is supported by the Youth Innovation Promotion Association of CAS (Grant No.2021151 153) and "Climbing" Program of IIE,AA (F370081 101).
Chen, Y., Shang, J., Zhang, Z., Cui, S., Liu, T., Wang, S., Sun, Y., and Wu, H. LEMON: Reviving stronger and smaller LMS from larger LMs with linear param- eter fusion. In Ku, L.-W., Martins, A., and Srikumar, V. (eds.), Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 4005-8019, Bangkok, Thailand, August 2024a. Association for Computational Linguis- tics. doi: 10. 10.11665311/2022.aclllooo 434, URL https: //aclanthology: www.:ollooo.oy././oooo
acl- -ong. 434.
References
Chen, Y., Wang, G., Shang, J., Cui, S., Zhang, Z., Liu, T., Wang, S., Sun, Y., Yu, D., and Wu, H. Nacl: A general and effective kv cache eviction framework for Ilm at inference time, ArXiv, abs/2408.03675, 2024b. URL https: //ppi.pppp/////// semanticscholar. org/ wtgg/orrpoTT::21744555
Alizadeh, K., Mirzadeh, I., Belenko, D., Khatamifard, K., Cho, M., Del Mundo, C. C., Rastegari, M., and Farajtabar, M. Llm in a flash: Efficient large lan- guage model inference with limited memory. arXiv preprint arXiv:2312.11514, 2024. URL https:// trxiv.org/aby. Ntrg/abb/23111 11514.
Chen, Y., Zhang, L., Shang, J., Zhang, Z., Liu, T., Wang, S., and Sun, Y. Dha: Learning decoupled-head attention from transformer checkpoints via adaptive heads fusion. ArXiv, abs/2406.06557,, 2024c. URL https://app. org/CorrusTD::2703990
Anthropic. Anthropic: Introducing claude 2.1, 2023. URL https: /twa:://www.. www.rapittccccoo com/ index/cll fodex/claudes claude-2-1.
Barron, A. R. Universal approximation bounds for su-

---

Mixture of Hidden- sindden-Dimensions: Not All Hidden-States" Dimensions are Needed in Transformer
Clark, C., Lee, K., Chang, M.-W., Kwiatkowski, T., Collins, M., and Toutanova, K. Boolq: Exploring the surprising difficulty of natural yes/no questions. arXiv preprint arXiv:1905.1004, 2019.
Gu, N., Zhang, Z., Liu, X., Fu, P., Lin, Z., Wang, S., Sun, Y., Wu, H., Wang, W., and Wang, H. Beamlora: Beam- constraint low-rank adaptation, 2025. URL https:// arxiv. wwg.
arg/abs////e abs 211. 13604.
Clark, P., Cowhey, I., Etzioni, O., Khot, T.. Sabharwal, A., Schoenick, C., and Tafjord, O. Think you have solved question answering? try arc, the A12 reasoning challenge. CoRR, abs/1803.05557, 201 8a. URL http:://txxiiti org/abs/1888 05457.
He, X. Mixture of million experts. arXiv preprint arXiv:2407.04153, 2024. URL https: bttps://arxiy. /arxiv. org/ hts/aba/22400 2005/2007.4453.
Jain, G., Hegde, N., Kusupati, A., Nagrani, A., and Buch, Mixture of nested experts: Adaptive processing of visual tokens. arXiv preprint arXiv:2407. 2024. URL https://arxiv.orgg
Clark, P., Cowhey, I., Etzioni, O.,, Khot, T., Sabharwal, A., Schoenick, C., and Tafjord, O. Think you have solved question answering? try are, the ai2 reasoning challenge. arXiv preprint arXiv: arXiv:/803.0m 1803.004377 2018b.
Jiang, A. Q. Sablayrolles, A., Roux, A., Mensch, A., Savary, B., Bamford, C.,, Chaplot, D. S., Casas, D. d. I., Hanna, E. B., Bressand, F,, et al. Mixtral of experts. arXiv preprint arXiv:2401.04088 2024a.
Dai, D., Deng, C., Zhao, C.,, Xu, R. X., Gao, H., Chen, D., Li, J., Zeng, W., Yu, X., Wu, Y., Xie, Z., Li, Y. K., Huang, P., Luo, F., Ruan, C., Sui, Z., and Liang, W. Deepseekmoe: Towards ultimate expert specialization in mixture-of-experts language models, 2024. URL http:::: //arxiv. org/ abs / 2401. 06066.
Jiang, A. Q. Sablayrolles, A., Roux, A., Mensch, A..
Savary, B., Bamford, C.,, Chaplot, D. S., de las Casas, D., Hanna, E. B., Bressand, F,, Lengyel, G., Bour, G., Lample, G., Lavaud, L. R., Saulnier, L., Lachaux, M.- A., Stock, P., Subramanian, S., Yang, S., Antoniak, S., Scao, T. L., Gervet, T., Lavril, T., Wang, T., Lacroix, T., and Sayed, W. E. Mixtral of experts, 2024b. URL https:
Dao, T., Chen, B., Sohoni, N., Desai, A., Poli, M., Grogan, J., Liu, A., Rao, A., Rudra, A., and Ré, C. Monarch: Expressive Structured Matrices for Efficient and Accu- rate Training, April 2022. URL http: www.i.vvvv.o.....o.............. org/ abs/ 2204. 00595. arXiv:2204.0050 [cs].
04088.
Kaplan, J., McCandlish, S., Henighan, T., Brown, B., Chess, B., Child, R., Gray, S., Radford, A., Wu, J., and Amodei, D. Scaling laws for neural language mod- cls, 2020. URL https: /tta//iiivv/../////v/ org/ abs/ 08361.
Dean, J. Introducing pathways: A next-generation ai archi- tecture. Google Blog, 366, 2021.
Dong, P., Li, L.,, Tang, Z., Liu, X., Pan, X., Wang, Q. and Chu, X. Pruner-zero: Evolving symbolic pruming metric from scratch for large language mod- cls. ArXiv, abs/2406.02924, 2024. URL https: //api. 270257857.
Lee, J. et al. Cats: Training-free activation sparsity for swiglu-based llms. arXiv preprint, 2024. URL https: // arxav.. wwt.o/////aooo hbsg://b///// htt.s///44..... 12345.
wwwpi/aarrrrg//oopo Crg/CorppsID:
Li, M., Gururangan, S., Dettmers, T., Lewis, M., Althoff, T., Smith, N. A., and Zettlemoyer, L. Branch-train merge: -raim---e-----
Embarrassingly parallel training of expert language mod- cls. arXiv preprint arXiv:2208 (rX.::2088.03306, 2022.
Fedus, W., Zoph, B., and Shazeer, N. Switch transform- ers: Scaling to trillion parameter models with simple and efficient sparsity. The Journal of Machine Learning Research, 23(1):5232-5277 2022.
Li, Z., You, C., Bhojanapalli, S., Li, D., Rawat, A. S., Reddi, S. J., Ye, K., Chern, F, Yu, F, Guo, R., and Kumar, S. The lazy neuron phenomenon: On emer- gence of activation sparsity in transformers. arXiv preprint arXiv:2210.063333 2023. URL https:// arxiv wwl.voorg/aaae wttpaa////22..22222
Gao, I.,, Tow, J., Biderman, S., Black, S., DiPofi, A., Foster, C., Golding, L., Hsu, J., McDonell, K., Muen- nighoff, N., Phang, J., Reynolds, L., Tang, E., Thite, A. Wang, B., Wang, K., and Zou, A. A frame- work for few-shot language model evaluation. In Zen- odo. https://doi.org/10..o 12we://..000 September 2021.
Liu, Cui, L.,, Liu, H., Huang, D., Wang, Y., and Zhang, Y. Logiqa: A challenge dataset for machine reading comprehension with logical reasoning. arXiv preprint arXiv:2007.08124, 2020.
Gu, N., Fu, P., Liu, X.. Shen, B., Lin, Z., and Wang, W. Light-peft: Lightening parameter-efficient fine-tuning via early pruning, 2024. URL https:/ /aaxiv.vorr. org/ abs/2406. 03792.
Liu, J., Ponnusamy, P., Cai, T., Guo, H., Kim, Y., and Athiwaratkun, B. Training-free activation sparsity in large language models, 2024. URL https: pitoos///rriiv org/ 2albs/2400 14690.
10

---

Mixture of Hidden- seiden-Dimmensios::: Not All Hidden-States" Dimensions are Needed in Transformer
Liu, Z., Wang, J., Dao, T., Zhou, T., Yuan, B., Song, Z., Shrivastava, A., Zhang, C., Tian, Y., Re, C., and Chen, B. Deja vu: Contextual sparsity for efficient LLMs at inference time. In Krause, A., Brunskill, E., Cho, K. Engelhardt, B., Sabato, S., and Scarlett, J. (eds.), Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Ma- chine Learning Research, pp. 22137-22176. PMLR, 23- 29 Jul 2023a. URL https: httttts::/ppoeeeeeteeeeeee http:////ccceeeee...ee.o. mmlr.
144. URL hatp:s//aallaatthhloooooooo naac1-long.
Riquelme, C., Puigcerver, J., Mustafa, B., Neumann, M., Jenatton, R., Susano Pinto, A., Keysers, D., and Houlsby, N. Scaling vision with sparse mixture of experts. In Advances in Neural Information Processing Systems, vol- ume 34, pp. 3583-8595 2021.
Sakaguchi, K., Bras, R. L.,, Bhagavatula, C., and Choi, Y. Winogrande: An adversarial winograd schema challenge at scale. In The Thirty-Fourte AAAI Conference on Ar- rificial Intelligence, AAAI 2020, The Thirty-Second In- novative Applications of Artificial Intelligence Confer- ence, IAAI 2020, The Tenth AAAI Symposium on Educa- tional Advances in Artificial Intelligence, EAAI 2020, New York, NY, USA, February 7-12, 2020, pp. 8732- 8740. AAAI Press, 2020. doi: 10.1609/AAAL 6399. URL http::://doo..rg/111// aaai..
Fress/v202/110 /tt/////22112222 html.
Liu, Z., Wang. J., Dao, T., Zhou, T., Yuan, B., Song, Z., Shrivastava, A., Zhang, C., Tian, Y.. Re, C., and Chen, B. Deja vu: Contextual sparsity for efficient lims at inference time. arXiv preprint arXiv:2310.17157, 2023b. URL https:ttp/// /tt////axxiii wttpl//.rrgbbb////33e
Ma, X., Fang, G., and Wang. X. LIm-pruner: On
the structural pruning of large language models.
ArXiv, abs/2305.11627, 2023a. URL https:
v34i05. 6399.
//api..seantiicccoollaa htt.://cccoooaaaooroo/oroopppoas 258823276.
Snell, C.,, Lee, J., Xu, K., and Kumar, A. Scaling Ilm test-time compute optimally can be more effec- tive than scaling model parameters. arXiv preprint URL
Ma. X., Fang, G., and Wang, X. LLM-Pruner: On the Structural Pruning of Large Language Models, Septem- ber 2023b. URL http: /ttp:///aaxiv.v// http://xwwv..or/ab//222/0./ 11627.
arXiv:2408.0555 03314, 2024. org/ abs
03314.
So, D. et al. Squared relu: A simple and effective activation function. Advances in Neural Information Processing Systems, 2022.
Mirzadeh, I., Alizadeh, K., Mehta, S., Del Mundo, C. C.,
Tuzel, O., Samei, G., Rastegari, M., and Farajtabar, M. Relu strikes back: Exploiting activation sparsity in large 2023.
Song, C., Han, X., Zhang, Z., Hu, S., Shi, X., Li, K., Chen, C., Liu, Z., Li, G., Yang, T., and Sun, M. Prosparse: Introducing and enhancing intrinsic activation sparsity within large language models. arXiv preprint arXiv:2402. 2024a. URL https: //axivv. org/ a /t.b////422.m 13516..
language models. arXiv preprint
URL https:tt///// //a:://aaaa http//axxvvv....aaaa/// wrt. abs////////// 2311 04564.
Muennighoff, N., Yang. Z.. Shi, W., Li, X. L., Fei-Fei, L., Hajishirzi, H., Zettlemoyer, L., Liang, P., Candès, E., and Hashimoto, T. sl: Simple test-time scaling. arXiv preprint 2rXiv:2501.19399; 2025. URL https: //arxiv. awwrxii.org/abs///// .291119393. 19393.
Song, Y., Mi, Z., Xie, H., and Chen, H. Powerinfer: Fast large language model serving with a consumer-grade gpu, arXiv preprint arXiv:2312. //arxiv. wwt/ /ta.ab/2333 /2022 12456.
OpenAl. Openai: Gpt-4, 2023. URL htttt:://ppeeee com/ www.eercchhhee 4.
2023. URL https:
Paperno, D., Kruszewski, G., Lazaridou, A., Pham, Q. N., Bernardi, R., Pezzelle, S., Baroni, M., Boleda, G., and Fernández, R. The lambada dataset: Word prediction requiring a broad discourse context. arXiv preprint 2016.
Song, Y. et al. Powerinfer: Enhancing activation spar- sity in Ilm serving with consumer-grand grade gpus. arXiv preprint arXiv:2312. 2024b. URL https:// arxiv .allv.orge httgg/bbs/3311. 12456.
Team, T. M. M. composer. htttp:://ittubb.co mosai cml ccoppsee// 2021.
Qiu, Z., Huang, Z., and Fu, J. Unlocking emergent mod- ularity large language models. In Duh, K., Gomez, H., and Bethard, S. (eds.), Proceedings of the 2024 Con- ference of the North American Chapter of the Associa- tion for Computational Linguistics: Human Language Technologies ( (Volume 1: Long Papers), pp. 2638 (288-,660, Mexico City, Mexico, June 2024. Association for Compu- tational Linguisties. doi: 10.1 8653/v 1///v/1//20 /t//242222aacc naacl-long.
TogetherAI. Redpajama: An open source recipe to repro- duce llaaa training dataset, 2023.
Touvron, H. Martin, L., Stone, K., Albert, P., Almahairi, A., Babaei, Y., Bashlykov, N., Batra, S., Bhargava, P., Bhosale, S., Bikel, D., Blecher, L., Ferrer, C. C., Chen, M., Cucurull, G., Esiobu, D.,, Fernandes, J., Fu, J., Fu, W.,

---

Mixture of Grndden-Dimensinnn Not All Hidden-States" Dimensions are Needed in Transformer
Fuller, B., Gao, C., Goswami, V., Goyal, N., Hartshorn, A., Hosseini, S., Hou, R., Inan, H., Kardas, M., Kerkez,
International Conference on Learning Representations,
2022. URL https: / appi ipp..... www.anticccho.........m.. org/ /tt//CopuuD::2511388088
V., Khabsa, M., Kloumann, I., Korenev, A., Koura, P. S., Lachaux, M.-A., Lavril, T., Lee, J., Liskovich, D., Lu, Y., Mao, Y., Martinet, X., Mihaylov, T., Mishra, P., Molybog, I. Nie, Y., Poulton, A., Reizenstein, J., Rungta, R., Saladi, K., Schelten, A., Silva, R. Smith, E. M., Subramanian, R., Tan, X. E., Tang. B., Taylor, R., Williams, A., Kuan, J. X., Xu, P., Yan, Z., Zarov, I., Zhang, Y., Fan, A., Kambadur, M., Narang, S., Rodriguez, A., Stojnic, R., Edunov, S., and Scialom, T. Ilama 2: Open Foundation and Fine- Tuned Chat Models, July 2023a. URL RtLhttp:// http:///axxi atg/abs/aae 2307 09288.
Xue, F., Zheng. Z.., Fu, Y., Ni, J., Zheng, Z., Zhou, W., and You, Y. Openmoe: An early effort on open mixture- of-experts language models, 2024a. URL https:// arxiv. www.orggaaaa hac//abs//t/ /402 01739.
Xue, F. Zheng, Z.,, Fu, Y., Ni, J., and Zhou, W. Open- moe: An early effort on open mixture-of-expprts language models. arXiv preprint arXiv:2402.01739) 2024b. URL attps://arxis org wwwpaaa//2222017.
You, C.,, Bhojanapalli, S., Li, D., and Rawat, A. The lazy neuron phenomenon: On emergence of activation sparsity in transformers, arXiv preprint arXiv:2210.066: 2022. URL https:://///.. /arxiv orq/ abs 2210 06313.
Touvron, H., Martin, L., Stone, K., Albert, P., Almahairi, A., Babaei, Y., Bashlykov, N., Batra, S., Bhargava, P., Bhosale, S., et al. Llama 2: Open foundation and fine- tuned chat models. arXiv preprint arXiv:2307.09288, 2023b.
Zellers, R., Holtzman, A., Bisk, Y., Farhadi, A., and Choi, Y. Hellaswag: Can machine really finish your sentence? In Korhonen, A., Traum, D. R., and Mårquez, L. (eds.), Proceedings of the 57th Conference of the Association for Computational Linguistics, ACL. 2019, Florence, Itlly, July 28- August 2, 2019, Volume 1: Long Papers, pp. 4791-4800. Association for Computational Linguistics, 2019. doi: 10.18653/V11/199 1472. URL https:// doi. 10.
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L., and Polosukhin, I. Attention is all you need, 2023. URL https: abs 2os/17055 03762.
Wang, H. Ma. S., Wang, R., and Wei, F. Q-sparse: All large language models can be fully sparsely- activated. ArXiv, abs/2440.11096 2024a. URL. https: h/app.semantiic.......eaaaaaaaaa commanticccholar.. org/CorpusI:: 271212835.
Zhang, W. Liu, T., Song, M., Li, X., and Liu, T. SOTOPIA- Ω: Dynamic strategy injection learning and social instruc- tion following evaluation for social agents, 2025a. URL https: / arxiiv www.oorg/bbb..... /2002 11.5222.1555 15538.
Wang, L., Ma, L., Cao, S., Zhang, Q. Xue, J., Shi, Y., Zheng, N., Miao, Z., Yang, F., Cao, T., Yang, Y., and Yang, M. Ladder: Enabling efficient low-precision deep learning computing through hardware-aware tensor trans- formation. In 18th USENIX Symposiun on Operating Systems Design and Implementation (OSDI 24), 2024b. URL htthtts::.. htwp::/:://///wwwwwi.it/./. usenix org/ conference/ lei.
Zhang, W., Nie. S., Zhang, X., Zhang, Z., and Liu, T. S1- bench: simple benchmark for evaluating system 1 think- ing capability of large reasoning models, 2025b. URL htap//bb//500000 10368.
https /tr///axivvv
osdi 2 http///eeeeettttteeeeeeeeeeeeeee
Zhang, Z., Song, Y., Yu, G., Han, X., Lin, Y., Xiao, C., Song, C., Liu, Z., Mi, Z., and Sun, M. Relu2 wins: Discovering efficient activation functions for sparse llms. arXiv preprint arXiv:2402.03804 2024. URL https: / //arxiv. wwtg abs/ 03804.
Welbl, J., Liu, N. F., and Gardner, M. Crowdsoure- ing multiple choice science questions. arXiv preprint (007:v/777.66200000 2017.
Wu, X., Huang, S., and Wei, F. Multi-head mixture-of experts. arXiv preprint arXiv:2404. 15045, 2024. URL https:/ arxiv. org/ abs 2404.444400 15045.
Zhou, Y., Lei, T., Liu, H., Du, N., Huang, Y., Zhao, V., Dai, A. M., Le, Q. V., Laudon, J., et al. Mixture-of- experts with expert choice routing. In Advances in Neural Information Processing Systems, 2022.
Xia, M., Gao, T., Zeng, Z., and Chen, D. Sheared LLaMA: Accelerating Language Model Pre-training via Structured Pruning, October 2023. URL http:////////tiv.. abs/ 2310. 06694.
xiao Li, Z., You, C., Bhojanapalli, S., Li, D., Rawat, A. S., Reddi, S. J., Ye, K. Q., Chern, F., Yu, X., Guo, R., and Kumar, The lazy neuron phenomenon: On emergence of activation sparsity in transformers. In
12

---

Mixture of Tlidden-Dimensionss Not All Hidden-States" Dimensions are Needed in Transformer
A. Definition
A.3. Hidden Sparsified Attention
In this Section, we define the activation sparsity present in the hidden dimension of LLMs and use this to formulate sparsely activated FFN and Attention mechanisms.
For ht-ead Mult-Head--Atenniinn (MHA), we define WP, WK, wy E R WP E R as key, query, value and output projections for the i-th head, where d, denotes as the head dim, S [h.] With sparsely activated hidden state a.. a small parameter subset SA represents sparsely activated selection of rows from WP, W WV and columns from WP-
Let X E Raxd denote the embeddings of n tokens, and E RRddd represent the embedding of a single input token. The activation sparsity ó of a hidden state æ is defined as the proportion of zero-valued entries within the vector (Liu et al., 2024). We then define function S d - åd that selectively activates a subset of dimensions in x. The sparsely activated representation is denoted as I, S(x, 6), where v E R x8d, representing the selective activation of 0- -proportion of the dimensions in J.
MHASA (I.))
Σ Head, W.SAA
(12)
Head: =o - 1,5H (x. W.K W1.5
vdn 2. 1sp (13)
where a is the softmax function. Since x is sparse in the hidden dimension, we can find an approximation Js of a, such that, under the activation of the corresponding subset of parameters Sn, the outputs of the sparsified FFN and sparsified attention closely approximate the outputs of the dense model.
A.1. Hidden Dimension Sparsity
Considering the model's semantic modeling in Euclidean space, we define the magnitude m, of each dimension i as the square of its activation value:
B. Theoretical Proof for Mixed Sparse Activation Superiority
We define hidden dimension sparsity as:
1 à 11(
In this section, we show that a mixed sparse activation scheme-where some hidden dimensions are shared across (10)
Sparsity
1(., <)),
tokens and others are token-speciffc-yyeldds strictly better risk bounds than a purely token-only activation. We proceed in three steps:
where d is the total number of hidden dimensions, x; repre- sents the squared activation value of the i-th dimension, and 6 is small threshold used to identify near-zero activation values. The indicator function 1 (zz E) is equal to I if the activation value is below the threshold and 0 otherwise.
Show that applying random sparsification to an n- dimensional hidden vector is, in expectation with re- spect to second moments, equivalent to training a wider network of dimension n/p, where is the retention probability. ("Width Expansion")
A.2. Hidden Sparsified FFN
Define WUP, Waste dgd Widoom E Igd'xe xd as the up, gate, down matrix in one FFN block, where d' is the in- termediate size. In this context, the i-th row of the up, gate matrix is defined as WHP, Wasate E lxdd' and the i-th column of the down matrix is defined Welnen Rdxx1 Specifically, the sparsely activated hidden state v, under activation sparsity 8 only activates a subset of rows in the up, gate matrix and a corresponding subset of columns in the down matrix, denoted as Sy S [d]. Thus, the sparsified FFN computation can be described as follows:
2. Use classical Barron-space approximation results to argue that a network of larger width has smaller ap- proximation error. Then compare a purely token-level network (width n..)) to a mixed network (width n, nn.). ("Approximation Error Decay")
Show that parameter-tying in the shared dimensions reduces the Rademacher complexity compared hold- ing all parameters token-specific. ("Complexity Reduc- tion""
FFN (z.) weloon (o (Nx Wup @ x. Weate) SM
(11)
Combining these facts under explicit quantitative conditions, we obtain strictly lower risk bound for the mixed activation scheme.
where is the activation function. O is the element-win wise production. Due to the sparsification of the hidden state, the up and gate matrices share the same activation subset. To ensure the output remains sparsified, the down matrix is also sparsified, though its activation subset can differ from that of the up and gate matrix.
B.1. Notation and Definitions
Let n be the total number of hidden units in a given layer. We index them by j E (I....,n) n..
13

---

Mixture of Hidden-Diinn DinicnnoDDmmssin:::: Not All Hidden-States" Dimensions are Needed in Transformer
Denote by h E R deterministic hidden-layer repre- sentation. Its j-th coordinate is htj
B.2. Step 1: Sparse Activation Efffective Width Expansion (Lemma 1)
We apply random mask E {0, 1) whose coordi- nates ry are i.i.d. Bernoull(p)), i.e.
Lemma B.1 (Unbiased Sparse Forward Pass). Let h R" be any deterministic vector. Let 10, be an iid. Bernoulli(p) mask, and define
Pr(r, = 1) = P. Pr(r, = 0) = 1
independently for j ..... Define the sparsely activated version of h by
hj
ifr; = 1,
h =
p
0,
if 0.
h2
if 1,
h2
p
Then:
1
0,
rf O,
E h1
h,
B[[
p
that h (r b) and in particular (r 2 h)j ry hj.
where Elry) = p and r3 rj for each Bernoulli(y) vari- able.
In what follows, write expectations IE over the ran- domness of r. We note that r ry for each Bernoulli variable, and E[ry] = p.
Proof.
For each coordinate j I,..., n, since Tj Bernoulli(p),
We let
||2 denote the Euclidean norm on R".
hj
with probability
ho
For a function class H. and a sample size m, Radm (H) denotes its (empirical) Rademacher complexity, under the same norm-andd-Lipshinn constraints on parameters detailed below.
0,
with probability p-
Hence
Ehhy =p hi p
We use e(n)) to denote the approximation error of n-width network when approximating target f". In particular, if f belongs to a Barron space with Bar- ron norm Cý (i.e., its hidden-layer weight & norm is bounded by Cу) (Barron, 1993), then a classical result yields
+ (1 (1-P) p)- hy, for all j,
s0 [E h. Next, since hy = (ri/p) hy and rž ry.
E
Σ E 22 hz)
(n) S C₁ n
where the - norm is taken with respect to the data distribution on the input domain.
1
123 (p) = 11112
j=1
We will compare:
This completes the proof.
A token-only network of width n. (all hidden units token-specific). Its approximation error is denoted eg = e(n.).
Corollary B.2 (Effective Width Expansion). Fix hidden- layer dimension n and sparsity = (so that k = is the expected number of nonzero coordinates). Then, noting that || p > [|h one sees that terms of second- moment (Euclidean norm) statistics, a network of physical width n with per-step sparsity p and rescaling 1/p be- haves analogously to a full (dense) network of effective width
A mixed network of width n, nue where n. dimensions are shared across all tokens and nh. dimensions remain token-specific. Its approxima- tion error is c = c(n,, + na).
Finally, for a network h(.), we denote its population risk as
n²
n'
P
k/n
k
R(h) E(h) C. (Radm.H),
In particular, if < (so p 1), then n' n/p im- plying that sparse activation in expectation with respect to second moments expands the hidden-layer norm as though the network were wider. (This equivalence refers only to those two Euclidean norm properties and does not assert full functional equivalence under arbitrary nonlinearities.)
where ((h) is the approximation (Bayes) enror under the data distribution, and C' > 0 is a constant depend- ing on the Lipschitz constant of the loss. We use subseripts A or B to indicate mixed vs. token-only, respectively.
14

---

Mixture of Wlideen-iimension Not All Hidden-States" Dimensions are Needed in Transformer
Proof. By Lemma B. each hidden vector h E Rd is re- placed by satisfying B/h] = h and | = [1 dense network of width n/p whose hidden vector scaled by /p would satisfy / k [ = p || Thus, in terms of those two aconddmoment statistics, the sparse-activated width-n network simulates dense width-(n/p) network. Since n/p > n whenever p < L, this establishes the corol- lary.
A mixed network allocate ng units to g(s(x)) (shared) and nu units to Es u(xx) (token-specific), for total width ny + ny. By Lemma B.3, one obtains
Cg. n4
Ilfa -5ll2112 S E.
inf
T TC
inf [fu(z;) uxx:)lzz
B.3. Step 2: Width Expansion Reduces Approximation Error (Lemma? 2)
n4
We now invoke a classical result for Barron-type function spaces.
Therefore, the mixed network's 's overall approximation error satisfies
Lemma B.3 (Approximation Error Decay with Width). (Barron, 1993) Let F be a Barron Barroffunction class on the data domain, and suppose f F has Barron norm Cj. Then any two-layer ReLU network of width n approximate f" with uniform 2 error
EA
inf
-
fEHo,-T
Cg
TCo
+
n1x
nhu
e(n) = inf llf S C
In particular, if we set thu for some >> 0, then
(an)=inff]] SEHn
Cg
TC₂
Cg + TC.
cA S
+
where Hn, denotes the class of two-layer ReLU networks of width n, and the L norm is taken with respect to the data distribution D.
a nyy
n4e
nlu
nuu
Meanwhile,
TCg
Corollary B.4 (Mixed Activation Lowers Approximation Error), Suppose the target function can be decomposed as
>
Thu
Hence
T
TCg
Cg
T
fz)x) = g(s(z)) +
Eu(zi)
CB cA 2
a nl
nu
1=1
Cg
where
T - 1 - TTT GC,)
nv
a
= (EI,..., .T) contains T tokens,
Thus, provided
sax) E Rd is shared summary representation of all tokens,
T - 1
T Cv
i.e.
C,
T(1 Co/
9 Rd R nonconstant "shared-dimension" function with Barron norm Cg.
there exists a strictly positive gap
8 = cB c = 9 1
u. Ritoeen-dim) R is "token-specific" function with Barron norm C. (applied to each zz).
Proof. By definition, a token-only network of width n. cannot directly allocate any units the shared term g(s())) unless it replicates subnetwork parameters across T tokens, which yields an effective width for g of most n,/T. Hence
Let:
A token-only network allocate all width to modeling u(-) over each token, i.e. width nu in each token's 's subnetwork. Its total width is n = Tui hence its ap- proximation error satisfies (conservatively)
TC2
BB = <(n.) 2
nu
Meanwhile, mixed network of width n, + T ny can dedi- cate n, units to approximate g(s(z)) and n, units per token to approximate each u(xx). By Lemma B.3,
TC2
BB = c(m
2
nh
since in the worst case one must dedicate at most nu/ units per token to approximate g(x(x)))
Cg ng
T
C
TC₁
Eg <
cu S
Bu
n2
15

---

Mixture of Wlidden-Dimensionn Not All Hidden-States" Dimensions are Needed in Transformer
Thus
B.S. Putting It All Together: Risk Bound Comparison
C. n
C
6
+
Theorem B.6 (Mixed Activation Yields Strictly Lower Risk), Assume all networks in H (mixed) and H. (token- only) have the same total number of parameters and impose identical &g norm bounds B on their weights. Suppose the target function decomposes as in Corollary B.4, and that activation functions share the same Lipschitz constant L. Fix any sample size m.
nu
Setting n, = ?u yields the claimed expression for d. The stated condition T (1/o)= > T (Cn/Cg) ensures 8 > 0.
B.4. Step 3:
Mixed Activation Reduces Rademacher Complexity (Lemma 3)
Next, we argue that tying n, shared units across all T tokens lowers the Rademacher complexity compared keeping each of the T tokens' hidden units disjoint in token-only network.
If the following two conditions hold:
TC
(
TCa
+
>
0,
Tii
nu
Lemma B.5 (Rademacher Complexity Reduction via Param- eter Tying). Let H. be the hypothesis class corresponding to token-only networks of total width T nue where each token subnetwork has width ny. Let HA be the class of mixed networks with total width n n. + T Dww in which n, shared units are tied across all T tokens and the remain- ing T ny units remain token-specific. Suppose both classes impose the same ly norm bound B on their parameters, and both use activation functions with the same Lipschitz constant L. Then, for any sample size m,
820
C RRadm (HB) - Rad (Hl))
> 0,
CAS0
then
RA = cA + C C'Rad (HA)
< CB Cadm.
RB,
i.e. the mixed sparse activation network has strictly lower population risk than the token-only network.
Radm (Ha) < Radm (HB)
A,
A
O
ns Vm
Proof. By Corollary B.4, under the stated condition we have = ég - cA > 0. By Lemma B.5, under the norm- and-Lipschitz constraints,
Proof. Under a uniform & norm bound B on weights and Lipschitz activation constant L, standard covering-numble arguments (see, e.g., (Bartlett Mendelson, 2003)) yield
Rad (H.) I Radm (HB) A, A θ
0.
Hence, for the same constant C 0 in the risk decomposi- tion,
Radm (HR) = OBL T n2 m
RA = C + C' C'Rad (Ha))
S (en 8) + (Radm (Hn) A)
Rad,, (Ha)) o (BL. ns + Tn, (T - 1) nx
= (cB + C Rad,.. (HB)) (6 + CA).
m
Since both 8 > 0 and A > 0, it follows that
= o BL
T nu
m
RA < RB-
Therefore,
B.6. Discussion
A = Radm (H.) - A)LA)
Width Expansion (Lemma B.1 & Corollary B.2). By randomly dropping coordinates with retention prob- ability p = k/n and rescaling by 1/p, a network of "physical width" n behaves, in expectation with respect to second-moment statistics, like a full net- work of width n/p. Therefore, applying sparsity to a fixed-width network implicitly trains wider net- work in that sense, which by Barron's 's theorem (Lemma B.3) lowers the approximation error from ((1/n) to Opp/n) = Ok//n"), ). provided the target indeed lies in a Barron space.
.BB VTn, n2 VT nu
Vm
VT' vm
since T Tn, and Tn, differ by T(n. - nw) and n = n, + Tn, is held fixed. Hence parameter tying reduces the Rademacher complexity by an amount on the order of / nz/mm m.
16

---

Mixture of Hiidenn-Dimensionn Not All Hidden-States" Dimensions are Needed in Transformer
Approximation Error Gap (Corollary B.4). A purely token-only network of width nu cannot efficiently capture any shared global mapping g(s(x)) without replicating the same parameters for each of the T to- kens, leading to an approximation error lower bound EB 2 T Tg In contrast, a mixed network of total width n = ns + T The can devote Tis units to model g(s(z)) and ny units per token to model Es u(r;). yielding CA Cg/n, + T Cu//nu When- ever T / > (Cg/n, + T TC//n.)) there is strictly positive gap 8 = (11mnn).
complexity by A (((m/z)) which further lowers the generalization term in the risk.
Conclusion (Theorem B.6). Combining the strictly smaller approximation error with strictly smaller com- plexity term, the mixed activation scheme achieves strictly smaller population risk than token-only activa- tion, even when the two networks are constrained to have the same total number of parameters.
C. Extended Related Work
Complexity Reduction (Lemma B.5). Tying ns shared units across all T tokens reduces the effective hypothesis-class size: instead of having Tn, inde- pendent parameters to model (z)) per token, there are only ns global parameters. Under a uniform & norm bound and Lipschitz activations, this reduces the Rademacher complexity by A @(vn,, n//m/. /)). which further lowers the generalization term in the risk.
The design of the MoHD is indeed inspired by prior re- search (Fedus et al., 2022: Riquelme et al., 2021: Zhou et al., 2022: Liu al., 2024; Chen et al., 2024b;c). In addi- tion to some related work given in the article, we compared more research work related to this article and express our sincere thanks.
C.1. Activation Sparsity
Activation sparsity (xiao Li et al., 2022; Ma et al., 2023a; Dong et al., 2024; Wang et al., 2024a) refers to the phe- nomenon where significant proportion of a model's hidden states are zero- valued. This property naturally arises in the intermediate states of ReL.U-based MI.Ps, as demonstrated in prior work (You et al., 2022; Li al., 2023). Some studies have leveraged activation sparsity to improve the efficiency of LLMs during inference. Liu et al. (2023b) utilized activa- tion sparsity to accelerate LLM inference by omitting the transfer of weight channels corresponding to zero-valued entries to GPU registers. Additionally, Song et al. (2023) and Alizadeh et al. (2024) extended this concept to CPU offloading, significantly reducing memory transfer overhead between CPUs and GPUs. Recent works has reintroduced activation sparsity into LLM architectures to enhance effi- ciency. Mirzadeh et al. (2023) replaced SiLU and GeLU with ReLU, achieving sparsity through extended pretraining. Zhang et al. (2024) identified Squared ReLU (So et al., 2022) as a superior alternative for sparse activations. Song et al. (2024a:b) proposed regularization techniques to increase sparsity, while Wang et al. (2024b) combined pruning and quantized activations to establish scaling laws. Lee et al. (2024) introduced CATS, achieving training-fres free sparsity in SwiGLU-based LLMs. Liu et al. (2024) extended these concepts to training-free activation sparsity for large-scale language models. Recent approaches, such as Test-Time Scaling ("Slow-Thine (ailk-Thi (Muennighoff al., 2025; Snell et al., 2024: Zhang al., 2025b;a), aim to enhance perfor- mance by allocating more computation during the inference search process, but it also further increases the need for com- putational optimization during model inference. Building on prior studies, we investigate hidden dimension sparsity, focusing on continuous activation across tokens. Leveraging this, we design a sparse activation architecture that improves
Conclusion (Theorem B.6). Combining the strictly smaller approximation error EA < EB with strictly smaller complexity term Rad,.. (Ha) < Rad, (H the mixed activation scheme achieves strictly smaller population risk than the token-only activation, provided the explicit conditions on a, Cg. Cu. T, na Nuy m are satisfied.
B.7. Discussion
Width Expansion (Lemma B.1 & Corollary B.2). By randomly dropping coordinates with retention probabil- ity p = k/n and rescaling by l/p. network of "physi- cal width" behaves in expectation like a full network of width n/p. Therefore, applying sparsity to fixed- width network implicitly trains a wider network, which by Barron's 's theorem (Lemma B.3) lowers the approxi- mation error from (1/n) to Opp/n) /n) = Okknn))
Approximation Error Gap (Corollary B.4). A purely token-only network of width n, cannot efficiently cap- ture any shared global mapping ((s())) without repli- cating the same parameters for each token-thii this yields larger approximation error Eg 2 Cg/ng. In contrast, mixed network of width cate n, units to model g( (z))) and n, units to model Σ; u(z;), giving cA I Cg/ nn + Cu/nu. Choosing Th O(n.) ensures A < BB.
n, + ny can dedi-
Complexity Reduction (Lemma B.5). Tying n1s shared units across all T' tokens reduces the effective hypothesis-class size: instead of T n, independent pa- rameters to model g(s())) per token, there are only n, global parameters. This reduces the Rademacher
17

---

Mixture of Hidden-Dimenseee Not All Hidden-States" Dimensions are Needed in Transformer
parameter efficiency and enhances hidden dimension scala- bility.
High Value
Non-Zero Value
Zero
C.2. Sparsely-activated Transformer
X
Sparsely-activated Transformer models, such Sparse Mfxture-of- Expext (MoE) architectures, leverage input adap- tivity to achieve scalable and efficient computation. These models dynamically activate only a subset of specialized subnetworks, or "experts," for processing each input token, significantly reducing computational overhead (Fedus et al., 2022; Riquelme et al., 2021; Zhou et al., 2022; Jiang et al., 2024a; Xue al., 2024b). This mechanism enables ef- fective handling of diverse data domains (Li et al., 2022; Jain et al., 2024) while maintaining high performance. Re- cent advancements in sparsely-activatee activated Transformers have extended their capabilities by introducing heterogeneous experts (Wu et al., 2024; He, 2024), allowing networks to integrate experts with varying capacities and specializations (Dean, 2021: Zhou et al., 2022: Dai et al., 2024). Some recent studies(Qiu et al., 2024) have observed the activa- tion patterns in the intermediate dimensions of FFNs and explored sparsely-activate architectures based on these observations. However, no existing Transformer architec- ture has implemented sparse activation specifically in the hidden dimensions. Inspired by the work of (Qiu et al., 2024; Dai al., 2024; Cai et al., 2024a), we conducted an in-tepth analysis of the hidden dimensions and designed novel sparse activation strategy. This innovation opens new research avenue for sparsely-activated Transformer architectures.
Sparsification (MoE)
X
X
Hidden Sparsification (MoHD)
Intermediate Activation
Input
Sparse Matrix
Sparse Matrix Output
Wup
Wdown SM
X
SM
Xin
y
Figure 9. Using FFN example, the traditional method (MoE) exploits the sparsity of the intermediate dimension. Our method (MoHD) selectively activates only subset of hidden dimension parameters across all matrices to enhance efficiency.
Table S. Theoretical Forward-Pan Pass FLOPS for MoHDD Models Different Configurations (E+ 12+12 denotes x 10
Model Size
MeHlD 50%
MoHlD 75%
Raseline 160%
2x Width
3x Width
Width
355M
2.70E+12
3.63+112
4.58E+12
5.404412
6.24-122
7.07x122
495M L113
4.19E+12
5
7.00xa12
B.40E+12
9,73--12 1,55-113
1.11
6.93E-12
9.80
1.28Ex13
1,40E+13
Connection with Activation Compression: Some works also prune or quantize hidden dimension activations during inference. However, MoHD reduces training costs signif- icantly by activating only a small portion of parameters during training. This provides scalability that these opti- mization methods lack. Since MoHD maintains the same architecture for both training and inference, it avoids perfor- mance degradation. Additionally, MoHD can be combined with these methods for further optimization.
D. Connection With Other Methods
Connection with MoE: MoHD and MoE enhance model efficiency through distinct dimensions. While MoE im- proves memory capacity by sparsely activating experts in the intermediate dimension (Cai et al., 2024c; Dai et al., 2024; Xue et al., 2024a). MoHD optimizes parameter uti- lization by leveraging hidden dimension sparsity, directly strengthening representational capabilities through hidden dimension expansion, as shown in Figure 9, Experiments (in Table 2) demonstrate that MoHD achieves superior parame- ter efficiency under identical activation budgets compared to MoE, though it faces unique challenges like activation collapse and information degradation, necessitating tailored design solutions. The two mechanisms are complemen- tary-MoHD focuses on dynamic feature extraction in the hidden dimension, whereas MoF specialiis in parallel com- putation expansion in the intermediate dimension. Their theoretical compatibility enables synergistic integration, of- fering hybrid architecture that balances representational depth with computational resource optimization.
E. Computational Cost Analysis of MoHD Models
We provides a detailed analysis of the training and inference computational costs for various MoHD model scales, en- compassing both compression and expansion settings. We quantify these costs primarily through theoretical forward- pass Floating Point Operations (FLOPS), comparing them against baseline models, MoE architectures, and other preva- lent sparsification methodologies.
As delineated in Table 5, an increase in model width, while leading to a proportional rise in WTE-related FLOPS, only exerts modest relative impact on total
18

---

Mixture of Guadden-Dimensionee Not All Hidden-States" Dimensions are Needed in Transformer
computational cost. This characteristic is particularly pro- nounced in deeper MoHD models, where the ratio of ac- tivated parameters to total parameters becomes more effi- cient during scaling. We assert that MoHD's 's performance enhancements are not merely a consequence of increased WTE size or a broader parameter count. For instance, the 4x-width model exhibits significantly higher FLOPS than its 3x-width counterpart, yet its performance improvement is marginal. Conversely, the 1.13B model reaps greater perfor- mance benefits from MoHD than the 495M model, despite its FLOOS-to- L.----to-aseeliee ratio being lower. These findings substantiate the claim that MoHD achieves its efficiency predominantly through structural sparsity and expert spe- cialization, rather than through brute-force scaling. Further- more, the computational overhead introduced by MoHD's routing and group fusion layers is minimal.
Sifeer 95%
58%
Jatter wi 50%
weo 7%
(attr.out out
(wwlore dosen 30%
1.00%
1.00%
after WV 38%
Siter pate
S
Figure 10. Visualization of activation magnitude in LLaMA2-7B layer 30. In the Transformer, multiple layers show a consistent pattern of activation flow.
sification techniques due to training-inference mismatches. While these optimization methods can be combined with MoHD for further improvements, MoHD, by focusing on activation sparsity in the hidden dimension, achieves better generalization and stability at inference costs comparable to those of activation sparsification approaches.
In comparing MoHD with MoF. architectures, it is cru- cial to recognize that these methodologies enhance model efficiency along distinct dimensions. MoE typically im- proves memory capacity by sparsely activating experts in the intermediate dimension, whereas MoHD optimizes pa- rameter utilization by leveraging sparsity in the hidden di- mension, thereby directly augmenting representational capa- bilities through hidden dimension expansion. Their routing granularities also differ: MoHD routes at the hidden dimen- sion level, enabling the tailoring of token-specific subspace activations, while MoE routes at the expert (subnetwork) level, primarily within the FFN. Moreover, MoHD's ap- plicability spans both Attention and FFN components, in contrast to MoE's conventional confinement to FFN pro- jections. Empirical evidence indicates that MoHD attains superior parameter efficiency under identical activation bud- gets when compared to MoE. Despite inherent challenges in MoHD's 's routing across hidden dimensions, which ne- cessitate novel optimization and implementation strategies, MoHD demonstrates notable improvements over MoE when trained from scratch. We suggest that MoHD and MoE are complementary, with MoHD focusing dynamic feature extraction in the hidden dimension and MoE on parallel computation expansion in the intermediate dimension, im- plying potential for synergistic integration into more robust and efficient language models.
F. Extended Observations
In this section, we present several key findings that serve as the foundation for the design of the MoHD approach. In Section F.I, we observe the long-tail effect of hidden dimension activation values and define activation sparsity accordingly. We analyze the sparsity distribution and dif- ferences between attention and FFN across different layers. In Section F.2, we analyze activation flow in Transformers, highlighting compression patterns, stabilization by resid- uals and normalization, and functional layer differences. In Section F.3, we further identify the existence of shared continuous high activation behaviors and unique discrete high activation behaviors across tokens. Finally, we analyze these phenomena and propose motivations for designing feasible hidden dimension sparsification methods.
F.1. Sparsity in Tokens' Hidden Dimension
For a more comprehensive understanding, we observe the ac- tivation magnitudes of 4096 hidden dimensions in LLaMA2- 7B. As shown in the left panel of Figure 2, we visualize the relationship between dimension magnitudes and reordered dimension indices based on magnitude size.
When juxtaposed with other sparsification techniques such as pruning and quantization, MoHD presents dis- tinct advantages in terms of training costs and architec- tural consistency. MoHDD significantly reduces training expenses by intrinsically learning sparse activations during the training phase, thereby eliminating the need for itera- tive pruning cycles or post-training fine-tuning procedures characteristic of traditional methods. A key benefit is that MoHD maintains a consistent architecture across both train- ing and inference stages, which inherently mitigates the performance degradation often observed in post-hoc spar-
Similar to previous observations (Liu et al., 2024), the ac- tivation of hidden dimensions exhibits a long-tail sparsity phenomenon. For instance, in the input Attention activations of LLaMA2-7B's 16th layer, the cumulative magnitude of the top 1000 dimensions accounts for 71.96% of the total magnitude. In contrast, most dimensions have low activation values, indicating that the model does not utilize informa- tion from the majority of hidden dimensions, leading to substantial sparsity in activations.
We also visualized the sparsity of activations in the input and output of Attention and FFN components. Our obser-
19

---

Mixture of Hidden- Doaden--Dimensios:: Not All Hidden-States" Dimensions are Needed in Transformer
vations reveal significant differences in the magnitude of hidden activations across positions. Attention exhibits higher activation magnitudes, while FFN activations are comparatively lower. At the input stage, activation mag- nitudes are relatively high (median > 1), whereas at the output stage, activation magnitudes drop significantly (me- dian < 0.5). The sparsity of hidden dimensions in the input components is consistent across different modules, likely due to the influence of residual connections. However, at the output stage, the sparsity patterns of Attention and FFN differ markedly. As shown in the middle panel of Figure 2, Attention demonstrates significant fluctuations in sparsity, with alternating high and low sparsity distributions. In contrast, FFN sparsity remains relatively stable. These differences highlight the distinct functional roles and in- formation processing characteristies of Attention and FFN, prompting to consider differentiated activation designs for these components.
tered and reordered based on their activation patterns. Ap- proximately 400 dimensions commonly highly activated across all 5 tokens, while about 200 dimensions are uniquely highly activated within each token. This indicates that each token's activations contain shared sub-dimensions that are commonly activated and token-specific sub- dimensions that are independently activated. Shared high activations model the similarity information shared across tokens in hidden dimensions, while specialized unique acti- vations capture differences. These observations inspired the shared- -herad-specialized activation mechanism in the subsequent design of MoHD.
G. More Implementation Details
G.I. MoHD Block Implementation
In Sections 2.1 and 2.2, we observed activation differ- ences across components in various layers, prompting us to design separate routing mechanisms for the Attention and FFN components. Specifically, in one Transformer Block, Gate 6, ()) and Gatern (z, 6, N, Φ) produc- ing scores that determine the activation of dimension- specific sub-dimensions for the output:
E2. Activation Flow in Transformer
We also investigate the variations in activation magnitudes within single Transformer block, as illustrated in Figure 10. Consistent activation flow patterns were observed across different Transformer blocks. The Attention module com- presses input activations normalized to 100% through pro- jections (Wq. WK. Wv) and weighted averaging, reducing activation magnitudes 6.7% at Wo. This highlights its ability to suppress irrelevant information through weighted aggregation, while also showcasing significant functional differences between layers, as the output activation mag- nitudes vary to accommodate layer-specific roles. In con- the FFN module demonstrates stable activation pat- terns, with compression arising from high-dimensional pro- jections, nonlinearity that sparsifies activations, and dimen- sionality reduction through linear weighted summation, col- lectively reducing activation magnitudes.
a, Ir Gate (I, 6, NN, (p), m, 2, Gate (I. N. :p).
In practice, different components may employ distinct spar- sification settings. However, for simplicity, we use the same notation throughout this section to represent these settings in a unified manner. Based on the scores from the Router, MoHD applies synchronized sparsification to the hidden dimensions of all up-projection and down-projection ma- trices, as well the input From Equation 4, we trans- form W WK, WV, W Wup Waste, Wdown into MoHD N
's sub-dimensions 1 Q N
M i= K, l1 V O, and GATE,i, II DOWN; We substitute these into the sparsified Attention and FFN defined in Equation 13 and 11, yielding outputs ya and ym. respectively:
II= UP N I1=h,
E3. Continuous High Activation
h
MHAMMHD (#.) M.C.. Σ Head, (1S- 25 O((xa), (14)
We further investigate the temporal correlation of activation sparsity by observing high activation values across different tokens and analyzing the indices that are repeatedly acti- vated by multiple tokens. A clear correlation in activations is observed over consecutive tokens. Figure 2 Right shows the number of commonly highly activated dimensions across 2 to 9 consecutive tokens, with the x-axis representing the threshold for defining high activation. When using the top 20% of activation values as the threshold, 2672 dimensions are commonly activated across 2 consecutive tokens, and 673 dimensions remain commonly activated across 9 con- secutive tokens.
1=1
Head; = a 20 (115", (0(a)))
x.
(1150 ag K, (2x1)
1
vith
x. (1/= V, Vy(
(15)
FFNMOHD (xx) = Mm.mm
mi (x.)
x Ix
mc UP.
Figure further illustrates the correlated activation patterns over 5 tokens, where the 4096 hidden dimensions are clus-
my GATE ) (16)
20

---

Mixture of Tldden-Dimensiinn Not All Hidden-States" Dimensions are Needed in Transformer
fusion, which suppresses noise from redundant activations and enhances performance.
Table Parameter configurations of MoF. under compression and expansion experiments. We use the same settings for only FFN.
FFN Exhibits Greater Redundancy. The FFN layer shows more redundancy in hidden dimensions, causing minimal performance loss (and sometimes improvement) when spar- sified. In contrast, sparsifying Attention results in a more substantial performance drop. For the 50% sparsity setting, the FFN uses only 195M activated parameters compared 239M for Attention. This suggests the FFN is better suited for MoHD transformation. The FFN achieved a -0.30 PPL reduction at 50% sparsity, likely due to a decrease in redun- dant activations that reduces overfitting, whereas Attention sparsification led to a + +1.04 PPL increase.
MoH
50%
75%
x2
x3
shared experts
1
1
1
mated experts
7
7
11
16
num experts per tok
3
3
expert interdemiate vatio
12.5%
12.5%
25%
25%
25%
Transformer
14.0
Only MoHD FFN
13.5
Only MoHD ATTN
Eval PPL
13.0
MoHD ATTN MLP
12.5
Combining MOHD Attention and FFN Leads to En- hanced Performance. Joint sparsification of Attention and FFN yields the best parameter efficiency. The 50%ATTN- 50%FFN model achieved a PPI. of 12.05 with only 145M activated parameters-lower than both 50%FFN and 50%ATTN configurations. This joint sparsification outper- forms the 25% FFN configuration by 0.19 PPL. because enhances the consistency of activated hidden dimensions, preserving better learning capacity,
12.0
11.5
100
75
50
25
Activated Parameters Ratio(%)
Figure 11. The model's Eval PPL under different sparsity settings applied to Attention and FFN components at varying ratios.
We construct MOHD BLOCK based on MoHD spec- ified MHA and FFN components. Residual connections are designed to further mitigate information loss during the specified forward pass. Following the configuration of LLAMA, we apply LayerNorm layers before the input to both MHA and FFN; however, for simplicity, these are omit- ted in the formal equations. This process can be formalized as follows:
I. Discussion
L.. Connection with Mixture of Experts Methods
Mixture of Hidden-Dimensionn (MoHD) and Mixture of Experts (MoE) are both sparse activation architectures de- signed to enhance Transformer efficiency, but they operate on fundamentally different dimensions of the model archi- tecture. MOHD focuses on improving the efficiency of scaling the hidden dimension, which represents the embed- ding size of tokens and is related to the model's ' width or representational capacity per token. In contrast, MoE targets the efficiency of scaling the intermediate dimension within the Feed-Forward Networks (FFNs), effectively increasing the model's depth or memory capacity by activating a subset of specialized "experts" subnetworks.
BLOCKMOHD (x) FFNMoHD (MHAMOHD(*)++- +x) +xx)). +r. (17)
To train the model effectively, we combine cross-entropy loss LcB for language model pre-training and the load bal-
ance loss LB, resulting in the final training objective:
L == Lcr + La-
(18)
H. Extended Experiment Details
MoHD achieves efficiency by observing and exploiting the sparsity within the hidden dimension activations. It posits that not all hidden dimensions are equally important for every token, identifying both shared sub-dimensions ac- tive across multiple tokens and specialized sub-dimensions unique individual tokens. By introducing a routing mech- anism that dynamically selects and activates only subset of these hidden sub-dimenn aub-dimensions, MoHD reduces the number of activated parameters across all relevant matrices (including those in Attention and FFN layers). This approach allows for the expansion of the total hidden dimension, thereby enhancing the representational power of each token, without a proportional increase computational cost.
H.1. Decoupled MOHD Components Setting
To investigate MoHD sparsification, we built three models based on L.aMM2-355M: one with sparsification applied to Attention, one to the FFN, and one to both, as shown in Figure 11. All models were trained on 10B tokens. # Acti- vate refers to activated parameters with the total parameter remaining constan. Table 4 presents the effects of decoupled MoHD at 100%, 50%, and 25% sparsity.
MoHD Architecture Advantages. Even with 100% spar- sity (where activated parameters match the original model), MoHD outperformed the baseline. This is likely due to its optimized activation dimension allocation and grouped
21

---

Mixture of Hiidden-Dimensionn Not All Hidden-States" Dimensions are Needed in Transformer
Conversely, MoE architectures enhance efficiency by intro- ducing multiple expert networks in the intermediate dimen- sion of FFN layers. For each token, gating mechanism selects and activates only sparse subset of these experts. This effectively expands the total parameter count and mem- ory capacity available to the model by adding more paral- lel subnetworks, but the activation remains sparse across these experts, reducing the computational cost compared to activating all experts. MoE's primary mechanism is token- conditional routing to different intermediate pathways, typi- cally confined to the FFN module, rather than sparsifying the core hidden dimension across all components.
sions are consistently highly activated across sequence of tokens, termed shared sub-dimensions, while others are uniquely activated for individual tokens, referred to as specialized sub-dimensions. This distinction suggests a functional divergence within the hidden space: shared dimensions likely capture common features and representa- tional similarities that are relevant across different tokens, facilitating the processing of general linguistic patterns. In contrast, specialized dimensions appear to be crucial for encoding fine-grained, token-specific semantic differences and higher-level contextual information. This empirical finding directly inspires the structured routing approach in MoHD, which decouples the handling of these two types of dimensions. By partitioning the hidden space into shared and specialized subspaces and employing a routing mech- anism that guarantees consistent activation for the shared portion while dynamically selecting from the specialized portion based on token input, MoHD aims to effectively model both the commonalities and unique characteristics of tokens. This mixed activation strategy, balancing univer- sally relevant features with token-specific details, is central to MoHD's ability to enhance representational capacity and parameter efficiency simultaneously.
Thus, the core distinction lies in their dimensionality of fo- cus: MoHD addresses the efficiency of the model's width by sparsely activating hidden dimensions improve token representation, while MoE addresses the efficiency of the model's depth or intermediate capacity by sparsely activat- ing experts to enhance computational and memory scalabil- ity. Although they operate on different dimensions, these two approaches are complementary and can be considered for synergistic integration to potentially leverage sparsity in both hidden and intermediate spaces.
1.2. Design Principles of Sub-Dimension Routing
I.3, Activation Flow Maintenance and Load Balancing
The design of the Mixture of Hidden Dimensions (MoHD) architecture. particularly its sub-dimension routing mech- anism, is fundamentally driven by empirical observations regarding the activation patterns within Transformer hidden states. Analysis of trained Transformers reveals signifi- cant sparsity in the hidden dimensions, where a substantial portion of dimensions exhibit low activation values and con- tribute minimally to the overall activation magnitude. This long-tail distribution of activation magnitudes indicates that the model does not fully leverage all available dimensions uniformly across all tokens. Consequently, the traditional Transformer approach of treating all hidden dimensions equally results in considerable computational and memory overhead, especially as models scale. This observed sparsity motivates the necessity of a dynamic mechanism that can selectively activate only a relevant subset of hidden dimen- sions for each input token, thereby improving efficiency by concentrating computational resources on the most infor- mative subspaces rather than processing the entire, often re- dundant, dimensionality. The dynamic routing mechanism in MoHD serves this purpose by allowing the model to adaptively determine which dimensions are activated based on the specific characteristics of the input, avoiding the com- putational burden associated with full hidden dimension activation.
The design of Activation Flow Maintenance in MoHD motivated by empirical observations of consistent activation flow patterns within Transformer blocks. Specifically, it was noted that Attention modules compress input activations, while FFNs maintain stable activation patterns. However, applying sparse activation to hidden dimensions, as done in MoHD, can lead to information degradation. This degrada- tion arises from the router's 's softmax- normalized weights, which may cause some dimensions to receive disproportion- ately high weights while others neglected. Furthermore, the parallel concatenation of activated sub-dimension out- puts can suppress information in low-weighted dimensions without compensation. To counteract these issues and main- tain robust activation flow despite sparse activation, MoHD employs several strategies.
One key strategy is Sub-dimension Scaling. This mecha- nism addresses the suppression of sub-dimension activa- tions caused by softmax weight normalization. A scaling factor, a. is introduced to ensure that the sum of activation weights across all dimensions remains consistent with the input magnitude, allowing activated dimensions to retain their proportional influence and preserving stable activation magnitudes across dimensions.
Additionally, the Grouped Fusion Layer is introduced to mitigate information loss from the sparse output. This layer projects the sparsified to the original dimension d. To reduce computational over-
Further investigation into the highly activated dimensions across multiple tokens reveals a more nuanced structure than simple universal sparsity. We observe that some dimen-
output, ye. back
22

---

Mixture of
Not All Hidden-States" Dimensions are Needed in Transformer
head, a Monarch matrix is used for efficient grouped fusion mapping. Given a receptive field r, this mapping matrix M is structured to perform grouped filling and mapping after sparse activation. This process reconstructs information across the original hidden dimensions while maintaining computational efficiency, preserving information integrity, and improving dimension utilization without significantly increasing parameters.
in the routing can lead to skewed distributions, potentially suppressing useful but low-weighted sub-dimensions and reducing representational fidelity. Addressing these lim- itations is vital for the broader applicability and further enhancement of MoHD.
Exploring the application of MoHD to larger-scale Large Language Models (LLMs) presents significant potential despite associated challenges. Larger models typically ex- hibit higher parameter redundancy, making MoHD's 's sparse activation mechanism potentially more effective in reducing inefficiencies without sacrificing expressive power. Experi- ments on smaller scales indicate that MoHD's 's performance improvement over baselines increases with total parame- ter count, suggesting that the benefits may scale positively with model size. By combining shared and specialized sub- dimensions, MoHD enhances the model's capacity to cap- ture both general and fine-grained, token-specific patterns. Scaling up could increase the size of each uub-dimension, improving its ability to model complex language phenomena and potentially enhancing generalization across tasks. How- ever, pretraining larger L.LMs is computationally expensive, hindering immediate full-scale experiments. Future work should aim to validate these potential benefits and navigate the associated training complexities.
Beyond maintaining the activation flow within the com- ponents, it is crucial to ensure that the learned routing strategy effectively utilizes all available specialized sub- dimensions. Research on conditional computation architec- tures has shown that automatically learned routing can suffer from load imbalance, where only a few sub-dimensions are frequently selected, leaving others underutilized. To ad- dress this and encourage more even distribution of tokens among different sub-dimensions, the jub-Dimension Load Balance Loss is incorporated. This loss penalizes imbal- ances in sub-dimension assignments. By encouraging the gating mechanism to distribute assignments more uniformly across sub-dimensions, the balance loss helps prevent router collapse, where the router consistently selects only small subset of sub-dimensions, and ultimately leads to improved utilization and overall model efficiency.
The core principles of MoHD may also extend beyond Natural Language Processing (NIP) tasks, particularly to other domains that utilize Transformer architectures, such as Vision-Language Models (VLMs). Vision Transformers (ViTs), for instance, process image patches as tokens, simi- lar text tokens. These visual patches exhibit both global patterns shared across the image and local unique features, analogous to linguistic semantics. If the hidden dimensions in ViT's also demonstrate shared and token-specific acti- vation patterns, the principle of selectively activating sub- dimensions could enhance efficiency and expressiveness in vision tasks. However, such application would require architectural adaptation visual representations differ significantly from language. Tuning routing and partitioning strategies to align with visual characteristics is necessary. Furthermore, in multimodal settings like VLMs, incorporat- ing sparsity-aware mechanisms is crucial to avoid harming the alignment across modalities. Direct transfer is not triv- ial, but the generalizability of MoHD's principles warrants exploration in vision and multimodal settings.
1.4. Limitations and Future Research Directions
Despite its demonstrated effectiveness in improving param- eter efficiency and task performance, the Mixture of Hid- den Dimensions (MoHD) architecture has several existing limitations that present opportunities for future research. A primary challenge lies in the sensitivity of MoHD to hyperparameters, specifically the sparsity ratio 8, the pro- portion of shared sub-dimensions yp, and the total number of sub-dimensions N. Achieving the optimal balance between shared and specialized sub-dimensions is crucial for maxi- mizing generalization and specialization, requiring careful tuning. Furthermore, optimizing the routing mechanism poses a significant challenge. Unlike the routing in Mixture of Experts (MoE) models which activate experts in the inter- mediate dimension, MoHD activates sub-dimensions within single matrix, which inherently more complex. While effective, the current routing loss stability decreases as the number of sub-dimensions increases, potentially limiting the scalability of the routing design itself. Another limitation arises from the growth of the Word Token Embedding (WTE) layer. As MoHD expands the effective hidden dimension (model width), the WTE layer scales proportion- ally, contributing a neglgggl portion of total param- eters and potentially offsetting some efficiency gains, par- ticularly at very large scales. Finally, despite the inclusion of activation scaling and group fusion mechanisms, there is a risk of information degradation under extreme spar- sity. Large-scale downsampling and the softmax weighting
Future research directions for MoHD should focus on en- hancing the robustness and applicability of the architec- ture. Improving the stability and scalability of the rout- ing mechanism is a critical area, particularly as the num- ber of sub-dimensions increases. Investigating alternative routing strategies or loss functions could help address load imbalance issues and prevent router collapse, ensuring all sub- lub-dimensions are actively utilized. Further analysis and
23

---

Mixture of Hiidden-Dimennioos:: Not All Hidden-States" Dimensions are Needed in Transformer
Table Detailed configuration, activation parameters, and total pa- rameters of the models included in our study. L.2-355M represents the LaMMA-2 architecture model with 355M total parameters.
Model Setting
L.2-355M
1.2-495M
1.2-1.13B
hidden size
1024
1536
2048
internediate size
2560
2560
4096
attention heads
32
32
32
num kv heads loyers
32
16
32
24
24
24
# Activate
289M
396M
1B
# Params
355M
495M
1.13B
Table Parameter configurations of MoHD under compression and expansion experiments. We use the same settings for both Attention and FFN, For detailed reasoning behind these configura- tions, please refer to Analysis.
MoH
50%
75%
x2
x3
X
4
atm top
4
4
4
4
atn sub-dim num
8
12
8
12
16
ffn top k
4
4
4
4
4
sub-dim num
8
12
8
12
16
shared sub-dim num
group fusion dim
12
16
optimization of the WTE layer's growth relative to total computational cost are needed to maximize efficiency gains at very large scales. Exploring the potential for synergistic integration with other sparsity methods, such Mix- ture Experts (MoE), could unlock sparsity benefits across multiple dimensions, combining MoHD's focus hidden dimension dynamism with MoE's intermediate dimension expansion. Additionally, conducting more extensive abla- tion studies on the interaction between shared and special- ized sub-dimensions and analyzing the impact of different routing parameters (e.g., K-value selection) would provide deeper insights into MoHD's mechanisms and guide fur- ther design improvements. Finally, empirical validation larger-scale models and diverse non-NLP tasks is essen- tial to demonstrate the broader potential and identify specific challenges in different domains.
24


Reviews:
[{"review_id": "WsH7MIr1yr", "reviewer": "Reviewer_fdJw", "summary": "This paper proposes MoHD (Mixture of Hidden Dimensions), an architecture that optimizes hidden dimension usage via dynamic routing between shared and token-specific dimensions. Experiments demonstrate its superior parameter efficiency and task performance over existing models.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the valuable comments and suggestions.\n\n**Q1: On MoHD-495M’s performance on WinoGrande (WG)**\n\n**R1:** We respectfully clarify that MoHD-495M does **not consistently underperform** on WG. In fact, our MoHD 50%-495M model achieves **52.7%**, outperforming the **LLaMA2-495M baseline (51.3%)**. At larger scales (e.g., 1.13B), MoHD shows **clear gains** under multiple configurations (75% x2, x3, x4), indicating strong adaptability.\n\nThat said, performance dips in certain settings may stem from:\n\n- **WG-specific reasoning demands**: WG requires fine-grained commonsense inference, which may benefit from further adaptation of MoHD’s routing and activation strategies.\n- **Routing sensitivity**: Some configurations may have suboptimal routing for WG due to overfitting to other tasks. Ensuring robustness across specialized reasoning datasets like WG is a valuable direction for future work.\n\n**Q2: On integrating MoE into MoHD**\n\n**R2:** We appreciate this suggestion. **Yes, integration is theoretically feasible**. MoHD introduces sparsity along the **hidden dimension (width)**, while MoE sparsifies the **intermediate dimension (length)**. These operate orthogonally and, in principle, could be combined for **multi-dimensional sparsity**, potentially improving parameter efficiency.\n\nHowever, such integration would introduce **significant optimization and engineering complexity**. Further empirical studies are needed to assess whether their combination yields **synergistic or conflicting effects**.\n\n**Q3: MoHD vs. MoE**\n\n**R3:** Key differences:\n\n- **Routing granularity**: MoHD routes at the **hidden dimension** level, tailoring token-specific subspace activations. MoE routes at the **expert** (subnetwork) level in the FFN.\n- **Component scope**: MoHD applies to **both Attention and FFN**, while MoE is typically confined to FFN projections.\n- **Capacity and interpretability**: MoHD expands **model width**, increasing per-token representation capacity. MoE expands the **FFN depth**, aiding memory but not width.\n- **Efficiency**: MoHD reduces redundant activations in **both Attention and FFN**, offering **better scaling under fixed activation budgets** (see Table 2).\n- **Challenges**: MoHD’s routing across hidden dimensions is **less explored** and demands **novel optimization and implementation**.\n\nDespite these challenges, MoHD shows **notable improvements** over MoE when trained from scratch, highlighting its **promise for efficient scaling**.\n\n**Q4: Theoretical Proof**\n\n**R4:** We prove that mixed sparse activation achieves strictly better risk bounds. \n\n**Lemma 1 (Unbiased Sparse Forward Pass).**  \nLet $h \\in \\mathbb{R}^n$ be a hidden layer. Apply a mask $r \\in \\{0,1\\}^n$ with $\\mathbb{P}(r_j=1) = p$, and define the sparsely activated output as $\\hat{h} = \\frac{1}{p}(r \\odot h)$. Then:  \n$$ \\mathbb{E}[\\hat{h}] = h; \\quad \\mathbb{E}[||\\hat{h}||^2] = ||h||^2/p + (1-1/p) \\sum_{j=1}^n h_j^2 $$\n\n*Proof:* Linearity of expectation and variance decomposition.  \n\n**Corollary 1.1 (Effective Width).**  \nFor $p = k/n$, training with sparse activation is equivalent (in expectation) to training a full network with width $n' = \\frac{n}{k} \\cdot \\mathbb{E}[\\|h\\|_2^2] / \\|h\\|_2^2 \\geq n$. Thus, sparse activation expands effective width.\n\n**Lemma 2 (Approximation Error Decay with Width [Barron, 1993]).**  \nLet $f^* \\in \\mathcal{F}$, where $\\mathcal{F}$ is a Barron space. For a network with width $n$, the approximation error satisfies:  \n$$\n\\epsilon(n) \\leq \\frac{C_f}{\\sqrt{n}},\n$$  \nwhere $C_f$ is the Barron norm of $f^*$.  \n\n**Corollary 2.1 (Mixed Activation Lowers Error).**  \nDefine $s$ encodes global context, $g$ decodes shared features, $u$ models per-token contributions; $\\Theta$ denotes all parameters.\n\nIf $f^* = g(s(x)) + \\sum_i u(x_i)$ with non-constant $g$, then:  \n- Token-only networks ($n = n_u$) have $\\epsilon_B \\geq \\frac{C_g}{\\sqrt{n_u}}$.  \n- Mixed networks ($n = n_s + n_u$) have $\\epsilon_A \\leq \\frac{C_g}{\\sqrt{n_s}} + \\frac{C_u}{\\sqrt{n_u}}$.  \nChoosing $n_s = \\Theta(n_u)$ yields $\\epsilon_A < \\epsilon_B - \\delta$ for $\\delta = \\frac{C_g}{2\\sqrt{n_s}}$.\n\n**Lemma 3 (Rademacher Complexity of Shared Dimensions).**  \nLet $\\mathcal{H}_A$ (mixed) and $\\mathcal{H}_B$ (token-only) have equal total parameters. Then:  \n$$\n\\text{Rad}(\\mathcal{H}_A) \\leq \\text{Rad}(\\mathcal{H}_B) - \\Delta,\n$$  \nwhere $\\Delta = \\Omega\\left(\\sqrt{\\frac{n_s}{m}}\\right)$ for $m$ samples.  \n\n*Proof:* Shared Dimensions reducing the VC dimension [Bartlett, 1998]. Apply Dudley entropy integral.\n\n**Theorem (Risk Bound).**  \nFor risk $R(f) = \\epsilon(f) + C\\cdot\\text{Rad}(\\mathcal{H})$:  \n\n$$ R_A \\leq \\epsilon_A^{\\text{(Lower error)}} + C\\cdot\\left(\\text{Rad}(\\mathcal{H}_B) - \\Delta\\right)^{\\text{(Lower complexity)}} < R_B $$\n\nSparse activation (i) expands effective width, which (ii) lowers approximation error when shared dimensions are included, while (iii) shared dimensions reduce Rademacher complexity - collectively proving $R_A < R_B$.", "claims_and_evidence": "The paper provides empirical evidence to support claims. \nThe authors empirically show sparse activation patterns in hidden dimensions across tokens: some shared, and others token-specific. Extensive experiments validate the performance gains with low computational overhead. Ablation studies confirm component efficacy.", "methods_and_evaluation_criteria": "I think the method design is mostly reasonable. Though hidden dimension sparsity is known, cross-token activation modeling is novel and justifies shared sub-dimensions. Evaluations use standard NLP datasets and metrics, aligning with community norms.", "theoretical_claims": "I think the theoretical contributions are mostly based on the observations of activation patterns and routing strategies. While the motivations are intuitive, further analyses can be enhanced.", "experimental_designs_or_analyses": "I think the experiments are mostly convincing.", "relation_to_broader_scientific_literature": "NA", "essential_references_not_discussed": "NA. Most related works have been cited.", "other_strengths_and_weaknesses": "I think this paper has strengths as:\n1. Novel Approach.  \n2. Clear writing.  \n3. Robust generalization across NLP tasks  \n4. Persuasive ablation studies.  \nFor the shortcomings of this paper, please refer to Questions.", "other_comments_or_suggestions": "There are some terminology and symbols requiring more clarification:\n1. Terminology clarification: Does *dimension* refer to embedding size, specific dimensions, or all dimensions?  \n2. Symbols clarification, e.g., $\\otimes$ in Equation 3.", "questions_for_authors": "2. Why does the 495M MoHD underperform on WinoGrande (WG)?  \n3. Is it possible to apply the existing MoE model and Dense model to the MoHD architecture?\n4. Could you describe in more detail the differences between the MoHD and MoE approaches in terms of routing optimization, performance efficiency, etc.? Can these two approaches be combined?\n5. Can you provide some theoretical basis for MoHD?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "MzsYdrFVi7", "reviewer": "Reviewer_h4Rw", "summary": "This paper proposes a LLM sparsification method, namely Mixture of Hidden Dimension (MOHD), to improve the efficiency of Transformer-based LLMs. Based on the observation that only a small subset of hidden dimensions is shared and activated across tokens in given texts, MOHD selectively discern shared and specific sub-dimensions. In this way, not all hidden dimensions are utilized, and thereby improves the parameter efficiency, and utmost retain competitive performance with original LLMs. Experimental results demonstrate the effectiveness of the proposed method.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the valuable comments and suggestions.\n\n**Q1:  Comparison with MoE and other sparsification methods**\n\n**R1:** We respectfully refer to our responses to Reviewer fdJw (Q2, Q3) for a detailed comparison between MoHD and MoE, especially regarding routing design and efficiency. Comparisons with other sparsity-based methods are addressed in Reviewer 3ShX, Q1.\n\n**Q2: The illustration and motivation of routing mechanism and activation flow**\n\n**R2:** Thank you and we will further improve the writing. A detailed explanation is provided below:\n\n- The motivation for introducing **Sub-dimension Scaling** comes from our analysis of activation patterns in Transformers (Section 2.2, Figure 10). We observed that after sparse activation, activation magnitudes tend to decrease, potentially causing **information loss**. To address this, we introduced a scaling factor \\( \\alpha \\), which adjusts the outputs of activated sub-dimensions to restore the original magnitude. This factor is dynamically computed based on the number of activated sub-dimensions, helping to **preserve consistent activation flow**.\n\n- Our **Dynamic Routing** mechanism is motivated by the observation that many hidden dimensions have consistently low activation values, indicating **redundancy**. To reduce this, we designed a routing method that dynamically selects sub-dimensions based on input token characteristics, allowing the model to adaptively adjust its activations and enhance representational capacity. We also observed **shared activation patterns** across tokens and introduced **shared sub-dimensions** to capture these common features, reducing the complexity of routing and improving training efficiency. Experiments confirm that this approach yields **notable improvements**.\n\n\n**Q3: MoHD's applicability to larger LLMs**\n\n**R3:** Thank you for the suggestion. We believe applying MoHD to larger-scale language models offers several potential benefits:\n\n* **Improved parameter efficiency**: Larger models typically suffer from higher parameter redundancy. This makes MoHD’s sparse activation mechanism more effective, as it helps reduce inefficiencies without sacrificing expressive power.\n\n* **Accelerated performance gains with scale**: As shown in **Table 1**, under the same proportion of activated parameters, the performance improvement of MoHD over the baseline tends to grow as the total parameter count increases. This suggests that the advantages of MoHD may scale positively with model size.\n\n* **Enhanced representation and generalization capacity**: By combining shared and specialized sub-dimensions, MoHD captures both general features and fine-grained, token-specific patterns. Scaling up the model increases the size of each sub-dimension, potentially improving its ability to model complex language phenomena and enhancing generalization across tasks.\n\n\n\nDue to the significant time and computational costs associated with pretraining larger LLMs, we were unable to include full-scale experiments in the current version. However, we plan to include these extended results in a future revision to address your suggestion more thoroughly.\n\n\n\n**Q4: The computational costs at different scales can be further clarified**\n\n**R4:** Thank you for this valuable suggestion. To provide a clearer view of the computational overhead, the following table presents the theoretical forward-pass FLOPS for models equipped with Word Token Embeddings (WTE), routing, and group fusion layers under different configurations:\n\n|Model Size|MoHD 50%|MoHD 75%|Baseline 100%|2× Width|3× Width|4× Width|\n|-|-|-|-|-|-|-|\n|355M|2.70E+12|3.63E+12|4.56E+12|5.40E+12|6.24E+12|7.07E+12|\n|495M|4.19E+12|5.59E+12|7.00E+12|8.40E+12|9.73E+12|1.11E+13|\n|1.13B|6.93E+12|9.80E+12|1.26E+13|1.40E+13|1.55E+13|1.69E+13|\n\nAs shown above, although increasing model width leads to a proportional increase in WTE-related FLOPS, the relative impact on total FLOPS remains modest. This behavior is partly attributed to the architecture design: deeper models with a higher depth-to-width ratio benefit more from MoHD, as the ratio of activated parameters to total parameters becomes more efficient during scaling.\n\nIt’s also important to note that the performance gains of MoHD are **not merely a result of increased WTE size or parameter count**. For instance:\n\n* The FLOPS of the **x4-width model** is significantly higher than that of **x3**, yet the performance improvement is marginal.\n\n* The **1.13B model** benefits more from MoHD in terms of performance improvement compared to the 495M model, even though its **FLOPS-to-baseline ratio is actually lower**.\n\nThese observations reinforce our argument that MoHD achieves efficiency primarily through **structural sparsity and expert specialization**, not brute-force scaling.\n\n**Q5:**  The clarification of symbols.\n\n**R5:**  We apologize for the confusion and will provide clearer and more readable content in the next revision.", "claims_and_evidence": "The experiments are basically sufficient. The model is established based on the observation, which provides empirical evidence for the motivation. To achieve the sparsification, the paper further proposes a routing mechanism. Experimental results indicate the effectiveness of the proposed method.", "methods_and_evaluation_criteria": "The paper conduct experiments on benchmark datasets and tasks on the widely-used LLaMA architecture. The method evaluation is convincible to me.", "theoretical_claims": "The model is established based on the observation that not all hidden dimensions are utilized in LLMs for tasks. Though the underlying designations of the specific modules can be further illustrated, most parts are reasonable. Experimental results evaluate the method designs.", "experimental_designs_or_analyses": "The paper conducts comparison experiments, ablation studies, parameter analysis. The appendix also provides details of experimental design and analysis.", "relation_to_broader_scientific_literature": "The proposed method is related to general efficiency studies on LLMs. The research topic is important, and this paper may draw some interests in the community.", "essential_references_not_discussed": "Most are cited. More discussion and connection to existing MoE studies may be helpful in understanding the further contribution between them.", "other_strengths_and_weaknesses": "Strengths:\n1. The motivation is mostly clear. The preliminary experiments and observations are inspirable.\n2. The proposed method is somewhat interesting and mostly reasonable.\n3. The experiments demonstrate the effectiveness of the proposed method.\n\nWeaknesses:  \n1. Some details of the specific module are not clear. The routing mechanism and activation flow can be further illustrated. The specific motivation of modules can be enhanced.\n2. The computational costs at different scales can be further clarified.\n3. It is not clear whether the proposed method can be applicable to larger LLMs.\n4. The writing can be further improved. The clarification of symbols can be detailed for better understanding.", "other_comments_or_suggestions": "The paper is well-written, but technical explanations of the routing mechanism would benefit from further clarification. A more explicit comparison with MoE and other sparsification methods in terms of scalability would also enhance the work.", "questions_for_authors": "What are the computational costs of training and inference with MOHD compared to MoE and other sparsification methods?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {"ethical_review_concerns": "No."}}, {"review_id": "oz5gprRaQi", "reviewer": "Reviewer_3ShX", "summary": "-\tThe proposed Mixed Hidden Dimensions (MOHD) aims to address the inefficiency of hidden dimension scaling.\n-\tThe core insight lies in the observation that only subsets of dimensions are activated across tokens, with some dimensions shared globally and others allocated as \"private\" dimensions. \n-\tThe MOHD model reportedly achieves comparable or superior performance to standard Transformers while reducing activation parameters by up to 50%.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are greatly encouraged by the reviewer’s positive feedback. Below, we address each of your comments in detail.\n\n**Q1: Theoretical Grounding**\n\n**R1:**  Please see our response to **Reviewer fdJw (Q4)** for a more detailed explanation. In summary, we theoretically ground MoHD by showing that:\n\n1. Sparse activation leads to **effective width expansion**;\n2. This expanded width results in **lower empirical loss**;\n3. A **hybrid scheme** of private and shared sub-dimensions yields better performance than either alone.\n\n**Q2: Computational Costs Compared to MoE and Other Methods**\n\n**R2:**  We summarize MoHD’s computational characteristics relative to other approaches:\n\n- **Compared to MoE**: MoHD sparsifies the hidden dimension across **all matrices**, while MoE targets only the FFN. Thus, MoHD achieves higher sparsity coverage. It does, however, introduce modest overhead from additional WTE parameters and routing operations, particularly at inference.\n\n- **Compared to pruning and quantization**: MoHD avoids iterative pruning or post-training fine-tuning. It learns sparse activations **during training**, reducing overall training costs while maintaining performance.\n\n- **Compared to activation sparsification**: Unlike post-hoc sparsification, which may degrade performance due to training-inference mismatch, MoHD maintains consistency between phases. This leads to **better generalization and stability** at a similar inference cost.\n\n**Q3: Scaling Costs**\n\n**R3:** Please refer to **Reviewer h4Rw, Q4** for detailed cost analysis across model scales. In brief:\n\n- **Routing and fusion layers** add minimal parameter overhead.\n- The primary computational increase comes from the **WTE layer**, which scales with model width.\n- Notably, the performance benefits of MoHD **cannot be solely attributed to increased WTE size**, as demonstrated in our ablation and FLOPS analysis.\n\n**Q4: How MoHD Improves Performance While Reducing Activations**\n\n**R4:**\n\n1. **Sparsity and Redundancy in Hidden Dimensions**  \nEmpirical studies reveal that only a small subset of hidden dimensions are meaningfully activated. MoHD removes redundant activations, concentrating computation on **informative subspaces** and improving efficiency.\n\n2. **Component-wise Sensitivity to Sparsification**  \nConsistent with prior work, our experiments show that the FFN layers tolerate, and sometimes benefit from, sparsification due to regularization effects. In contrast, MHA layers are more sensitive. MoHD accommodates this by applying **structured sparsity selectively**, preserving performance.\n\n3. **Grouped Fusion Mechanism**  \nMoHD achieves efficiency via:\n\n- **Activation scaling**: Dynamically adjusts magnitudes of active units for representational stability.\n- **Grouped fusion**: Aggregates sparse activations to preserve expressive capacity.\n\nThese mechanisms enable MoHD to reduce activations without compromising—and often **improving**—performance. Our findings suggest that **well-structured sparsity** improves generalization by filtering noise and reinforcing salient patterns.\n\n**Q5: Limitations of MoHD**\n\n**R5:**  We appreciate the chance to elaborate on current limitations:\n\n- **Hyperparameter Sensitivity**  \nMoHD relies on key hyperparameters: sparsity ratio  $\\delta$, shared dimension proportion $\\phi$, and total number of sub-dimensions $N$. Tuning the balance between **shared and specialized** sub-dimensions is especially important for optimizing generalization vs. specialization.\n\n- **Routing Optimization Challenges**  \nUnlike MoE’s routing, MoHD activates sub-dimensions **within a single matrix**, making routing more complex. While effective, our current routing loss becomes less stable as the number of sub-dimensions increases, **limiting scalability**.\n\n- **WTE Layer Growth**  \nAs MoHD expands the effective width, the **WTE layer scales proportionally**, contributing to a non-negligible share of total parameters. This can offset efficiency gains at very large scales.\n\n- **Information Degradation at Scale**  \nDespite using activation scaling and group fusion, large-scale downsampling and softmax weighting may still cause **skewed distributions**, suppressing useful but low-weighted sub-dimensions. This can reduce representational fidelity, especially under extreme sparsity.\n\n**Q6: Extend Reference**\n\n**R6:**  We sincerely thank the reviewer for the additional references. We will include a detailed discussion of them in the next revision.", "claims_and_evidence": "-\tExperiments supports the claims for smaller models (e.g., 355M/495M parameters). \n-\tResults are largely convincing, though extrapolation to larger models (e.g., 1.13B) raises questions.", "methods_and_evaluation_criteria": "-\tMOHD is evaluated across diverse tasks.", "theoretical_claims": "The theoretical foundation requires further elaboration. \n\n-\tWhile the shared/specialized sub-dimension concept is appealing, rigorous analysis of their interactions is lacking. \n-\tThe activation scaling mechanism, though practical, lacks theoretical justification for mitigating information loss. Deeper mathematical insights would ground the method more firmly.", "experimental_designs_or_analyses": "-\tAblations offer useful insights but could better explain design choices (e.g., group fusion, balance loss) and their specific contributions. \n-\tClaims of MOHD superiority over MoE feel underdeveloped. A nuanced comparison of trade-offs (e.g., model scale, task complexity) would provide a more balanced perspective.", "relation_to_broader_scientific_literature": "-\tMOHD contributes to LLM Architecture, and Efficient-LLM", "essential_references_not_discussed": "Existing works also observe the sparsity of Transformers [1,2], and this paper is also related to the pruning methods of LLMs [3,4], which are recommended for citation and discussion.\n\n[1] Li et al. The Lazy Neuron Phenomenon: On Emergence of Activation Sparsity in Transformers. ICLR 2023. \n\n[2] Wang et al. Q-Sparse: All Large Language Models can be Fully Sparsely-Activated. CoRR abs/2407.10969 (2024).\n\n[3] Ma et al. LLM-Pruner: On the Structural Pruning of Large Language Models. NeurIPS 2023.\n\n[4] Dong et al. Pruner-Zero: Evolving Symbolic Pruning Metric From Scratch for Large Language Models. ICML 2024.", "other_strengths_and_weaknesses": "Pros:\n- The writing and organization are relatively clear.  \n- A unique solution to inefficient hidden dimension scaling in Transformers.  \n- Convincing results for smaller models.  \n- The shared/specialized sub-dimension concept is novel and promising.  \nCons:\n- Insufficient theoretical grounding.  \n- Limited comparisons with other sparse/MoE architectures.  \n- Computational costs for scaling MOHD need clearer articulation.", "other_comments_or_suggestions": "This paper would benefit from deeper analysis of limitations.", "questions_for_authors": "- How does MOHD compare to other sparse models (e.g., structured/activation pruning) in training efficiency and convergence?  \n- What is the theoretical basis for activation scaling? Why does it preserve activation flow despite reduced compute?  \n- How does MoHD improve performance while reducing activations?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "JzsZiibjY0", "reviewer": "Reviewer_GY6F", "summary": "This paper presents a novel Transformer architecture to address the challenges associated with scaling hidden dimensions. The proposed MoHD (Mixture of Hidden Dimensions), leverages the observation of high hidden dimension sparsity to enhance computational efficiency and model performance. In experiments, MoHD outperforms vanilla Transformers and MoE models on efficiency and performance. Several techniques such as shared sub-dimension are introduced to dynamically activate sub-dimensions through a routing mechanism, ensuring that both common and specialized token are effectively modeled.\n\n## update after rebuttal\n\nThe rebuttal addresses my questions well. I am pleased to keep my positive scores. I also found that all the reviewers have the positive rating, and therefore I believe we have reached a consensus on this submission.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you very much for your high evaluation of our work!\n\n**Q1:** Providing theoretical grounding for MoHD’s sparsity and routing mechanisms.\n\n**R1:** We refer to our response to Reviewer fdJw Q4, where we discuss how sparse mixed activation expands effective width, reduces complexity, and improves loss. Due to space constraints, we will include the full theoretical derivation in the final version.\n\n\n**Q2:** Extending evaluations to other domains/larger models and real-world applications.\n\n**R2:** As noted in h4Rw Q3, MoHD can scale to larger models, particularly those based on Transformer architectures. For cold-started models, MoHD is expected to apply similarly. We plan to extend evaluation to more domains in future work.\n\n**Q3:** Why finer-grained routing harms performance while higher expert specialization improves outcomes.\n\n**R3:**  Our experiments show that performance peaks when hidden representations are divided into 16 sub-dimensions (each of size 256). Finer routing granularity leads to degradation due to:\n\n- **Redundancy**: Over-partitioning causes sub-dimensions to capture overlapping or low-importance features, reducing capacity efficiency.\n\n- **Routing instability**: Finer granularity makes routing decisions more sensitive and harder to stabilize during training.\n\n- **Fusion mismatch**: Our Group Fusion Layer re-integrates sparsely activated sub-dimensions. A moderate number of active groups is easier to optimize, while excessive fragmentation hinders training.\n\nAs for **expert specialization**, we interpret it as each sub-dimension focusing on a narrower subspace or data distribution. This expands representational diversity and enables the model to better capture fine-grained patterns, improving generalization across tasks.\n\n**Q4:** How sensitive is MoHD to routing variations and stability tested under different routing configurations?\n\n**R4:** Thank you for the insightful question. We analyzed MoHD’s sensitivity and stability under varied routing configurations, with key findings as follows:\n\n- **Routing stability**: As shown in *Figure 7*, sub-dimension selection probabilities mostly stay within 0.2–0.3, indicating active and balanced usage. Our **Sub-Dimension Load Balance Loss** further improves distribution and efficiency.\n\n- **Shared subspaces enhance stability**: Best performance occurs when **75% of sub-dimensions are shared** across tokens. This mitigates instability in private subspaces. Beyond a certain partitioning level, additional sub-dimensions yield diminishing returns or harm performance due to routing instability.\n\n- **Component sensitivity**: MoHD is more sensitive in **Attention** layers than **FFN** layers. Sparsification in Attention causes larger performance drops, suggesting that routing in this component requires finer, task-specific design—an area we aim to explore further.\n\nWe agree that routing stability is key to scaling MoHD and appreciate the reviewer’s focus on this aspect.\n\n**Q5:** Can MoHD adapt to non-NLP tasks without architectural modifications?\n\n**R5:** Thank you for the thoughtful question. While MoHD is developed for NLP, its core ideas may extend to other domains—particularly Vision-Language Models (VLMs), which also use Transformer architectures. Our perspective:\n\n- **Shared structure**: Vision Transformers (ViTs) treat image patches as tokens, analogous to text tokens. These patches share global patterns while retaining local uniqueness, similar to linguistic semantics.\n\n- **Core applicability**: If ViT hidden dimensions show both shared and token-specific activations, MoHD’s principle of selective sub-dimension activation may enhance efficiency and expressiveness.\n\n- **Architectural adaptation**: Despite structural parallels, visual representations differ from language. Routing and partitioning strategies would need tuning to align with visual characteristics.\n\n- **Multimodal fusion**: In VLMs, MoHD-induced sparsity could affect alignment across modalities. Incorporating sparsity-aware mechanisms without harming cross-modal interaction remains an open direction.\n\nIn short, while direct transfer isn’t trivial, MoHD’s principles are generalizable and worth exploring in vision and multimodal settings.\n\n**Q6:** Extend Reference\n\n**R6:** We sincerely thank the reviewer for the additional references. We will include a detailed discussion of them in the next revision.", "claims_and_evidence": "MoHD reportedly achieves a 1.7% performance gain with 50% fewer active parameters and a 3.7% improvement at 3× parameter scaling. These claims are empirically supported through rigorous testing. Results are convincing.", "methods_and_evaluation_criteria": "The MoHD architecture integrates shared/specialized sub-dimensions and dynamic routing, a well-grounded approach for intermediate dimensions. The proposed method aligns well with issues identified in observational studies. Evaluations in NLP tasks are convincing Parameter efficiency and task performance as metrics are practical choices.", "theoretical_claims": "The method is built upon the observations. Experimental results demonstrate the effectiveness of the proposed method.", "experimental_designs_or_analyses": "Experiments are comprehensive across many baselines to validate MoHD.", "relation_to_broader_scientific_literature": "The work is related to the design of efficient LLM architectures:\n1. Sparsity among dimension, heads, activation, etc.\n2. Conditional computation such as MoE.", "essential_references_not_discussed": "While key sparse and adaptive architecture studies are cited, recent MoE advancements requires further discussion (1).  \n1. R. Cai et al. Flextron: Many-in-One Flexible Large Language Model.", "other_strengths_and_weaknesses": "**Strengths**:  \n\n-\tThe paper is mostly well-written and easy to follow.\n\n-\tThe proposed method is technically sound with well-visualized observations.\n\n-\tExperiments are conducted with numerous model size. The ablation studies effectively validating components of MoHD. \n\n**Weaknesses**:  \n\n-\tThe experiments were conducted on models with up to 1B activated parameters. It could be helpful to show the sparsity patterns in larger models, which may be different. \n-\tExploring the combination of MoE and MoHD could help unlock sparsity benefits across multiple dimensions.\n-\tFurther ablation studies to isolate the effects of shared vs. specialized sub-dimensions and routing mechanisms would clarify why finer-grained routing harms performance while higher expert specialization improves outcomes.\n-\tHow sensitive is MoHD to routing variations (e.g., K-value selection, thresholds)? Was stability tested under different routing configurations?  \n-\tCan MoHD adapt to non-NLP tasks without architectural modifications? Or are design adjustments necessary?", "other_comments_or_suggestions": "The paper can be improved by:  \n\n- Providing theoretical grounding for MoHD’s sparsity and routing mechanisms.  \n- Extending evaluations to other domains/larger models and real-world applications.  \n- Clarifying limitations/failure modes when deploying MoHD across tasks/architectures.", "questions_for_authors": "See Weaknesses.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper introduces MoHD (Mixture of Hidden Dimensions), a novel Transformer architecture that improves parameter efficiency by leveraging hidden dimension sparsity. The key insight is that only a subset of hidden dimensions are highly activated, with some dimensions commonly activated across tokens while others are token-specific. MoHD combines shared sub-dimensions for common features with dynamically routed specialized sub-dimensions per token. The authors demonstrate strong empirical results across 10 NLP tasks, achieving 1.7% higher performance with 50% fewer activated parameters and 3.7% improvement with 3× parameter expansion at constant activation cost.

The reviewers found the paper's motivation clear and the methodology sound, with comprehensive experiments validating the approach. They particularly appreciated the novel perspective on scaling model architectures through hidden dimension sparsity. The theoretical foundations were strengthened through the authors' rebuttal, which provided detailed proofs of how mixed sparse activation achieves better risk bounds.

Several constructive suggestions were raised regarding clarification of the routing mechanism, computational costs at different scales, and comparisons with other sparsification methods. The authors provided thorough responses and committed to incorporating these improvements. The paper would benefit from expanded discussion of limitations, scalability to larger models, and potential applications beyond NLP.

The revised version is expected to fully address these points, particularly by including clearer explanations of the routing dynamics, comprehensive cost analyses across model scales, and extended comparisons with related approaches. With these additions, this work will make a valuable contribution to efficient scaling of language models.
Decision: Accept (poster)