🛠️ Gemma 4 31B SwiGLU Fusion Toolkit
This toolkit provides a robust pipeline for generating and injecting custom, functionally correct MLP weights into a Gemma 4 31B safetensors model. It is specifically designed to handle Gemma 4's unique architectural features—namely its interleaved sliding-window attention (SWA) and global full-context layers—while preserving model stability through advanced alignment and fusion techniques.
✨ Features
- Targeted Architecture: Explicitly accounts for Gemma 4's interleaved attention structure (5 SWA + 1 Global per period).
- Double-Wide MLP Support: Automatically detects and generates intermediate sizes for global full-context layers.
- Sign-Symmetric Alignment: Generates weights using paired frame vectors ($+gc+gc$ and $-gc-gc$) to ensure the net activation remains zero-centered, preventing distribution shifts.
- Shape-Contoured Fusion (SCF): Instead of naive addition, the fusion process smoothly merges synthetic weights into the model's existing learned contours, minimizing destabilization.
- Controlled Adjustment: Implements dynamic scaling ($\alpha$) and clamping ($\gamma$-cap) to control the magnitude of weight updates.
- Selective Fusion: Allows users to target specific layers for modification.
⚙️ Pipeline Overview
The toolkit operates in two stages, assuming you have generated the source Neuron and derivitive Neurons which are re-usable:
1. The Generator (weights.py)
This script creates synthetic, functionally correct MLP projections (gate_proj, up_proj, down_proj) based on reference source neurons. It maps piece-wise linear behaviors into the non-linear SwiGLU block without introducing mean-shift destabilization.
- Key Function: Generates balanced, sign-symmetric delta weights.
- Output: Saves the generated delta weights to an output directory.
2. The Fuser (applyweights.py)
This script applies the safetensor delta files generated by weights.py to your existing Gemma 4 31B model using Shape-Contoured Fusion (SCF).
down_proj(Contoured Multiplicative Delta): The update is scaled by the existing weight profile and variance-normalized to ensure smooth integration. $$\text{W}{\text{down}} = \text{W}{\text{down}} + (\alpha_{\text{dynamic}} \cdot \Delta_{\text{down}} \cdot \text{W}_{\text{down}})$$gate_proj(Multiplicative Gamma Scaling): Applies a clamped fractional adjustment to prevent unbounded activation spikes. $$\text{W}{\text{gate}} = \text{W}{\text{gate}} \cdot (1 + \text{clamp}(\Delta_{\text{gate}}, -\gamma, \gamma))$$up_proj(Untouched): The linear value path is intentionally skipped during fusion.
🚀 Usage
Stage 1: Generating Delta Weights
Use weights.py to generate the synthetic delta weights.
python weights.py \
--base-model /path/to/original_gemma_model \
--output-dir generated_weights \
--seed 42
Stage 2: Applying Weights (Fusion)
Use applyweights.py to merge the generated weights into your target model.
python applyweights.py \
--model /path/to/original_gemma_model \
--weights generated_weights \
--output /path/to/fused_gemma_model \
--alpha 0.02 \
--gamma-cap 0.05
💡 Helpful Flags
| Flag | Description | Default |
|---|---|---|
--dry-run |
Simulates the fusion process, identifying missing keys, calculating partial coverage matrices, and verifying sharded indexing without writing files. | N/A |
--layers |
A space-separated list of specific layer indices to fuse (e.g., --layers 5 11 17), skipping the rest. |
None |
--alpha |
Controls the variance scale multiplier for the down_proj update. |
0.02 |
--gamma-cap |
Sets the maximum fractional adjustment allowed for the gate_proj. |
0.05 |
Note:
The default values for Alpha and Gamma are extremely conservative, they will not influence model behavior much at the defaults; it can and should, for pronounced change closer to Claude be pushed to around Alpha ~32 and Gamma-cap ~25 without breaking.