Research map / Architectures and efficiency / Attention and transformers

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

ICML 2025 · Architectures and efficiency › Attention and transformers · approach: Transformer block and depth design

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes MoHD, a sparse architecture that routes hidden-state dimensions into shared and token-specialized sub-dimensions, enabling hidden-dimension expansion at constant activated parameters, and evaluates it by pretraining LLaMA2-style models from scratch.

Problem

Scaling Transformer hidden dimensions increases parameters and compute uniformly for all tokens; the sparsity of hidden-dimension activations has not been exploited to scale hidden dimensions at controllable cost.

Approach

Observes that a small subset of hidden dimensions is highly activated, with some shared across tokens and some token-specific; partitions hidden dimensions into sub-dimensions, always activates shared ones and top-K routes specialized ones per token, then restores activation magnitude via sub-dimension scaling and a Monarch-based grouped fusion layer.

Key results

• MoHD x3 at 1.48B total parameters (396M activated) improves average accuracy by 2.2% over LLaMA2-495M baseline on the 10-task suite.
• MoHD 50% at 355M retains roughly 99% of baseline performance with about half the activated parameters (145M vs 289M).
• MoHD outperforms the DeepSeek-V2 MoE baseline by about 2% on average in both compression and expansion settings at 355M scale.

Measured effects reported in the paper

Surprising findings

• Sparsifying the FFN hidden dimensions gives minimal loss or even slight perplexity improvement (-0.30 PPL at 50% sparsity), suggesting redundancy in FFN.
• Sparsifying attention is more sensitive to sparsification than FFN, with a +1.04 PPL increase at 50% sparsity.

Limitations

• Experiments limited to models up to 4.55B total parameters pretrained from scratch on 50B tokens; scaling to large LLMs is untested.
• Sparse routing and Monarch fusion add implementation complexity and may not yield wall-clock speedups (not reported).

Experimental setup

Datasets: RedPajama

Benchmarks: SciQ PIQA WinoGrande ARC-Easy ARC-Challenge HellaSwag LogiQA BoolQ LAMBADA MMLU

Metrics: accuracy perplexity

Baselines: LLaMA2-355M LLaMA2-495M LLaMA2-1.13B LLaMA2-2.7B DeepSeek-V2 MoE

Builds on: LLaMA2 DeepSeek-V2 MoE Sheared-LLaMA Monarch matrices

Closest papers by meaning

• SwitchHead: Accelerating Transformers with Mixture-of-Experts Attention (NeurIPS 2024 · similarity 0.86)
• Predictive Geometry of Hidden Trajectories in Transformers (NeurIPS 2026 · similarity 0.85)
• TokenFormer: Rethinking Transformer Scaling with Tokenized Model Parameters (ICLR 2025 · similarity 0.85)
• Transformer Block Coupling and its Correlation with Generalization in LLMs (ICLR 2025 · similarity 0.85)
• Understanding Scaling Laws with Statistical and Approximation Theory for Transformer Neural Networks on Intrinsically Low-dimensional Data (NeurIPS 2024 · similarity 0.85)
• Separations in the Representational Capabilities of Transformers and Recurrent Architectures (NeurIPS 2024 · similarity 0.84)