VersatileFFN: Achieving Parameter Efficiency in LLMs via Adaptive Wide-and-Deep Reuse (2512.14531v1)

Abstract: The rapid scaling of LLMs has achieved remarkable performance, but it also leads to prohibitive memory costs. Existing parameter-efficient approaches such as pruning and quantization mainly compress pretrained models without enhancing architectural capacity, thereby hitting the representational ceiling of the base model. In this work, we propose VersatileFFN, a novel feed-forward network (FFN) that enables flexible reuse of parameters in both width and depth dimensions within a fixed parameter budget. Inspired by the dual-process theory of cognition, VersatileFFN comprises two adaptive pathways: a width-versatile path that generates a mixture of sub-experts from a single shared FFN, mimicking sparse expert routing without increasing parameters, and a depth-versatile path that recursively applies the same FFN to emulate deeper processing for complex tokens. A difficulty-aware gating dynamically balances the two pathways, steering "easy" tokens through the efficient width-wise route and allocating deeper iterative refinement to "hard" tokens. Crucially, both pathways reuse the same parameters, so all additional capacity comes from computation rather than memory. Experiments across diverse benchmarks and model scales demonstrate the effectiveness of the method. The code will be available at https://github.com/huawei-noah/noah-research/tree/master/VersatileFFN.

• The paper introduces VersatileFFN, a drop-in FFN module that decouples LLM capacity from parameter count via adaptive width and depth reuse.
• It employs a virtual expert paradigm and recursive processing to dynamically allocate compute based on token difficulty, boosting accuracy.
• Empirical results demonstrate higher accuracy on reasoning tasks with minimal parameter increase compared to standard FFN and MoE methods.

VersatileFFN: Parameter-Efficient Feed-Forward Architectures for LLMs

Motivation and Problem Setting

The continued expansion in the parameter count of LLMs, motivated by empirical scaling laws, directly increases memory costs and impedes model deployment under realistic hardware and economic constraints. Existing strategies for maximizing parameter efficiency, including post hoc pruning, quantization, and low-rank adaptation, reduce storage requirements but do not substantially increase the model's representational power; instead, they compress existing architectures, bumping up against the upper limit induced by the original network's topology. Conversely, methods that enhance capacity—such as Mixture-of-Experts (MoE) designs—typically achieve computational sparsity but incur heavyweight parameter duplication due to the instantiation of real experts. This paper introduces VersatileFFN, which challenges the typical paradigm by seeking to decouple model capacity from parameter count via systematic and unified parameter reuse along two complementary dimensions: width (multi-expert emulation by virtualizing a single FFN) and depth (recursive MLP application for adaptive compute).

Architectural Contributions

VersatileFFN is presented as a drop-in replacement for standard Transformer FFN blocks. All innovations occur within the FFN portion, leaving self-attention and residual connectivity untouched.

Figure 1: VersatileFFN integrates a width-versatile pathway for rapid execution and a depth-versatile pathway for deep reasoning, both derived from a shared base FFN. This design enables a flexible trade-off between the two computational dimensions.

Width-Versatile Pathway: Adopts the MoE philosophy but with strictly virtual experts. It partitions the output of a single FFN into multiple non-overlapping subspaces (“sub-experts”) via strided slicing through the shared parameters, as formalized in the paper. Each sub-expert is functionally orthogonal and covers a disjoint slice of the hidden dimension, ensuring expert behaviors are diverse but the storage footprint is identical to the baseline MLP.

Figure 2: Decoupling a large real expert into multiple lightweight virtual experts facilitates expert diversity without parameter proliferation.

Route selection is governed by a token-wise gating MLP and standard top-$k$ activation, emulating sparsity as in physical MoE networks but without the parameter cost.

Figure 3: Trade-off and scaling behaviors for various expert counts and numbers of activated experts in the width-versatile pathway.

Depth-Versatile Pathway: Applies the same shared FFN recursively per token, mirroring recursive architectures such as Universal Transformers but introducing token-level adaptivity via Gumbel-Softmax-based loop count prediction. Each token's representation is refined for a learned number of steps (subject to a maximum cap). This structure is fully differentiable using straight-through estimation and soft aggregation during training, with discrete execution during inference for maximal efficiency.

Difficulty-Aware Fusion: The expected loop count serves as a proxy for token difficulty. The final output is a linear fusion between the width-wise and depth-wise branches, where the fusion coefficient is a function of the predicted loop count for each token. Efficient tokens primarily utilize the width branch; more complex ones allocate increased computation to recursive processing. This dynamic resource allocation is realized under a strictly fixed parameter regime.

Empirical Analysis

The experimental validation encompasses rigorous continued pretraining and extensive zero-shot evaluation across standard QA, reasoning, and commonsense benchmarks (e.g., PIQA, ARC-easy/ARC-challenge, HellaSwag, OpenBookQA, SciQ, CommonsenseQA, WinoGrande).

VersatileFFN is compared against:

• Dense Baseline: Fixed-size FFNs.
• Physical MoE: Eight real experts (with top-2 routing), inducing high parameter overhead.
• $k$-Loop Baseline: Standard FFN applied $k$ times per token; introduces computational cost linearly with $k$.

Strong numerical outcomes are reported:

• On the 720M parameter scale, VersatileFFN achieves an average accuracy of 57.03% across benchmarks, surpassing MoE (55.87%) and $k$-Loop (56.55%) variants at equal or lower computational cost.
• Parameter efficiency is highlighted: VersatileFFN adds less than 0.1% to the baseline parameter count, compared to a 60%+ increase for MoE.
• In terms of FLOPs, VersatileFFN provides a sublinear increase relative to the same-layer $k$-Loop architectures with comparable performance.

Claim: VersatileFFN empirically achieves superior accuracy under a tightly controlled parameter and compute budget, particularly excelling in reasoning-heavy tasks and demonstrating robustness in generalization.

Model Behavior and Feature Analysis

Loop count diagnostics on ARC-challenge elucidate adaptive compute allocation per token and per layer. The smaller model centralizes recurrence in later layers, while the larger model allocates depth in mid-network layers, suggestive of emergent hierarchical processing strategies.

Figure 4: The average predicted loops per layer on ARC-c dataset. Left: 354M base model. Right: 720M base model.

Representation analysis at layer 0 shows nontrivial divergence between the width-wise and depth-wise output features, despite parameter sharing. These branches capture complementary information, reinforcing the utility of joint fusion.

Figure 5: Heatmap of cosine similarity at Layer 0. Left: 354M base model. Right: 720M base model.

Token-level gating analysis using word clouds confirms that low loop-count tokens are dominated by high-frequency syntactic elements, while high loop-count tokens concentrate on semantically pivotal or rare vocabulary.

Figure 6: Word cloud for various gating coefficients on ARC-c, mapping higher compute to more task-relevant, semantically complex tokens.

Component Effectiveness and Ablations

A comprehensive ablation reveals that individual pathways each yield gains over baseline, while their combination with adaptive fusion offers the best trade-off. Increasing the virtual expert count or the recursion depth past moderate levels does not linearly enhance accuracy and can induce diminishing returns or overfitting, highlighting the value of parameter-efficient balancing. VersatileFFN’s effectiveness persists even without continued pretraining.

Practical and Theoretical Implications

Practically, VersatileFFN enables deployment of LLMs with resource demands closely tracking those of the base model, yet provides capacity benefits normally associated with architectural scaling or MoE expansion. Given that both modalities (width and depth) reuse parameters entirely, all gains in expressivity are “compute-derived” rather than “memory-derived.” This design is highly conducive to low-memory, high-FLOP accelerator nodes and to distributed settings where communication constraints dominate.

Theoretically, VersatileFFN suggests that intrinsic parameter reuse—especially in the FFN block—allows for a richer class of functions without violating parameter budgets. This extension of the looped transformer and virtual expert paradigms invites further investigation into hybrid architectures, curriculum- or token-aware compute allocation, or even more granular sharing schemes.

Conclusion

VersatileFFN demonstrates that parameter-efficient, adaptive FFNs can robustly improve LLM performance under severe memory constraints by leveraging recursive application and virtualized, structured expert diversity. The evidence supports a strategic shift toward compute-heavy, memory-light architectures, where “scaling” is accomplished via dynamic, token-level path selection and not just gross parameter accumulation. Future developments will likely extend this approach with more granular path and expert definitions, multi-modal gating, and optimized compiler/runtime support for recursive/tiled compute in resource-bounded environments.

---

Citation: "VersatileFFN: Achieving Parameter Efficiency in LLMs via Adaptive Wide-and-Deep Reuse" (2512.14531)