PaperlayerView exact revision on arXiv

Source-linked AI summary

VersatileFFN: Achieving Parameter Efficiency in LLMs via Adaptive Wide-and-Deep Reuse

Ying Nie, Kai Han, Hongguang Li, Hang Zhou, Tianyu Guo, Enhua Wu, Xinghao Chen, Yunhe Wang

arXiv:2512.14531v2cs.CL

TL;DR

Large language models’ parameter growth brings prohibitive memory costs, while compression methods do not extend the base model’s architectural capacity. VersatileFFN reuses a shared FFN across width- and depth-oriented pathways, and experiments report consistent advantages over other methods, including higher average accuracy than MoE and 6-Loop at 1.21B.

  • Problem

    LLM scaling raises memory costs, while pruning and quantization mainly compress pretrained models without expanding their architectural capacity.

    Source: §1 · p. 1

  • Method

    VersatileFFN reuses one FFN across width-wise virtual sub-experts and recursive depth processing, with a difficulty-aware gate balancing the pathways.

    Source: §1 · pp. 1–2, §3.2 · p. 4, §3.3 · p. 4

  • Results

    VersatileFFN consistently outperforms other methods across experiments; at 1.21B it reaches 60.47% average accuracy versus MoE’s 59.65% and 6-Loop’s 60.05%.

    Source: §4.2 · p. 7, §5 · p. 8

  • Takeaways & Limitations

    The results support intelligent weight reuse and adaptive computation as viable alternatives to simply scaling model size.

    Source: §5 · p. 8

    The difficulty-aware approach assumes required recursive depth is an intrinsic proxy for semantic difficulty, from trivial tokens needing one loop to complex tokens benefiting from deeper refinement.

    Source: §3.4 · p. 5

Abstract

from arXiv · show

The rapid scaling of Large Language Models (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 is available at https://github.com/huawei-noah/noah-research/tree/master/VersatileFFN.

1. Introduction

VersatileFFN targets memory-constrained scaling by reusing a shared FFN across width- and depth-versatile computation; the paper reports consistent gains over parameter- or FLOPs-matched methods.

  • Parameter growth strains deployment through memory, cost, communication, and engineering overhead, while compression methods preserve rather than extend the base architecture’s capacity.Source: §1 · p. 1
  • VersatileFFN combines width- and depth-versatile pathways derived from a shared base FFN to trade computation across dimensions.Source: §1 · pp. 1–2, Figure 1 · p. 2
  • The width pathway derives a mixture of sub-experts from one FFN by composing non-overlapping hidden subspaces, avoiding separate expert parameters.Source: §1 · p. 2
  • The depth pathway recursively reapplies the same FFN for token-specific iterative refinement rather than allocating separate parameters at each step.Source: §1 · p. 2
  • VersatileFFN consistently outperforms other parameter-matched or FLOPs-matched methods across benchmarks and model scales.Source: §1 · p. 2
  • Its difficulty-aware gate uses expected depth iterations as a token-difficulty proxy, favoring width processing for easy tokens and depth refinement for hard tokens.Source: §1 · p. 2

2. Related Work

Related work covers routed expert models and parameter-efficient MoE variants, alongside recursive methods that adapt computation to input complexity. The paper identifies width-wise reuse as an underexplored complement to recursive computation.

2.1. Mixture-of-Experts

2.2. Recursive Computation and Adaptive Inference

3. Method

VersatileFFN replaces the Transformer FFN with width- and depth-versatile pathways that reuse shared weights while retaining canonical self-attention. A difficulty-aware gate balances rapid width-wise processing against deeper iterative refinement for harder tokens.

3.1. Architectural Overview

  • VersatileFFN retains self-attention and replaces the standard FFN with two pathways that share its weights: width-wise virtual experts and depth-wise recursive refinement.The pathways are fused dynamically according to token difficulty.Source: §3.1 · p. 3

3.2. Width-Versatile Mechanism

  • Virtual experts are non-overlapping hidden-subspace views of shared weights, avoiding separate expert matrices while preserving functional orthogonality.A learnable router uses Top-K selection to combine the selected experts; training also uses an auxiliary load-balancing loss to prevent expert collapse.Source: §3.2 · p. 4, Figure 2 · p. 4

3.3. Depth-Versatile Mechanism

  • The depth pathway repeatedly applies the full shared FFN to refine representations, while a Gumbel-Softmax controller predicts token-specific iteration counts.Training uses a straight-through estimator; inference makes a discrete loop-count decision and returns the corresponding state.Source: §3.3 · pp. 4–5

3.4. Difficulty-Aware Fusion

  • Expected loop count acts as a difficulty proxy: the fusion gate favors width for easy tokens and shifts toward depth for hard tokens.Figure 3 plots training loss and average accuracy across configurations of the two mechanisms; the method also uses early exit and can prune the width branch when its contribution is negligible.Source: §3.4 · p. 5, Figure 3 · p. 6

4. Experiments

Across the tested model scales, VersatileFFN achieves strong zero-shot accuracy while adding negligible parameters; its adaptive computation and ablations examine how tokens, loops, and pathway components affect results.

4.1. Implementation Details

  • Experiments train 354M, 720M, and 1.21B models on 40B, 70B, and 100B tokens, then compare continued pre-training against MoE and k-Loop.The models use FineWeb-Edu and OLMo2 settings; both baselines continue training from the same pretrained model on the same corpus.Source: §4.1 · p. 6, §4.2 · p. 6

4.2. Evaluation Results

  • 60.47% average accuracy on the 1.21B model exceeds MoE at 59.65% and 6-Loop at 60.05%.VersatileFFN also leads MoE by 2.81% on ARC-e.Source: §4.2 · p. 7, Table 1 · p. 7
  • At 350M, VersatileFFN uses approximately 45% fewer FLOPs than 6-Loop while achieving superior accuracy, with negligible parameter overhead.MoE increases parameters from 1.21B to 1.97B (+63%).Source: §4.2 · p. 7, Table 2 · p. 7

4.3. Visualized Analysis

  • Predicted loop allocation shifts from late layers at 354M, to middle layers at 720M, to a front-loaded profile at 1.21B.Figure 4 plots average predicted loops per layer on ARC-c for these three model sizes.Source: §4.3 · p. 7, Figure 4 · p. 8
  • On ARC-c, lower gating coefficients indicate more loops; the word-cloud analysis associates lower coefficients with action verbs and higher ones with generic terms.Examples include “clean” and “cut” at the lower-coefficient end, and “make” and “use” at the higher-coefficient end.Source: §4.3 · pp. 7–8, Figure 5 · p. 8

4.4. Ablation Studies

  • Each pathway individually outperforms the baseline, while combining all modules yields the lowest loss and highest average accuracy.The ablation reports simple averaging when fusion gating is removed and characterizes the two branches as complementary.Source: §4.4 · p. 8, Table 3 · p. 8
  • Four loops achieve the best zero-shot average accuracy among tested loop counts, while the (8, 2) expert configuration achieves the best accuracy among tested configurations.Six loops minimize training loss, but average accuracy peaks at four loops; the authors suggest excessive recurrence may overfit.Source: §4.4 · p. 8, Table 4 · p. 8
  • Trained from scratch, VersatileFFN reaches 51.14% accuracy, improving 3.16% over the Base and outperforming the other methods.Table 5 reports performance without continued pre-training.Source: §4.4 · p. 8, Table 5 · p. 8

5. Conclusion

The conclusion presents VersatileFFN as a compute-flexible alternative to parameter scaling, combining width and depth reuse within a constrained parameter budget.

  • VersatileFFN combines a width-versatile path for efficient computing with a depth-versatile path for iterative reasoning.The conclusion frames the design as prioritizing computational flexibility over memory expansion.Source: §5 · p. 8
  • The conclusion reports that VersatileFFN consistently outperforms other methods in extensive experiments.Source: §5 · p. 8
  • The authors present intelligent weight reuse and adaptive computation as alternatives to scaling model size, with resource-constrained deployment as a target context.Source: §5 · p. 8

A. Inference Optimization

Inference computes a token’s loop prediction and difficulty-aware fusion weight, conditionally runs the width pathway, and iterates the shared FFN until the selected loop count.

  • The controller predicts a discrete loop count and computes the difficulty-aware fusion scalar before pathway computation.The algorithm uses the maximum-probability loop count and derives the fusion scalar from the controller.Source: A. Inference Optimization · p. 12
  • When the fusion weight is positive, inference constructs virtual experts from shared weights and applies Top-k routing; otherwise, it prunes the width branch.Source: A. Inference Optimization · p. 12
  • The depth pathway recursively updates the hidden state using the shared weights through the predicted loop count, then fuses pathway outputs into Y.The algorithm selects the final recursive state and returns the fused output.Source: A. Inference Optimization · p. 12

B. Implementation Details

The appendix specifies continued-pretraining settings and the zero-shot benchmarks used to evaluate the models.

Experiments Setup

  • Continued pre-training uses sequence length 4,096, AdamW, weight decay 0.1, gradient clipping at 1.0, and cosine learning-rate decay after a 5% warm-up.Peak learning rates and global batch sizes vary across the 354M, 720M, and 1.21B models.Source: B. Implementation Details · p. 12

Evaluation Benchmark Setup

  • Evaluation covers eight benchmarks, including PIQA, HellaSwag, OpenBookQA, SciQ, ARC-e, ARC-c, CommonsenseQA, and Winogrande.OLMES is used for zero-shot assessment.Source: B. Implementation Details · p. 12

C. Features of two branches

At Layer 0, the width- and depth-versatile branches produce complementary representations despite inheriting the same base parameters.

D. Evolution of Downstream Accuracy During Training

Across three model scales, accuracy generally rises with training tokens, and larger models achieve higher peaks and steeper learning curves, though some trajectories oscillate.

Loading 2512.14531v2…