Sparse MLP Tuning

Sparse MLP tuning operates exclusively on the Multi-Layer Perceptron (MLP) or feed-forward network blocks within a transformer architecture. It leaves the self-attention mechanisms, embedding layers, and layer normalization parameters completely frozen. This is based on the empirical finding that for many downstream tasks, the knowledge adaptation required is largely captured by modifying the feature transformations within these dense, non-linear layers, while the attention patterns for contextual understanding remain largely transferable from the pre-trained base.

The method does not update all parameters within the targeted MLP layers. Instead, it employs a sparsity-inducing technique to select a critical subset. Common mechanisms include:

• Learned Binary Masks: A trainable gating function determines which weights receive non-zero gradients.
• Magnitude-Based Pruning: Weights with the smallest absolute values in the pre-trained model are frozen, assuming they are less critical.
• Gradient Masking: Gradients are computed for all MLP parameters, but only the top-k% by gradient magnitude are applied, setting the rest to zero. This results in a sparse parameter update, where often <10-30% of the MLP's weights are actually modified.

By freezing most of the model and sparsely updating only a fraction of the MLP parameters, this method achieves significant efficiency gains:

• Reduced GPU Memory: Only the gradients and optimizer states for the active subset of parameters need to be stored, lowering peak memory usage during training.
• Faster Training Steps: Backpropagation and optimizer updates are computationally cheaper due to the sparse computation graph.
• Minimal Storage Overhead: The fine-tuned model is represented by the base model checkpoint plus a small sparse delta (the changed weights and their indices), not a full duplicate. This is crucial for managing many task-specific adaptations.

Sparse MLP tuning aims to preserve the model's general world knowledge encoded during pre-training while adapting to a new domain. By limiting updates to a sparse subset of MLP weights, it inherently regularizes the model, reducing catastrophic forgetting of the base task. Performance is often on par with or slightly below full fine-tuning for in-domain tasks, but with far better retention of out-of-domain generalization capabilities. It is particularly effective for tasks that are semantically related to the model's pre-training corpus.

Sparse MLP tuning is distinct from other popular PEFT approaches:

• vs. Adapters/LoRA: These methods add new trainable parameters (small matrices) to the model. Sparse MLP tuning directly updates a subset of the existing base model weights.
• vs. Prefix/Prompt Tuning: These methods modify the input space (prepend trainable embeddings). Sparse MLP tuning modifies the internal model weights.
• vs. Full Fine-Tuning: The obvious contrast is in the order-of-magnitude reduction in trainable parameters and the explicit sparsity constraint. It represents a weight-based, selective approach rather than an additive or input-based one.

This technique is favored in scenarios requiring efficient adaptation with strong knowledge retention:

• Domain Adaptation: Tuning a general LLM (e.g., Llama 3) for specialized jargon in legal, medical, or financial documents.
• Multi-Task Serving: Hosting a single base model instance that can serve multiple downstream tasks via different sparse masks, reducing serving infrastructure costs.
• Edge/On-Device AI: Enabling model personalization on devices with limited memory by storing and updating only small sparse deltas.
• Continual Learning: Sequentially learning new tasks with minimal interference, as the sparse updates for each task occupy largely non-overlapping parameter subsets.

Sparse fine-tuning is the overarching parameter-efficient adaptation technique where only a small, strategically selected subset of a pre-trained model's weights are updated. This contrasts with dense fine-tuning, which updates all parameters.

• Core Principle: Achieves efficiency by identifying and training a sparse set of 'important' parameters.
• Methods: Includes parameter masking, gradient masking, and learned sparse masks.
• Goal: Maintains most of the model's pre-trained knowledge while adapting to a new task with minimal compute.

Selective fine-tuning is a strategy that identifies and trains only the most task-relevant parameters within a pre-trained model. It is often synonymous with sparse fine-tuning but emphasizes the criteria for selection.

• Selection Heuristics: Parameters are chosen based on importance scores derived from gradient magnitude, weight magnitude, or second-order information like the Hessian.
• Granularity: Can be applied at different levels: sparse layer tuning (entire layers), sparse attention tuning (attention mechanisms), or sparse neuron tuning (individual units).
• Application: Used when prior knowledge or analysis suggests specific model components are most relevant for the target domain.

Parameter masking is a core technique in sparse fine-tuning where a binary mask is applied to the model's weights or gradients to selectively freeze or enable updates.

• Mechanism: A mask M ∈ {0,1} is element-wise multiplied with the weight matrix W. Where M=0, the gradient is zeroed, preventing updates.
• Types: Can be static (defined before training) or dynamic (learned or updated during training). A sparse learned mask uses a parameterized gating function.
• Advantage: Explicitly enforces sparsity in the parameter update, directly controlling the number of trainable parameters.

Sparse optimization refers to a class of gradient-based algorithms, like Sparse SGD or Sparse Adam, designed to handle models where a large proportion of gradients are zero due to masking or inherent sparsity.

• Efficiency: Optimizers can skip computations for masked parameters, leading to faster training steps and reduced memory footprint.
• Challenge: Standard optimizers like Adam maintain momentum and variance estimates for all parameters; sparse variants must correctly handle state updates for frozen weights.
• Use Case: Essential for making sparse fine-tuning computationally beneficial, not just parameter-efficient.

Sparse importance scoring is the process of ranking a model's parameters using a heuristic to determine which weights are most critical for adaptation to a new task. This guides the creation of the sparse mask.

• Common Metrics:
  • Sparse magnitude pruning: Uses the absolute value of pre-trained weights.
  • Gradient-based scoring: Uses the L2 norm of gradients over a batch.
  • Sparse Fisher Information: Estimates parameter importance based on the expected gradient variance.
  • Sparse Hessian-based selection: Uses the diagonal of the Hessian matrix to estimate loss sensitivity.
• Outcome: Produces an ordered list, allowing selection of the top-k% most important parameters to tune.

Sparse task vectors are the difference between a fine-tuned model's weights and its pre-trained base weights (W_task - W_base), where this delta vector is constrained or encouraged to be sparse.

• Property: Represents the minimal change needed for task adaptation. Sparsity enables efficient storage and model composition.
• Application in Model Merging: Sparse task vectors are key to sparse model merging and sparse TIES-Merging, where multiple task-specific deltas are combined to create a multi-task model without interference.
• Theoretical Link: Supports the sparse intrinsic dimension hypothesis, suggesting effective adaptation exists in a low-dimensional, sparse subspace.