Package: jaxDiversity
jaxDiversityThis is an updated implementation for Neural networks embrace diversity
paperAuthorsAnshul Choudhary, Anil Radhakrishnan, John F. Lindner, Sudeshna Sinha,
and William L. DittoLink to paperarXivKey ResultsWe construct neural networks with learnable activation functions and
sere that they quickly diversify from each other under training.These activations subsequently outperform theirpurecounterparts on
classification tasks.The neuronal sub-networks instantiate the neurons and meta-learning
adjusts their weights and biases to find efficient spanning sets of
nonlinear activations.These improved neural networks provide quantitative examples of the
emergence of diversity and insight into its advantages.InstallpipinstalljaxDiversityHow to useThe codebase has 4 main components: * dataloading: Contains tools for
loading the datasets mentioned in the manuscript. We use pytorch
dataloaders with a custom numpy collate function to use this data in
jax.losses: We handle both traditional mlps and hamiltonian neural
networkss with minimal changes with our loss implementations.mlp: Contains custom mlp that takes in multiple activations and uses
themintralayerto create a diverse network. Also contains the
activation neural networks.loops: Contains the inner and outer loops for metalearning to optimize
the activation functions in tandem with the supervised learning taskMinimum exampleimportjaximportoptaxfromjaxDiversity.utilclassesimportInnerConfig,OuterConfig# simple utility classes for configuration consistencyfromjaxDiversity.dataloadingimportNumpyLoader,DummyDatasetfromjaxDiversity.mlpimportmlp_afunc,MultiActMLP,init_linear_weight,xavier_normal_init,savefromjaxDiversity.baselineimportcompute_lossascompute_loss_baselinefromjaxDiversity.hnnimportcompute_lossascompute_loss_hnnfromjaxDiversity.loopsimportinner_opt,outer_optinner optimzation or standard training loop with the baseline activationdev_inner_config=InnerConfig(test_train_split=0.8,input_dim=2,output_dim=2,hidden_layer_sizes=[18],batch_size=64,epochs=2,lr=1e-3,mu=0.9,n_fns=2,l2_reg=1e-1,seed=42)key=jax.random.PRNGKey(dev_inner_config.seed)model_key,init_key=jax.random.split(key)afuncs=[lambdax:x**2,lambdax:x]train_dataset=DummyDataset(1000,dev_inner_config.input_dim,dev_inner_config.output_dim)test_dataset=DummyDataset(1000,dev_inner_config.input_dim,dev_inner_config.output_dim)train_dataloader=NumpyLoader(train_dataset,batch_size=dev_inner_config.batch_size,shuffle=True)test_dataloader=NumpyLoader(test_dataset,batch_size=dev_inner_config.batch_size,shuffle=True)opt=optax.rmsprop(learning_rate=dev_inner_config.lr,momentum=dev_inner_config.mu,decay=dev_inner_config.l2_reg)model=MultiActMLP(dev_inner_config.input_dim,dev_inner_config.output_dim,dev_inner_config.hidden_layer_sizes,model_key,bias=False)baselineNN,opt_state,inner_results=inner_opt(model=model,train_data=train_dataloader,test_data=test_dataloader,afuncs=afuncs,opt=opt,loss_fn=compute_loss_baseline,config=dev_inner_config,training=True,verbose=True)metalearning with Hamiltonian Neural Networksinner_config=InnerConfig(test_train_split=0.8,input_dim=2,output_dim=1,hidden_layer_sizes=[32],batch_size=64,epochs=5,lr=1e-3,mu=0.9,n_fns=2,l2_reg=1e-1,seed=42)outer_config=OuterConfig(input_dim=1,output_dim=1,hidden_layer_sizes=[18],batch_size=1,steps=2,print_every=1,lr=1e-3,mu=0.9,seed=24)train_dataset=DummyDataset(1000,inner_config.input_dim,2)test_dataset=DummyDataset(1000,inner_config.input_dim,2)train_dataloader=NumpyLoader(train_dataset,batch_size=inner_config.batch_size,shuffle=True)test_dataloader=NumpyLoader(test_dataset,batch_size=inner_config.batch_size,shuffle=True)opt=optax.rmsprop(learning_rate=inner_config.lr,momentum=inner_config.mu,decay=inner_config.l2_reg)meta_opt=optax.rmsprop(learning_rate=outer_config.lr,momentum=outer_config.mu)HNN_acts,HNN_stats=outer_opt(train_dataloader,test_dataloader,compute_loss_hnn,inner_config,outer_config,opt,meta_opt,save_path=None)Link to older pytorch codebase with classification problem:DiversityNN