[Paper Review] Linearity Grafting: Relaxed Neuron Pruning Helps Certifiable Robustness

Tianlong Chen, Huan Zhang+5 more|arXiv (Cornell University)|Jun 15, 2022

Adversarial Robustness in Machine Learning4 citations

TL;DR

This paper introduces linearity grafting, a relaxed neuron pruning method that replaces insignificant and unstable ReLU neurons with learnable linear activations to enhance certifiable robustness and verification scalability. By optimizing slope and intercept parameters, it achieves up to 82.30% higher certified accuracy than baseline models without certified robust training, while enabling complete verification on large models with up to 17M parameters.

ABSTRACT

Certifiable robustness is a highly desirable property for adopting deep neural networks (DNNs) in safety-critical scenarios, but often demands tedious computations to establish. The main hurdle lies in the massive amount of non-linearity in large DNNs. To trade off the DNN expressiveness (which calls for more non-linearity) and robustness certification scalability (which prefers more linearity), we propose a novel solution to strategically manipulate neurons, by "grafting" appropriate levels of linearity. The core of our proposal is to first linearize insignificant ReLU neurons, to eliminate the non-linear components that are both redundant for DNN performance and harmful to its certification. We then optimize the associated slopes and intercepts of the replaced linear activations for restoring model performance while maintaining certifiability. Hence, typical neuron pruning could be viewed as a special case of grafting a linear function of the fixed zero slopes and intercept, that might overly restrict the network flexibility and sacrifice its performance. Extensive experiments on multiple datasets and network backbones show that our linearity grafting can (1) effectively tighten certified bounds; (2) achieve competitive certifiable robustness without certified robust training (i.e., over 30% improvements on CIFAR-10 models); and (3) scale up complete verification to large adversarially trained models with 17M parameters. Codes are available at https://github.com/VITA-Group/Linearity-Grafting.

Motivation & Objective

• To address the scalability bottleneck in neural network verification caused by excessive non-linearity in deep ReLU networks.
• To improve certifiable robustness without relying on time-consuming certified robust training.
• To relax traditional neuron pruning by introducing tunable linear activations instead of zeroing out neurons.
• To enable complete verification on large-scale, adversarially trained models by reducing the number of unstable ReLU neurons.
• To maintain high model performance through optimized linear parameters, avoiding performance degradation from over-pruning.

Proposed method

• Identify insignificant and unstable ReLU neurons using activation gradient (GAP) as a significance proxy.
• Replace selected ReLU neurons with learnable linear functions of the form $ a \times x + b $, where $ a $ and $ b $ are optimized during fine-tuning.
• Apply a one-shot or gradual grafting schedule to progressively replace neurons, with hyperparameters tuned via standard training.
• Use linear relaxations for verification, which are computationally efficient and tighter when non-linearity is reduced.
• Integrate grafting into pre-trained adversarially trained models to improve certifiable robustness without retraining from scratch.
• Optimize both slope and intercept to preserve model expressiveness and avoid inactive neuron saturation.

Experimental results

Research questions

• RQ1 Can replacing redundant ReLU neurons with learnable linear functions improve certifiable robustness without certified training?
• RQ2 Does reducing non-linearity via linearity grafting tighten verification bounds and improve scalability?
• RQ3 How does learnable linear grafting compare to classical neuron pruning (i.e., zeroing out neurons) in terms of performance and robustness?
• RQ4 Can linearity grafting enable complete verification on large-scale, adversarially trained models?
• RQ5 What is the trade-off between standard accuracy, verified accuracy, and training efficiency when applying grafting?

Key findings

• Linearity grafting improved certified accuracy by up to 82.30% on CIFAR-10 models compared to baseline models without certified robust training.
• The method achieved 32.30% certified accuracy at $ \epsilon = 8/255 $ on a 17M-parameter model without any certified defense training.
• Grafting with learnable parameters ($ a \times x + b $) outperformed grafting with zeroed-out neurons (i.e., standard pruning), improving verified accuracy by 2.96% and standard accuracy by 1.63%.
• The approach reduced training time by 92.87% compared to state-of-the-art certified robust training, while achieving higher verified and standard accuracy.
• Gradual grafting preserved slightly higher standard accuracy but slightly lower verified accuracy than one-shot grafting, indicating a trade-off in stability.
• The method enabled complete verification on large models previously intractable due to high non-linearity, demonstrating improved scalability.