Package: pytorch-cpr
Constrained Parameter RegularizationThis repository contains the PyTorch implementation ofConstrained Parameter Regularization.Installpipinstallpytorch-cprGetting startedUsage ofapply_CPROptimizer WrapperTheapply_CPRfunction is a wrapper designed to apply CPR (Constrained Parameter Regularization) to a given optimizer by first creating parameter groups and the wrapping the actual optimizer class.Argumentsmodel: The PyTorch model whose parameters are to be optimized.optimizer_cls: The class of the optimizer to be used (e.g.,torch.optim.Adam).kappa_init_param: Initial value for the kappa parameter in CPR depending on tge initialization method.kappa_init_method(default'warm_start'): The method to initialize the kappa parameter. Options include'warm_start','uniform', and'dependent'.reg_function(default'l2'): The regularization function to be applied. Options include'l2'or'std'.kappa_adapt(defaultFalse): Flag to determine if kappa should adapt during training.kappa_update(default1.0): The rate at which kappa is updated in the Lagrangian method.normalization_regularization(defaultFalse): Flag to apply regularization to normalization layers.bias_regularization(defaultFalse): Flag to apply regularization to bias parameters.embedding_regularization(defaultFalse): Flag to apply regularization to embedding parameters.**optimizer_args: Additional optimizer arguments to pass to the optimizer class.Example usageimporttorchfrompytorch-cprimportapply_CPRmodel=YourModel()optimizer=apply_CPR(model,torch.optim.Adam,kappa_init_param=1000,kappa_init_method='warm_start',lr=0.001,betas=(0.9,0.98))Run examplesWe provide scripts to replicate the experiments from our paper. Please use a system with at least 1 GPU. Install the package and the requirements for the example:python3-mvenvvenvsourcevenv/bin/activate
pipinstall-rexamples/requirements.txt
pipinstallpytorch-cprModular Addition / Grokking ExperimentThe grokking experiment should run within a few minutes. The results will be saved in thegrokkingfolder.
To replicate the results in the paper, run variations with the following arguments:For AdamW:pythonexamples/train_grokking_task.py--optimizeradamw--weight_decay0.1For Adam + Rescaling:pythonexamples/train_grokking_task.py--optimizeradamw--weight_decay0.0--rescale0.8For AdamCPR with L2 norm as regularization function:pythonexamples/train_grokking_task.py--optimizeradamcpr--kappa_init_methoddependent--kappa_init_param0.8Image Classification ExperimentThe CIFAR-100 experiment should run within 20-30 minutes. The results will be saved in thecifar100folder.For AdamW:pythonexamples/train_resnet.py--optimizeradamw--lr0.001--weight_decay0.001For Adam + Rescaling:pythonexamples/train_resnet.py--optimizeradamw--lr0.001--weight_decay0--rescale_alpha0.8For AdamCPR with L2 norm as regularization function and kappa initialization depending on the parameter initialization:pythonexamples/train_resnet.py--optimizeradamcpr--lr0.001--kappa_init_methoddependent--kappa_init_param0.8For AdamCPR with L2 norm as regularization function and kappa initialization with warm start:pythonexamples/train_resnet.py--optimizeradamcpr--lr0.001--kappa_init_methodwarm_start--kappa_init_param1000CitationPlease cite our paper if you use this code in your work:@misc{franke2023cpr,
      title={Constrained Parameter Regularization}, 
      author={Jörg K. H. Franke and Michael Hefenbrock and Gregor Koehler and Frank Hutter},
      year={2023},
      eprint={2311.09058},
      archivePrefix={arXiv},
      primaryClass={cs.LG}
}