modified: __init__.py

	modified:   cmhe_torch.py
	modified:   cmhe_utilities.py

auton_survival/models/cmhe/__init__.py

@@ -109,6 +109,9 @@ class DeepCoxMixturesHeterogenousEffects:
   layers: list
       A list of integers consisting of the number of neurons in each
       hidden layer.
+  gate_l2_penalty: float
+      Strength of the l2 penalty term for the gate layers.
+      Higher means stronger regularization. 
   random_seed: int
       Controls the reproducibility of called functions.
 
@@ -122,6 +125,7 @@ class DeepCoxMixturesHeterogenousEffects:
 
   def __init__(self, k, g, layers=None, gamma=100,
                smoothing_factor=1e-4,
+               gate_l2_penalty=1e-4,
                random_seed=0):
 
     self.k = k
@@ -130,6 +134,7 @@ def __init__(self, k, g, layers=None, gamma=100,
     self.fitted = False
     self.gamma = gamma
     self.smoothing_factor = smoothing_factor
+    self.gate_l2_penalty = gate_l2_penalty
     self.random_seed = random_seed
 
   def __call__(self):
@@ -205,6 +210,7 @@ def _gen_torch_model(self, inputdim, optimizer):
                          layers=self.layers,
                          gamma=self.gamma,
                          smoothing_factor=self.smoothing_factor,
+                         gate_l2_penalty=self.gate_l2_penalty,
                          optimizer=optimizer)
 
   def fit(self, x, t, e, a, vsize=0.15, val_data=None,

auton_survival/models/cmhe/cmhe_torch.py

@@ -44,7 +44,8 @@ def _init_dcmhe_layers(self, lastdim):
     self.omega = torch.nn.Parameter(torch.rand(self.g)-0.5)
 
   def __init__(self, k, g, inputdim, layers=None, gamma=100,
-               smoothing_factor=1e-4, optimizer='Adam'):
+               smoothing_factor=1e-4, gate_l2_penalty=1e-4, 
+               optimizer='Adam'):
 
     super(DeepCMHETorch, self).__init__()
 
@@ -65,6 +66,8 @@ def __init__(self, k, g, inputdim, layers=None, gamma=100,
 
     self._init_dcmhe_layers(lastdim)
 
+    self.gate_l2_penalty = gate_l2_penalty 
+
     self.embedding = create_representation(inputdim, layers, 'Tanh')
 
 

auton_survival/models/cmhe/cmhe_utilities.py

@@ -174,6 +174,9 @@ def m_step(model, optimizer, x, t, e, a, log_likelihoods, typ='soft'):
 
   optimizer.zero_grad()
   loss = q_function(model, x, t, e, a, log_likelihoods, typ)
+  gate_regularization_loss = (model.phi_gate.weight**2).sum()
+  gate_regularization_loss += (model.z_gate.weight**2).sum()
+  loss += (model.gate_l2_penalty)*gate_regularization_loss
   loss.backward()
   optimizer.step()