Initialize flows to near identity

davidnabergoj · Aug 10, 2024 · e82627e · e82627e
1 parent 03fe60f
commit e82627e
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/normalizing_flows/bijections/finite/autoregressive/conditioning/transforms.py b/normalizing_flows/bijections/finite/autoregressive/conditioning/transforms.py
@@ -234,7 +234,7 @@ def __init__(self,
         self.sequential = nn.Sequential(*layers)
 
     def predict_theta_flat(self, x: torch.Tensor, context: torch.Tensor = None):
-        return self.sequential(self.context_combiner(x, context))
+        return self.sequential(self.context_combiner(x, context)) / 1000.0
 
 
 class Linear(FeedForward):
@@ -257,7 +257,7 @@ def __init__(self, event_size: int, hidden_size: int, block_size: int, nonlinear
             self.sequential = nn.Sequential(*layers)
 
         def forward(self, x):
-            return x + self.sequential(x)
+            return x + self.sequential(x) / 1000.0
 
     def __init__(self,
                  input_event_shape: torch.Size,
@@ -289,7 +289,7 @@ def __init__(self,
         self.sequential = nn.Sequential(*layers)
 
     def predict_theta_flat(self, x: torch.Tensor, context: torch.Tensor = None):
-        return self.sequential(self.context_combiner(x, context))
+        return self.sequential(self.context_combiner(x, context)) / 1000.0
 
 
 class CombinedConditioner(nn.Module):

diff --git a/normalizing_flows/bijections/finite/autoregressive/layers_base.py b/normalizing_flows/bijections/finite/autoregressive/layers_base.py
@@ -180,7 +180,7 @@ def __init__(self, transformer: ScalarTransformer, fill_value: float = None):
         )
 
         if fill_value is None:
-            self.value = nn.Parameter(torch.randn(*transformer.parameter_shape))
+            self.value = nn.Parameter(torch.randn(*transformer.parameter_shape)) / 1000.0
         else:
             self.value = nn.Parameter(torch.full(size=transformer.parameter_shape, fill_value=fill_value))