Package: hat-cl
HAT-CLRedesigned Hard-Attention-to-the-Task for Continual LearningHAT-CL is a comprehensive reimagining of the Hard-Attention-to-the-Task (HAT) mechanism, designed specifically to combat catastrophic forgetting during Continual Learning (CL).
Originally proposed in the paperOvercoming catastrophic forgetting with hard attention to the task, HAT has been instrumental in enabling neural networks to learn successive tasks without erasure of prior knowledge.
However, the original implementation had its drawbacks, notably incompatibility with PyTorch's optimizers and the requirement for manual gradient manipulation.
HAT-CL aims to rectify these issues with a user-friendly design and a host of new features:Seamless compatibility with all PyTorch operations and optimizers.Automated gradient manipulation through PyTorch hooks.Simple transformation of PyTorch modules to HAT modules with a single line of code.Out-of-the-box HAT networks integrated withtimm.Quick StartInstallationTo install via pip:pipinstallhat-clOr, if you are using poetry:poetryaddhat-clBasic UsageTo use HAT modules, simply swap the generic PyTorch modules for their HAT counterparts.
For example, replacetorch.nn.Linearwithhat.modules.HATLinear, andtorch.nn.Conv2dwithhat.modules.HATConv2d.
The HAT modules take instances ofhat.HATPayloadas input and output, wrapping the tensor, task ID, and other essential variables for the HAT mechanism.Below is a basic example of a 2-layer MLP:importtorchimporttorch.nnasnnfromhatimportHATPayload,HATConfigfromhat.modulesimportHATLinearhat_config=HATConfig(num_tasks=5)classMLP(nn.Module):def__init__(self,input_dim,hidden_dim,output_dim):super().__init__()self.linear1=HATLinear(input_dim,hidden_dim,hat_config)self.relu=nn.ReLU()self.linear2=HATLinear(hidden_dim,output_dim,hat_config)defforward(self,x:HATPayload):x=self.linear1(x)# You can still pass the payload to the non-HAT modules like thisx=x.forward_by(self.relu)x=self.linear2(x)returnxmlp=MLP(input_dim=128,hidden_dim=32,output_dim=2)input_payload=HATPayload(torch.rand(10,128),task_id=0,mask_scale=10.0)output_payload=mlp(input_payload)output_data=output_payload.dataWith these steps, you've created a 2-layer MLP with the HAT mechanism and successfully conducted a forward pass through the model.
Just like any other PyTorch modules, it's ready to be trained, evaluated, and more—all under-the-hood operations are handled by the HAT modules.Additionally, HAT-CL provides ready-to-use HAT networks with timm integration. Creating a HAT model is as simple as creating any othertimmmodel:importtimmimporthat.timm_models# This line is necessary to register the HAT models to timmfromhatimportHATConfighat_config=HATConfig(num_tasks=5)hat_resnet18=timm.create_model('hat_resnet18',hat_config=hat_config)ModulesHere's a handy table of PyTorch modules and their HAT counterparts:PyTorch moduleHAT moduletorch.nn.Linearhat.modules.HATLineartorch.nn.Conv1dhat.modules.HATConv1dtorch.nn.Conv2dhat.modules.HATConv2dtorch.nn.Conv3dhat.modules.HATConv3dtorch.nn.BatchNorm1dhat.modules.TaskIndexedBatchNorm1dtorch.nn.BatchNorm2dhat.modules.TaskIndexedBatchNorm2dtorch.nn.BatchNorm3dhat.modules.TaskIndexedBatchNorm3dtorch.nn.LayerNormhat.modules.TaskIndexedLayerNormNetworksHere are the currently available timm-compatible HAT networks:HAT Network NameHas pretrained weightsDescriptionhat_resnet18NoHAT ResNet-18hat_resnet18sNoHAT ResNet-18 for smaller imageshat_resnet34NoHAT ResNet-34hat_resnet34sNoHAT ResNet-34 for smaller imageshat_vit_tiny_patch16_224YesHAT ViT-Tiny (16, 224)ExamplesContinual LearningFeature ImportanceLimitationsWhile HAT-CL is designed to be compatible with all PyTorch features, there are some inherent limitations due to the nature of the HAT mechanism:Optimizer Re-initialization: To prevent carryover of momentum from previous tasks, it is recommended to refresh the optimizer state after each task. This can be achieved by simply re-initializing the optimizer.Weight Decay (L2 Regularization): Weight decay is not compatible with the HAT mechanism due to its gradient modification process. As this occurs after the HAT mechanism modifies the gradient, it may interfere with parameters that should be locked out by the HAT mechanism, leading to potential forgetting. This includes theweight_decayparameter in the optimizer, or any optimizer that utilizes weight decay (e.g. AdamW).TODOAdd example notebook for pruningPackage paper for implementation detailsAuthorsXiaotian Duan (xduan7 at gmail dot com)