Fix float32 MLE optimization with autograd gradients

netcal/scaling/AbstractLogisticRegression.py

@@ -565,17 +565,32 @@ def convex(self, data: torch.Tensor, y: torch.Tensor, tensorboard: bool, log_dir
         # compute NLL loss - fix weights given of the model for the current iteration step
         def MLE(w, x, y):
 
-            data = {}
+            conditioned_values = {}
+            parameters = []
             start = 0
             for name, site in self._sites.items():
                 num_weights = len(site['init']['mean'])
-                data[name] = torch.from_numpy(w[start:start+num_weights]).to(
+                parameter = torch.tensor(
+                    w[start:start+num_weights],
                     dtype=dtype,
                     device=self._device,
+                    requires_grad=True,
                 )
+                conditioned_values[name] = parameter
+                parameters.append(parameter)
                 start += num_weights
 
-            return loss_op(torch.squeeze(pyro.condition(self.model, data=data)(x)), y).item()
+            loss = loss_op(
+                torch.squeeze(pyro.condition(self.model, data=conditioned_values)(x)),
+                y,
+            )
+            loss.backward()
+
+            gradient = np.concatenate([
+                parameter.grad.detach().cpu().numpy().astype(np.float64)
+                for parameter in parameters
+            ])
+            return loss.item(), gradient
 
         initial_weights = np.concatenate(
             [site['init']['mean'].cpu().numpy() for site in self._sites.values()]
@@ -600,7 +615,7 @@ def loss_op(x, y):
         optim_bounds = self._get_scipy_constraints()
 
         # invoke SciPy's optimization function as this is very light-weight and fast
-        result = minimize(fun=MLE, x0=initial_weights, args=(data, y), bounds=optim_bounds)
+        result = minimize(fun=MLE, x0=initial_weights, args=(data, y), bounds=optim_bounds, jac=True)
 
         # assign weights to according sites
         start = 0
@@ -618,7 +633,7 @@ def loss_op(x, y):
             if bounds:
                 # rerun minimization routine
                 initial_weights[masked_weights] = 0.0
-                result = minimize(fun=MLE, x0=initial_weights, args=(data, y), bounds=bounds)
+                result = minimize(fun=MLE, x0=initial_weights, args=(data, y), bounds=bounds, jac=True)
 
         # get intercept and weights after optimization
         start = 0

tests/test_mle_precision.py

@@ -0,0 +1,56 @@
+import unittest
+
+import numpy as np
+
+from netcal.scaling import LogisticCalibration, TemperatureScaling
+
+
+def softmax(logits: np.ndarray) -> np.ndarray:
+    logits = logits - np.max(logits, axis=1, keepdims=True)
+    probabilities = np.exp(logits)
+    return probabilities / np.sum(probabilities, axis=1, keepdims=True)
+
+
+def negative_log_likelihood(probabilities: np.ndarray, labels: np.ndarray) -> float:
+    selected = probabilities[np.arange(len(labels)), labels]
+    return float(-np.mean(np.log(np.clip(selected, 1e-15, 1.0))))
+
+
+class TestMLEPrecision(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        rng = np.random.default_rng(42)
+        num_samples, num_classes = 2000, 5
+        latent_logits = rng.normal(size=(num_samples, num_classes))
+        true_probabilities = softmax(latent_logits)
+        cls.labels = np.array([
+            rng.choice(num_classes, p=probabilities)
+            for probabilities in true_probabilities
+        ])
+
+        # Make the predictions deliberately overconfident. Both scaling methods
+        # should move away from their identity initialization and lower NLL.
+        cls.predictions = softmax(3.0 * latent_logits)
+
+    def test_mle_optimizes_float32_and_float64_inputs(self):
+        for calibrator_type in (TemperatureScaling, LogisticCalibration):
+            losses = {}
+            for dtype in (np.float32, np.float64):
+                with self.subTest(calibrator=calibrator_type.__name__, dtype=dtype.__name__):
+                    predictions = self.predictions.astype(dtype)
+                    before = negative_log_likelihood(predictions, self.labels)
+
+                    calibrator = calibrator_type(method="mle", use_cuda=False)
+                    calibrator.fit(predictions, self.labels, random_state=42, tensorboard=False)
+                    calibrated = calibrator.transform(predictions)
+                    after = negative_log_likelihood(calibrated, self.labels)
+                    losses[dtype] = after
+
+                    self.assertLess(after, before - 0.2)
+                    self.assertGreater(np.max(np.abs(calibrated - predictions)), 1e-3)
+
+            self.assertAlmostEqual(losses[np.float32], losses[np.float64], places=4)
+
+
+if __name__ == "__main__":
+    unittest.main()