5 Shaping model architectures via width pruning

Width pruning reshapes a transformer by trimming the MLP’s expansion layer rather than removing entire blocks, making it a more surgical alternative to depth pruning. The chapter shows how to reduce the intermediate width of GLU-based modules by identifying less useful neurons and rebuilding the affected layers so the model keeps its original input and output interfaces intact. It begins with a simple, static strategy that scores neuron pairs using the range of their weights in the gate and up projections, then keeps the most important pairs and copies their weights into smaller layers.

The results show that this kind of pruning can significantly cut parameters and speed up inference, but not without trade-offs. A 40% reduction in the expansion ratio removes a large portion of parameters while preserving depth, yet it sharply harms perplexity, reasoning, and broad knowledge tasks. At the same time, some behaviors improve unexpectedly: the pruned model becomes more concise, more instruction-following, and often less prone to confident hallucinations. This reveals that pruning does not simply make a model worse; it changes the balance of capabilities, sometimes favoring precision over creativity.

The chapter then advances to a data-driven method that combines structural weight information with real activation patterns gathered through PyTorch hooks. By capturing the L2 norm of activations at the input to the down projection during calibration, the pruning process can preserve neurons that are both structurally important and actually used by the chosen dataset. This makes it possible to create specialized variants, such as models better suited to WikiText or SMS, while retaining the same physical architecture and efficiency gains. The overall message is that width pruning is not just compression—it is a way to deliberately shape model behavior for a target domain and deployment constraint.

The complete GLU structure. Information enters with dimension hidden_size (2048 in Llama-3.2-1B) and expands to intermediate_size (8192) via gate_proj and up_proj. The gate_proj output, modulated by SiLU, is element-wise multiplied with up_proj’s output. The result is contracted back to hidden_size by down_proj.

Impact of width pruning on perplexity (PPL). Since this is a metric where "lower is better," the bars on the right show severe degradation: removing 40% of the expansion neurons causes perplexity to spike, indicating that the model has lost much of its ability to predict text sequences fluently.

Degradation across reasoning and knowledge benchmarks. Mathematical reasoning (GSM8K -68.9%) and common sense understanding (HellaSwag -41.3%) suffer the most severe losses, while reading comprehension tasks (BoolQ -1.9%, WinoGrande -4.8%) remain relatively stable.

Unexpected performance gains in instruction-following and truthfulness benchmarks. The pruned model dramatically improves at following precise instructions (IFEval +46.5%) and multi-step reasoning under constraints (MUSR +26.1%), while also becoming more truthful and less prone to generating plausible but incorrect answers (TruthfulQA +13.9% on MC2).

Hybrid importance calculation in GLU architecture. The static component combines weight ranges from gate_proj, up_proj, and down_proj. The dynamic component captures activation norms at the input of down_proj, where gated information flows. Both components multiply to produce the final importance score.

Hybrid importance combines structural capacity with activation patterns. Neuron [0] survives with moderate weights (0.9) but high activation (32). Neurons [1] and [2] are removed despite having identical or higher structural scores due to low activations (0.7 and 0.2).

Activation capture pipeline. The hooks registered in down_proj (1-2) are automatically triggered during the forward pass (3-4), intercept X_d, calculate L2 norms and accumulate them in CPU (5), finally returning an importance dictionary per layer (6-7).

Model surgery using hybrid importance. The update_model function (Listing 5.7) iterates over each layer (1-2), calls prune_neuron_pairs (3) that calculates hybrid scores by multiplying structural magnitudes by activation norms (4), selects the top-K neurons and creates new layers (5), replaces the original layers (6), and returns the specialized model (7).

The dichotomy visualized in Llama-3.2-3B. Panel A shows the progressive collapse of fragile capabilities (knowledge and reasoning) as the expansion ratio decreases. Panel B reveals the paradox: robust capabilities (instruction-following and truthfulness) improve dramatically, with IFEval reaching a peak of 174% of baseline at 1.6x expansion (40% pruning). Performance is normalized to the base model (2.67x expansion = 100%).

• Our width pruning technique modifies LLMs by removing specific neurons from MLP expansion layers in GLU architectures, offering more fine-grained control than depth pruning's aggressive block elimination. This technique can reshape model behavior without destroying all capabilities.
• In GLU architectures, the gate_proj and up_proj layers work as a synchronized pair. Neurons must be pruned simultaneously in both projections based on a combined importance score.
• The reconstruction process creates new, smaller Linear layers while preserving connections with attention modules. By maintaining hidden_size dimensions at module boundaries, the surgery remains internal to the MLP block without breaking the Transformer architecture.
• Data-driven neuron selection combines structural importance (weight magnitudes) with dynamic importance (activation norms). This hybrid approach multiplies a neuron's structural score by its actual usage during calibration, ensuring only neurons that are both structurally significant and actively used survive.
• PyTorch hooks capture activations at down_proj's input, after the gating mechanism but before dimensionality reduction, using register_forward_hook(). The L2 norm measures activation intensity, rewarding specialist neurons with sharp peaks over those with constant low-level noise.
• Domain-specialized models emerge from identical pruning percentages applied with different calibration data. A WikiText-calibrated model preserves neurons for formal text processing, while an SMS-calibrated model retains neurons for conversational language, despite sharing the same architecture externally.
• The efficiency gains from width pruning are most visible in compute-bound scenarios.
• Energy efficiency gains result from architectural changes, not behavioral modifications. However, width pruning's benefits are limited by the attention module bottleneck.
• Aligning the resulting intermediate_size to be divisible by 64, 128, or 256 maximizes GPU Tensor Core utilization in compute-bound scenarios. This optimization matters most in production deployments with high traffic on modern GPUs (A100/H100).
• Research papers "Fragile Knowledge, Robust Instruction-Following" (Martra, 2025) and "CFSP" (Wang et al., 2024) validate the dichotomy between degrading and improving capabilities, and the effectiveness of activation-based importance metrics. The practical implementation adapts their core insights into a simpler hybrid approach that combines structural capacity with dynamic usage.

FAQ