VexelNet

Official implementation for paper Vexel-Net: Vectors-Composed Features Oriented Neural Network.

---

Abatract

Extracting features that is robust to geometric transformations and adversarial perturbations remains a critical challenge in deep learning research and applications. Among existing methods, \textbf{Capsule}-based approaches address this challenge through vectorized feature representations and dynamic routing mechanisms, demonstrating notable resilience against affine transformations and adversarial attacks. However, these methods exhibit two limitations: 1) prolonged training periods caused by complex iterative dynamic routing processes, and 2) substantial computational overhead resulting from dense matrix transformations and auxiliary activation scalars. In this research, we propose a novel framework to addresses these limitations through three key innovations. First, we design compact vector-grouped neurons which eliminate the need for individual activation scalars. Second, we introduce an efficient non-iterative voting algorithm that reduces routing complexity while maintaining spatial relationships. Third, we implement adaptive weight-sharing strategies to optimize the trade-off between model's representational capacity and the computational efficiency. Comprehensive experimental evaluations across image classification, affine transformations test, and adversarial attacks show that our method achieves superior classification accuracy, enhanced robustness performance with a lighter computational burden compared to existing capsule netowrks.

Acknowledgement

Our codes are implemented benefiting from SR-CapsNet, we recommend that interested readers will gain a lot after reading their paper Self-Routing Capsule Networks and looking into their repository.

Requirement

pytorch >= 2.0.1+cu118
pytorch lightning >= 2.1.0
pytorch vision >= 0.15.2+cu118
einops >= 0.6.1
numpy >= 1.23.5

Other python libraries need for this repository will be installed when installing pytorch lightning. If met problem with libraries for this repository, you are welcome to submit issues.

About args

Args are defined in utils/configUtils.py , if users want to specify different values for args, we suggest modifying args in utils/configUtils.py, just cleanly run main.py. And other arguments for models and optimization are defined in main.py as python dictoinaries.

Image Classification

The configurations and parameters of networks for image classifications are already recorded in our paper, please refer to it if interested.

Affine Robustness

Just modify args for viewpoints by assign values for --exp and famaliar, look at data/data.py, data/utils.py for how these args work.

Adversarial Attack

Specify values for args --typeAttack, --epsAttack and --targeted in utils/configUtils/py in order to generate adversarial samples and run adversarial_attack.py for adversarial attack test. Interested readers can refer to adversarial_attack.py for loading checkpoints for LightningModule.

Ablation

All ablation experiments are conducted by changing args in main.py and changing network architecture in models/networks/VexNet.py. Please pay attention to channels, kernel_size, and padding of Vexel conv layers when modifying network architecture. We listed used configurations as follows.

Vex-1

	kernel_size	stride	padding	in_channels	out_channels
VexConv1	input spatial size	1	0	initVexChannels	numClasses

Vex-2

	kernel_size	stride	padding	in_channels	out_channels
VexConv1	5	2	2	initVexChannels	initVexChannels//2
VexConv2	input spatial size	1	0	initVexChannels //2	numClasses

As for ablation on initial Vexel channels and dimension, just refer to main.py and modify corresponding values in argsNet.

Logs and Checkpoints

We adopt Pytorch and Pytorch Lightning for implementation our methods as well as other networks, please refer to doc of load_from_checkpoint for loading trained LightningMotule checkpoints.

Logs for approaches in our experiments will be uploaded to google drive folders afterwards.