Criteria

Which parameter is important in a neural network ?

The criteria implemented come from this paper.

---

source

Reducer

def Reducer(
    *args, **kwargs
):

---

source

Normalizer

def Normalizer(
    *args, **kwargs
):

---

source

Criteria

def Criteria(
    f:Callable[[torch.Tensor], torch.Tensor], # Transform applied to the weights
    reducer:Callable=mean, normalizer:Callable | None=None, needs_init:bool=False, needs_update:bool=False,
    output_fn:Callable[[torch.Tensor, torch.Tensor], torch.Tensor] | None=None, # Combines current and reference weights
    return_init:bool=False, scale:dict[nn.Module, torch.Tensor] | None=None, needs_data:bool=False,
    data_fn:str='l2_norm', # Activation statistic to collect: 'l2_norm', 'max', 'mean'
):

Evaluates neural network parameters based on various criteria for pruning

Magnitude Based Criteria

---

source

magnitude_criteria

def magnitude_criteria(
    transform_fn, **kwargs
):

Create a criteria based on weight magnitude transformation.

Random

demo_model(random)

Large Final Value

demo_model(large_final)

Squared Final Value

demo_model(squared_final)

Small Final Value

demo_model(small_final)

Init based criteria

---

source

init_based_criteria

def init_based_criteria(
    transform_fn, output_fn:NoneType=None, return_init:bool=False, **kwargs
):

Create a criteria that compares current weights to initial weights.

Large Init Value

demo_model(large_init)

Small Init Value

demo_model(small_init)

Large Init Large Final Value

demo_model(large_init_large_final, 0.8)

Small Init Small Final Value

demo_model(small_init_small_final)

Increasing Magnitude

demo_model(magnitude_increase, 0.6)

Movement Pruning

demo_model(movement)

movmag = init_based_criteria(noop, output_fn=lambda x,y: torch.abs(torch.mul(x, torch.sub(x,y))))

demo_model(movmag)

Update based criteria

---

source

update_based_criteria

def update_based_criteria(
    transform_fn, output_fn:NoneType=None, **kwargs
):

Create a criteria that compares current weights to previous iteration weights.

The following criteria use an updating value of the weights, i.e. the value from the previous iteration of training, instead of the initialization value to better capture the training dynamics.

Updating Magnitude Increase

demo_model(updating_magnitude_increase)

Updating Movement

demo_model(updating_movement, 0.5)

Updating mov-magnitude

demo_model(updating_movmag)

---

source

available_criterias

def available_criterias():

Return the list of available criteria names

Activation-Based Criteria

The following criteria use input activation statistics collected during a calibration pass, providing data-aware importance scoring.

Wanda

Wanda (Sun et al. ICLR 2024) scores weight importance as |W| × ‖X‖₂ — the product of weight magnitude and input activation L2 norm. Best for one-shot post-training sparsification. Requires calibration data passed via Sparsifier(data=...).

---

source

activation_criteria

def activation_criteria(
    transform_fn, data_fn:str='l2_norm', **kwargs
):

Create a criteria that uses activation statistics to weight scores

New Ideas

updating_magnitude_increase = Criteria(torch.abs, needs_update=True, output_fn= lambda x,y: torch.abs(torch.sub(x,y)))

demo_model(updating_magnitude_increase)

updating_magnitude_increase = Criteria(torch.abs, needs_update=True, output_fn= lambda x,y: torch.sub(x,y))

demo_model(updating_magnitude_increase)

updating_magnitude_increase = Criteria(torch.square, needs_update=True, output_fn= lambda x,y: torch.abs(torch.sub(x,y)))

demo_model(updating_magnitude_increase)

updating_movmag = Criteria(noop, needs_update=True, output_fn=lambda x,y: torch.abs(torch.mul(x, torch.sub(x,y))))
demo_model(updating_movmag)

updating_movmag = Criteria(noop, needs_update=True, output_fn=lambda x,y: torch.abs(torch.mul(torch.square(x), torch.sub(x,y))))
demo_model(updating_movmag)

updating_movmag = Criteria(torch.square, needs_update=True, output_fn=lambda x,y: torch.abs(torch.mul(x, torch.sub(x,y))))
#updating_movmag = Criteria(noop, needs_update=True, output_fn=lambda x,y: torch.mul(x, torch.sub(x,y)))
demo_model(updating_movmag)

updating_movmag = Criteria(torch.abs, needs_update=True, output_fn=lambda x,y: torch.abs(torch.mul(x, torch.sub(x,y))))
#updating_movmag = Criteria(noop, needs_update=True, output_fn=lambda x,y: torch.mul(x, torch.sub(x,y)))
demo_model(updating_movmag, 0.3)

updating_movmag = Criteria(torch.abs, needs_update=True, output_fn=lambda x,y: torch.mul(x, torch.sub(x,y)))

demo_model(updating_movmag, 0.8)

updating_movmag = Criteria(torch.square, needs_update=True, output_fn=lambda x,y: torch.mul(x, torch.sub(x,y)))

demo_model(updating_movmag)

updating_movmag = Criteria(noop, needs_update=True, output_fn=lambda x,y: torch.mul(x, torch.sub(x,y)))

demo_model(updating_movmag)

updating_movement = Criteria(noop, needs_update=True, output_fn= lambda x,y: torch.abs(torch.sub(-x,y)))
demo_model(updating_movement, 0.5)

updating_movement = Criteria(torch.abs, needs_update=True, output_fn= lambda x,y: torch.abs(torch.sub(-x,y)))
demo_model(updating_movement)

updating_movement = Criteria(torch.abs, needs_update=True, output_fn= lambda x,y: torch.abs(torch.cosh(torch.sub(x,y))))
demo_model(updating_movement)

updating_movement = Criteria(torch.square, needs_update=True, output_fn= lambda x,y: torch.abs(torch.sub(x,y)))
demo_model(updating_movement)

updating_movement = Criteria(noop, needs_update=True, output_fn= lambda x,y: torch.sub(x,y))
demo_model(updating_movement)

mine = partial(torch.pow, exponent=4)

large_final = Criteria(torch.frac)
demo_model(large_final)

First order Taylor expansion on the weight (as per Nvidia Taylor Pruning)

---

source

grad_crit

def grad_crit(
    m:nn.Module, g:str, # granularity specification
)->torch.Tensor:

First order Taylor expansion criterion for weight importance (Nvidia Taylor Pruning)

scores = torch.randn(100).abs()
normed = Normalizer.standardization(scores)

---

See Also

• Sparsifier - Apply sparsification using these criteria
• Pruner - Structured pruning with importance scoring
• Granularity - Control what gets pruned (weights, filters, etc.)
• Parametrize - The master weight these criteria score when something computes the weight

Criteria Selection Guide

Criteria	Data-Aware	Best For	Requires
large_final	No	General-purpose magnitude pruning	Nothing
wanda	Yes	Post-training one-shot pruning	Calibration data
movement	No	During-training pruning	Initial weights
grad_crit	No	Gradient-informed pruning	Gradients
random	No	Baseline comparison	Nothing

Tests live in nbs/tests/test_criteria.ipynb.