Mixture of Hidden-Dimensions: Not All Hidden-States’ Dimensions are Needed in Transformer

Thank you very much for your high evaluation of our work!

Q1: Providing theoretical grounding for MoHD’s sparsity and routing mechanisms.

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.

Q2: Extending evaluations to other domains/larger models and real-world applications.

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.

Q3: Why finer-grained routing harms performance while higher expert specialization improves outcomes.

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:

• Redundancy: Over-partitioning causes sub-dimensions to capture overlapping or low-importance features, reducing capacity efficiency.

• Routing instability: Finer granularity makes routing decisions more sensitive and harder to stabilize during training.

• 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.

As 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.

Q4: How sensitive is MoHD to routing variations and stability tested under different routing configurations?

R4: Thank you for the insightful question. We analyzed MoHD’s sensitivity and stability under varied routing configurations, with key findings as follows:

• 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.

• 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.

• 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.

We agree that routing stability is key to scaling MoHD and appreciate the reviewer’s focus on this aspect.

Q5: Can MoHD adapt to non-NLP tasks without architectural modifications?

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:

• 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.

• 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.

• Architectural adaptation: Despite structural parallels, visual representations differ from language. Routing and partitioning strategies would need tuning to align with visual characteristics.

• 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.

In short, while direct transfer isn’t trivial, MoHD’s principles are generalizable and worth exploring in vision and multimodal settings.

Q6: Extend Reference

R6: We sincerely thank the reviewer for the additional references. We will include a detailed discussion of them in the next revision.

We are greatly encouraged by the reviewer’s positive feedback. Below, we address each of your comments in detail.

Q1: Theoretical Grounding

R1: Please see our response to Reviewer fdJw (Q4) for a more detailed explanation. In summary, we theoretically ground MoHD by showing that:

1. Sparse activation leads to effective width expansion;
2. This expanded width results in lower empirical loss;
3. A hybrid scheme of private and shared sub-dimensions yields better performance than either alone.

Q2: Computational Costs Compared to MoE and Other Methods

R2: We summarize MoHD’s computational characteristics relative to other approaches:

• 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.

• 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.

• 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.

Q3: Scaling Costs

R3: Please refer to Reviewer h4Rw, Q4 for detailed cost analysis across model scales. In brief:

• Routing and fusion layers add minimal parameter overhead.
• The primary computational increase comes from the WTE layer, which scales with model width.
• Notably, the performance benefits of MoHD cannot be solely attributed to increased WTE size, as demonstrated in our ablation and FLOPS analysis.

Q4: How MoHD Improves Performance While Reducing Activations

R4:

1. Sparsity and Redundancy in Hidden Dimensions
   Empirical 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.

2. Component-wise Sensitivity to Sparsification
   Consistent 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.

3. Grouped Fusion Mechanism
   MoHD achieves efficiency via:

• Activation scaling: Dynamically adjusts magnitudes of active units for representational stability.
• Grouped fusion: Aggregates sparse activations to preserve expressive capacity.

These 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.

Q5: Limitations of MoHD

R5: We appreciate the chance to elaborate on current limitations:

• Hyperparameter Sensitivity
  MoHD 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.

• Routing Optimization Challenges
  Unlike 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.

• WTE Layer Growth
  As 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.

• Information Degradation at Scale
  Despite 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.

Q6: Extend Reference

R6: We sincerely thank the reviewer for the additional references. We will include a detailed discussion of them in the next revision.

Thank you for the valuable comments and suggestions.

Q1: Comparison with MoE and other sparsification methods

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.

Q2: The illustration and motivation of routing mechanism and activation flow

R2: Thank you and we will further improve the writing. A detailed explanation is provided below:

• 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.

• 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.

Q3: MoHD's applicability to larger LLMs

R3: Thank you for the suggestion. We believe applying MoHD to larger-scale language models offers several potential benefits:

• 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.

• 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.

• 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.

Due 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.

Q4: The computational costs at different scales can be further clarified

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:

As 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.

It’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:

• The FLOPS of the x4-width model is significantly higher than that of x3, yet the performance improvement is marginal.

• 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.

These observations reinforce our argument that MoHD achieves efficiency primarily through structural sparsity and expert specialization, not brute-force scaling.

Q5: The clarification of symbols.

R5: We apologize for the confusion and will provide clearer and more readable content in the next revision.

Thank you for the valuable comments and suggestions.

Q1: On MoHD-495M’s performance on WinoGrande (WG)

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.

That said, performance dips in certain settings may stem from:

• WG-specific reasoning demands: WG requires fine-grained commonsense inference, which may benefit from further adaptation of MoHD’s routing and activation strategies.
• 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.

Q2: On integrating MoE into MoHD

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.

However, such integration would introduce significant optimization and engineering complexity. Further empirical studies are needed to assess whether their combination yields synergistic or conflicting effects.

Q3: MoHD vs. MoE

R3: Key differences:

• Routing granularity: MoHD routes at the hidden dimension level, tailoring token-specific subspace activations. MoE routes at the expert (subnetwork) level in the FFN.
• Component scope: MoHD applies to both Attention and FFN, while MoE is typically confined to FFN projections.
• Capacity and interpretability: MoHD expands model width, increasing per-token representation capacity. MoE expands the FFN depth, aiding memory but not width.
• Efficiency: MoHD reduces redundant activations in both Attention and FFN, offering better scaling under fixed activation budgets (see Table 2).
• Challenges: MoHD’s routing across hidden dimensions is less explored and demands novel optimization and implementation.

Despite these challenges, MoHD shows notable improvements over MoE when trained from scratch, highlighting its promise for efficient scaling.

Q4: Theoretical Proof

R4: We prove that mixed sparse activation achieves strictly better risk bounds.

Lemma 1 (Unbiased Sparse Forward Pass).
Let $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:
$\mathbb{E}[\hat{h}] = h; \quad \mathbb{E}[||\hat{h}||^2] = ||h||^2/p + (1-1/p) \sum_{j=1}^n h_j^2$

Proof: Linearity of expectation and variance decomposition.

Corollary 1.1 (Effective Width).
For $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.

Lemma 2 (Approximation Error Decay with Width [Barron, 1993]).
Let $f^* \in \mathcal{F}$, where $\mathcal{F}$ is a Barron space. For a network with width $n$, the approximation error satisfies:

$$\epsilon(n) \leq \frac{C_f}{\sqrt{n}},$$

where $C_f$ is the Barron norm of $f^*$.

Corollary 2.1 (Mixed Activation Lowers Error).
Define $s$ encodes global context, $g$ decodes shared features, $u$ models per-token contributions; $\Theta$ denotes all parameters.

If $f^* = g(s(x)) + \sum_i u(x_i)$ with non-constant $g$, then:

• Token-only networks ($n = n_u$) have $\epsilon_B \geq \frac{C_g}{\sqrt{n_u}}$.
• Mixed networks ($n = n_s + n_u$) have $\epsilon_A \leq \frac{C_g}{\sqrt{n_s}} + \frac{C_u}{\sqrt{n_u}}$.
  Choosing $n_s = \Theta(n_u)$ yields $\epsilon_A < \epsilon_B - \delta$ for $\delta = \frac{C_g}{2\sqrt{n_s}}$.

Lemma 3 (Rademacher Complexity of Shared Dimensions).
Let $\mathcal{H}_A$ (mixed) and $\mathcal{H}_B$ (token-only) have equal total parameters. Then:

$$\text{Rad}(\mathcal{H}_A) \leq \text{Rad}(\mathcal{H}_B) - \Delta,$$

where $\Delta = \Omega\left(\sqrt{\frac{n_s}{m}}\right)$ for $m$ samples.

Proof: Shared Dimensions reducing the VC dimension [Bartlett, 1998]. Apply Dudley entropy integral.

Theorem (Risk Bound).
For risk $R(f) = \epsilon(f) + C\cdot\text{Rad}(\mathcal{H})$:

$R_A \leq \epsilon_A^{\text{(Lower error)}} + C\cdot\left(\text{Rad}(\mathcal{H}_B) - \Delta\right)^{\text{(Lower complexity)}} < R_B$

Sparse 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$.