Make MLP module residual

cmd/conf/score_model/mlp.yaml

@@ -1,8 +1,9 @@
 _target_: fdiff.models.score_models.MLPScoreModule
 _partial_: true
-d_model: 512
+d_model: 128
+d_mlp: 1024
 num_layers: 10
-lr_max: 1.0e-3
+lr_max: 1.0e-4
 fourier_noise_scaling: ${fourier_transform}
 likelihood_weighting: False
 

src/fdiff/models/score_models.py

@@ -190,7 +190,8 @@ def __init__(
         max_len: int,
         noise_scheduler: DDPMScheduler | SDE,
         fourier_noise_scaling: bool = True,
-        d_model: int = 60,
+        d_model: int = 72,
+        d_mlp: int = 512,
         num_layers: int = 3,
         num_training_steps: int = 1000,
         lr_max: float = 1e-3,
@@ -216,9 +217,12 @@ def __init__(
         self.unembedder = nn.Linear(
             in_features=d_model, out_features=max_len * n_channels
         )
-        self.backbone = MLP(
-            in_channels=d_model,
-            hidden_channels=[d_model] * num_layers,
+
+        self.backbone = nn.ModuleList(
+            [
+                MLP(in_channels=d_model, hidden_channels=[d_mlp, d_model], dropout=0.1)
+                for _ in range(num_layers)
+            ]
         )
         self.pos_encoder = None
 
@@ -245,7 +249,8 @@ def forward(self, batch: DiffusableBatch) -> torch.Tensor:
         X = self.time_encoder(X, timesteps, use_time_axis=False)
 
         # Backbone
-        X = self.backbone(X)
+        for layer in self.backbone:
+            X = X + layer(X)
 
         # Channel unembedding
         X = self.unembedder(X)