refactor(losses): compute the regularization loss over regularized layers only

aether/losses/categorical_crossentropy.py

@@ -59,7 +59,14 @@ def get_config(self) -> dict:
         Used by Model.save() / Model.load()."""
         return {}
     def remember_trainable_layers(self, trainable_layers):
-        self.trainable_layers = trainable_layers
+        # Regularizers are fixed at construction, so filter once here, not every step.
+        self.regularized_layers = [
+            layer for layer in trainable_layers
+            if getattr(layer, "weight_regularizer_l1", 0.0) > 0
+            or getattr(layer, "weight_regularizer_l2", 0.0) > 0
+            or getattr(layer, "bias_regularizer_l1", 0.0) > 0
+            or getattr(layer, "bias_regularizer_l2", 0.0) > 0
+        ]
 
     def calculate(self, output, y, *, include_regularization= False, training = True):
         xp = config.get_array_module(output)
@@ -85,24 +92,23 @@ def calculate_accumulated(self, *, include_regularization = False):
     def regularization_loss(self):
         regularization_loss = 0             # if we don't do this, we risk overfitting.
                                             # We will have to denote partials for this too...
-        for layer in getattr(self, "trainable_layers", []):        
+        # Params are always fp32 (config.PARAM_DTYPE), so no cast is needed.
+        for layer in getattr(self, "regularized_layers", []):
             xp = config.get_array_module(layer.weights)
-
-            weights_fp32 = layer.weights.astype(xp.float32, copy=False)
-            biases_fp32 = layer.biases.astype(xp.float32, copy=False)
+            weights, biases = layer.weights, layer.biases
 
             if layer.weight_regularizer_l1 > 0:
                 regularization_loss += layer.weight_regularizer_l1 * \
-                                        xp.sum(xp.abs(weights_fp32))
+                                        xp.sum(xp.abs(weights))
             if layer.weight_regularizer_l2 > 0:
                 regularization_loss += 0.5 * layer.weight_regularizer_l2 * \
-                                        xp.sum(weights_fp32 * weights_fp32)
+                                        xp.sum(weights * weights)
             if layer.bias_regularizer_l1 > 0:
                 regularization_loss += layer.bias_regularizer_l1 * \
-                                        xp.sum(xp.abs(biases_fp32))
+                                        xp.sum(xp.abs(biases))
             if layer.bias_regularizer_l2 > 0:
                 regularization_loss += 0.5 * layer.bias_regularizer_l2 * \
-                                        xp.sum(biases_fp32 * biases_fp32) 
+                                        xp.sum(biases * biases)
         return regularization_loss
 
 

aether/model.py

@@ -786,13 +786,7 @@ def get_batch(epoch_indices, start_idx, end_idx):
         step_optimizer = self._step_optimizer
         advance_rng = self._rng_clock.advance
         shuffle_rng = self._shuffle_rng
-        has_reg = hasattr(loss, "regularization_loss") and any(
-            getattr(layer, "weight_regularizer_l1", 0.0) > 0.0
-            or getattr(layer, "weight_regularizer_l2", 0.0) > 0.0
-            or getattr(layer, "bias_regularizer_l1", 0.0) > 0.0
-            or getattr(layer, "bias_regularizer_l2", 0.0) > 0.0
-            for layer in self.trainable_layers
-        )
+        has_reg = bool(getattr(loss, "regularized_layers", ()))
         reg_loss_fn = loss.regularization_loss if has_reg else (lambda: 0.0)
 
         def run_step(batch_X, batch_y):

tests/losses/test_categorical_crossentropy.py

@@ -1,5 +1,6 @@
 import aether.config as config
 import tests.base_case as base_case
+from aether.layers.linear import Dense
 from aether.losses.categorical_crossentropy import CategoricalCrossEntropy
 
 
@@ -269,5 +270,29 @@ def test_new_pass_resets_accumulators(self):
         self.assertEqual(self.layer.accumulated_sum, 0)
         self.assertEqual(self.layer.accumulated_count, 0)
 
+    # ---- Regularization ---------------------
+
+    def test_regularization_loss_known_values(self):
+        """One regularizer per layer, weights [1, 2] and biases [3, -1] each, so dropping
+        any of the four terms changes the total; the unregularized layer adds nothing."""
+        regularizers = [
+            {"l1": (0.2,)},         # 0.2 * (1 + 2)                = 0.6
+            {"l2": (0.1,)},         # 0.5 * 0.1 * (1 + 4)          = 0.25
+            {"l1": (0.0, 0.3)},     # 0.3 * (3 + 1)                = 1.2
+            {"l2": (0.0, 0.4)},     # 0.5 * 0.4 * (9 + 1)          = 2.0
+            {},                     # unregularized                = 0
+        ]
+        layers = []
+        for kwargs in regularizers:
+            layer = self.make_built_layer(Dense, input_shape=(1,), n_inputs=1, n_neurons=2, **kwargs)
+            layer.weights = self.xp.array([[1.0, 2.0]], dtype=self.xp.float32)
+            layer.biases = self.xp.array([[3.0, -1.0]], dtype=self.xp.float32)
+            layers.append(layer)
+        layers[-1].weights = self.xp.array([[100.0, 100.0]], dtype=self.xp.float32)
+
+        self.layer.remember_trainable_layers(layers)
+
+        self.assertAlmostEqual(float(self.layer.regularization_loss()), 4.05, places=5)
+
 
 base_case.register_test_suites(globals(), TestLossCCE)

tests/model/test_train.py

@@ -1,9 +1,11 @@
 import contextlib
 import io
+from unittest import mock
 
 import numpy as np
 
 import aether.config as config
+from aether._utils.progress import make_progress
 from aether.layers import Dense, ReLU
 from aether.losses import CategoricalCrossEntropy, Loss
 from aether.metrics import CategoricalAccuracy, RegressionAccuracy
@@ -101,17 +103,22 @@ def test_train_without_shuffle(self):
         self.assertFalse(self.xp.allclose(initial_weights, trainable_layer.weights))
 
     def test_train_with_regularization(self):
-        model = Model()
-        model.to(self.backend_name)
-        model.add(Dense(n_inputs=4, n_neurons=self.NUM_CLASSES, l2=(0.01)))
-        model.configure(
-            loss=CategoricalCrossEntropy(),
-            optimizer=AdamW(lr=0.01),
-            accuracy=CategoricalAccuracy()
-        )
-        model.finalize(input_shape=(self.NUM_FEATURES,))
+        """train() tracks the regularization loss only when some layer is regularized."""
+        for l2, expected in (((0.01,), True), ((), False)):
+            with self.subTest(l2=l2):
+                model = Model()
+                model.to(self.backend_name)
+                model.add(Dense(n_inputs=4, n_neurons=self.NUM_CLASSES, l2=l2))
+                model.configure(
+                    loss=CategoricalCrossEntropy(),
+                    optimizer=AdamW(lr=0.01),
+                    accuracy=CategoricalAccuracy()
+                )
+                model.finalize(input_shape=(self.NUM_FEATURES,))
 
-        model.train(self.X, self.y, epochs=1, batch_size=16, verbose=False)
+                with mock.patch("aether.model.make_progress", wraps=make_progress) as progress:
+                    model.train(self.X, self.y, epochs=1, batch_size=16, verbose=False)
+                self.assertIs(progress.call_args.args[3], expected)  # has_reg
 
     def test_train_with_validation_data(self):
         self.model.finalize(input_shape=(self.NUM_FEATURES,))