diff --git a/symbolic_nn_tests/experiment2/model.py b/symbolic_nn_tests/experiment2/model.py
index 305c0ca..841bf75 100644
--- a/symbolic_nn_tests/experiment2/model.py
+++ b/symbolic_nn_tests/experiment2/model.py
@@ -9,8 +9,8 @@ class Model(nn.Module):
 
         self.return_module_y = return_module_y
 
-        self.x0_encoder = nn.TransformerEncoderLayer(7, 7)
-        self.x1_encoder = nn.TransformerEncoderLayer(10, 10)
+        self.x0_encoder = nn.TransformerEncoderLayer(7, 7, dim_feedforward=512)
+        self.x1_encoder = nn.TransformerEncoderLayer(10, 10, dim_feedforward=1024)
         self.encode_x0 = self.create_xval_encoding_fn(self.x0_encoder)
         self.encode_x1 = self.create_xval_encoding_fn(self.x1_encoder)
         self.ff = nn.Sequential(
diff --git a/symbolic_nn_tests/experiment2/semantic_loss.py b/symbolic_nn_tests/experiment2/semantic_loss.py
index 6e8369d..3717ae5 100644
--- a/symbolic_nn_tests/experiment2/semantic_loss.py
+++ b/symbolic_nn_tests/experiment2/semantic_loss.py
@@ -1,4 +1,4 @@
-from symbolic_nn_tests.experiment2.math import linear_fit, linear_residuals, sech
+from symbolic_nn_tests.experiment2.math import linear_fit, linear_residuals
 from random import random
 from torch import nn
 import torch
@@ -22,7 +22,7 @@ import torch
 class PositiveSlopeLinearLoss(nn.Module):
     def __init__(self, wandb_logger=None, version="", device="cuda", log_freq=50):
         super().__init__()
-        self.params = [random()]
+        self.params = [random(), random(), random()]
         self.wandb_logger = wandb_logger
         self.version = version
         self.device = device
@@ -54,17 +54,15 @@ class PositiveSlopeLinearLoss(nn.Module):
         m, c = linear_fit(diff_electronegativity, y_pred)
 
         # To start with, we want to calculate a penalty based on deviation from a linear relationship
-        residual_penalty = (
-            (1 / sech(linear_residuals(diff_electronegativity, y_pred, m, c)))
-            .abs()
-            .float()
-            .mean()
-        )
+        r = linear_residuals(diff_electronegativity, y_pred, m, c)
+        residual_penalty = ((self.params[0] * r * r) + 1).float().mean()
 
         # We also need to calculate a penalty that incentivizes a positive slope. For this, im using relu
         # to scale the slope as it will penalise negative slopes without just creating a reward hack for
         # maximizing slope.
-        slope_penalty = (nn.functional.relu(self.params[0] * (-m)) + 1).mean()
+        slope_penalty = (
+            nn.functional.softplus(self.params[1] * (-m), beta=self.params[2]) + 1
+        ).mean()
 
         if self.wandb_logger and (self.steps_since_log >= self.log_freq):
             self.wandb_logger.log_metrics({f"{self.version}-a": self.params})
diff --git a/symbolic_nn_tests/experiment2/train.py b/symbolic_nn_tests/experiment2/train.py
index 2f5d150..85c05ed 100644
--- a/symbolic_nn_tests/experiment2/train.py
+++ b/symbolic_nn_tests/experiment2/train.py
@@ -8,7 +8,7 @@ class TrainingWrapper(_TrainingWrapper):
     def __init__(self, *args, loss_rate_target=-10, **kwargs):
         super().__init__(*args, **kwargs)
         self.loss_optimizer = Optimizer(
-            [(0.0, 1000.0)],
+            [(0.0, 512.0), (0.0, 1024.0), (0.0, 512.0)],
             base_estimator=RandomForestRegressor(
                 n_jobs=-1,
             ),