experiment: add one-sided residual stability margin

1 files changed, 23 insertions, 3 deletions

diff --git a/sdil/conv.py b/sdil/conv.py
index 285ceb1..d9364a2 100644
--- a/sdil/conv.py
+++ b/sdil/conv.py
@@ -769,7 +769,8 @@ class CIFARKPMixedTrafficResNet(CIFARKPResNet):
         return squared_error / units
 
     @torch.no_grad()
-    def predictor_closed_form_fit(self, hiddens, min_variance=1e-12):
+    def predictor_closed_form_fit(self, hiddens, min_variance=1e-12,
+                                  stability_margin=0.0):
         """Fit the local affine neutral relation by per-cell least squares.
 
         Each spatial cell uses only its own soma and instruction-off apical
@@ -778,10 +779,15 @@ class CIFARKPMixedTrafficResNet(CIFARKPResNet):
         """
         if min_variance < 0:
             raise ValueError("minimum predictor variance must be nonnegative")
+        if stability_margin < 0:
+            raise ValueError("predictor stability margin must be nonnegative")
         traffic = self.traffic_fields(hiddens)
         residual_power = 0.0
         traffic_power = 0.0
         maximum_residual_slope = 0.0
+        maximum_positive_residual_slope = 0.0
+        minimum_residual_slope = 0.0
+        maximum_applied_margin = 0.0
         squared_error = 0.0
         units = 0
         for hidden, target, slope, bias in zip(
@@ -796,8 +802,10 @@ class CIFARKPMixedTrafficResNet(CIFARKPResNet):
             fitted_slope = torch.where(
                 active, covariance / variance.clamp_min(min_variance),
                 torch.zeros_like(variance))
-            fitted_bias = target_mean - fitted_slope * hidden_mean
-            slope.copy_(fitted_slope)
+            applied_margin = stability_margin * (1.0 + fitted_slope.abs())
+            stabilized_slope = fitted_slope + applied_margin
+            fitted_bias = target_mean - stabilized_slope * hidden_mean
+            slope.copy_(stabilized_slope)
             bias.copy_(fitted_bias)
             residual = target - slope * hidden - bias
             residual_covariance = (centered_h * (
@@ -808,6 +816,13 @@ class CIFARKPMixedTrafficResNet(CIFARKPResNet):
                 torch.zeros_like(variance))
             maximum_residual_slope = max(
                 maximum_residual_slope, float(residual_slope.abs().max()))
+            maximum_positive_residual_slope = max(
+                maximum_positive_residual_slope,
+                float(residual_slope.max().clamp_min(0.0)))
+            minimum_residual_slope = min(
+                minimum_residual_slope, float(residual_slope.min()))
+            maximum_applied_margin = max(
+                maximum_applied_margin, float(applied_margin.max()))
             power = float(residual.square().sum())
             residual_power += power
             traffic_power += float(target.square().sum())
@@ -820,6 +835,11 @@ class CIFARKPMixedTrafficResNet(CIFARKPResNet):
             "residual_traffic_rms_ratio": math.sqrt(
                 residual_power / traffic_power),
             "max_absolute_residual_soma_slope": maximum_residual_slope,
+            "max_positive_residual_soma_slope": (
+                maximum_positive_residual_slope),
+            "min_residual_soma_slope": minimum_residual_slope,
+            "max_applied_stability_margin": maximum_applied_margin,
+            "stability_margin": float(stability_margin),
             "observations": int(hiddens[0].shape[0]),
         }