Minor fixes

changelog_entry.yaml

@@ -0,0 +1,5 @@
+- bump: patch
+  changes:
+    added:
+    - Added a parameter to adjust the learning rate of the sparse optimizer.
+    - Fixed label in dashboard that incorrectly displayed 'estimate' instead of 'target'.

microcalibration-dashboard/src/components/LossChart.tsx

@@ -179,7 +179,7 @@ export default function LossChart({ data }: LossChartProps) {
             <Tooltip 
               formatter={(value: number, name: string) => [
                 formatValue(value), 
-                name === 'target' ? 'Target' : 'Estimate'
+                name === 'Target' ? 'Target' : 'Estimate'
               ]}
               labelFormatter={(label) => `Epoch: ${label}`}
             />

src/microcalibrate/calibration.py

@@ -27,6 +27,7 @@ def __init__(
         device: str = "cpu",  # fix to cpu for now to avoid user device-specific issues
         l0_lambda: float = 5e-6,  # best between 1e-6 and 1e-5
         init_mean: float = 0.999,  # initial proportion with non-zero weights, set near 0
+        sparse_learning_rate: float = 0.2,
         temperature: float = 0.5,  # usual values .5 to 3
         regularize_with_l0: Optional[bool] = False,
     ):
@@ -49,6 +50,7 @@ def __init__(
             l0_lambda (float): Regularization parameter for L0 regularization. Defaults to 5e-6.
             init_mean (float): Initial mean for L0 regularization, representing the initial proportion of non-zero weights. Defaults to 0.999.
             temperature (float): Temperature parameter for L0 regularization, controlling the sparsity of the model. Defaults to 0.5.
+            sparse_learning_rate (float): Learning rate for the regularizing optimizer. Defaults to 0.2.
             regularize_with_l0 (Optional[bool]): Whether to apply L0 regularization. Defaults to False.
         """
         if device is not None:
@@ -77,6 +79,7 @@ def __init__(
         self.l0_lambda = l0_lambda
         self.init_mean = init_mean
         self.temperature = temperature
+        self.sparse_learning_rate = sparse_learning_rate
         self.regularize_with_l0 = regularize_with_l0
 
         self.estimate_matrix = None
@@ -151,6 +154,7 @@ def calibrate(self) -> None:
             l0_lambda=self.l0_lambda,
             init_mean=self.init_mean,
             temperature=self.temperature,
+            sparse_learning_rate=self.sparse_learning_rate,
             regularize_with_l0=self.regularize_with_l0,
         )
 

src/microcalibrate/reweight.py

@@ -25,6 +25,7 @@ def reweight(
     init_mean: float,
     temperature: float,
     regularize_with_l0: bool,
+    sparse_learning_rate: Optional[float] = 0.2,
     dropout_rate: Optional[float] = 0.05,
     epochs: Optional[int] = 2_000,
     noise_level: Optional[float] = 10.0,
@@ -45,6 +46,7 @@ def reweight(
         l0_lambda (float): Regularization parameter for L0 regularization.
         init_mean (float): Initial mean for L0 regularization, representing the initial proportion of non-zero weights.
         temperature (float): Temperature parameter for L0 regularization, controlling the sparsity of the model.
+        sparse_learning_rate (float): Learning rate for the regularizing optimizer.
         regularize_with_l0 (bool): Whether to apply L0 regularization.
         dropout_rate (float): Optional probability of dropping weights during training.
         epochs (int): Optional number of epochs for training.
@@ -200,7 +202,7 @@ def dropout_weights(weights: torch.Tensor, p: float) -> torch.Tensor:
         # NOTE: Results are pretty sensitve to learning rates
         # optimizer breaks down somewhere near .005, does better at above .1
         optimizer = torch.optim.Adam(
-            [weights] + list(gates.parameters()), lr=0.2
+            [weights] + list(gates.parameters()), lr=sparse_learning_rate
         )
         start_loss = None