l0-python 0.3.0__py3-none-any.whl → 0.4.0__py3-none-any.whl

l0/calibration.py

@@ -30,10 +30,21 @@ class SparseCalibrationWeights(nn.Module):
         Lower bound of stretched concrete distribution
     zeta : float
         Upper bound of stretched concrete distribution
-    init_keep_prob : float
-        Initial probability of keeping each weight active
-    init_weight_scale : float
-        Initial scale for log weights (controls initial weight magnitude)
+    init_keep_prob : float or array-like
+        Initial probability of keeping each weight active.
+        If float, all weights use the same probability.
+        If array, must have shape (n_features,)
+    init_weights : float, array-like, or None
+        Initial weight values (on original scale, not log).
+        If float, all weights initialized to this value.
+        If array, must have shape (n_features,).
+        If None, defaults to 1.0 for all weights.
+    log_weight_jitter_sd : float
+        Standard deviation of noise added to log weights at start of fit() to break symmetry.
+        Set to 0 to disable jitter. Default is 0.0 (no jitter).
+    log_alpha_jitter_sd : float
+        Standard deviation of noise added to log_alpha at initialization to break symmetry.
+        Set to 0 to disable jitter. Default is 0.01 (following Louizos et al.).
     device : str or torch.device
         Device to run computations on ('cpu' or 'cuda')
     """
@@ -44,8 +55,10 @@ class SparseCalibrationWeights(nn.Module):
         beta: float = 2 / 3,
         gamma: float = -0.1,
         zeta: float = 1.1,
-        init_keep_prob: float = 0.5,
-        init_weight_scale: float = 1.0,
+        init_keep_prob: float | np.ndarray = 0.5,
+        init_weights: float | np.ndarray | None = None,
+        log_weight_jitter_sd: float = 0.0,
+        log_alpha_jitter_sd: float = 0.01,
         device: str | torch.device = "cpu",
     ):
         super().__init__()
@@ -53,25 +66,65 @@ class SparseCalibrationWeights(nn.Module):
         self.beta = beta
         self.gamma = gamma
         self.zeta = zeta
+        self.log_weight_jitter_sd = log_weight_jitter_sd
+        self.log_alpha_jitter_sd = log_alpha_jitter_sd
         self.device = torch.device(device)
 
-        # Log weights to ensure positivity via exp transformation
-        self.log_weight = nn.Parameter(
-            torch.normal(
-                mean=0.0,
-                std=init_weight_scale,
-                size=(n_features,),
-                device=self.device,
+        # Initialize weights (on original scale)
+        if init_weights is None:
+            # Default: all weights start at 1.0
+            weight_values = torch.ones(n_features, device=self.device)
+        elif isinstance(init_weights, (int, float)):
+            # Scalar: all weights start at this value
+            weight_values = torch.full(
+                (n_features,), float(init_weights), device=self.device
             )
-        )
+        else:
+            # Array: use provided values
+            weight_values = torch.tensor(
+                init_weights, dtype=torch.float32, device=self.device
+            )
+            if weight_values.shape != (n_features,):
+                raise ValueError(
+                    f"init_weights array must have shape ({n_features},), "
+                    f"got {weight_values.shape}"
+                )
+
+        # Convert to log space to ensure positivity via exp transformation
+        # Add small epsilon to avoid log(0)
+        self.log_weight = nn.Parameter(torch.log(weight_values + 1e-8))
 
         # L0 gate parameters
-        mu = torch.log(torch.tensor(init_keep_prob / (1 - init_keep_prob)))
-        self.log_alpha = nn.Parameter(
-            torch.normal(
-                mu.item(), 0.01, size=(n_features,), device=self.device
+        # Handle init_keep_prob as scalar or array
+        if isinstance(init_keep_prob, (int, float)):
+            # Scalar: broadcast to all features
+            keep_prob_values = torch.full(
+                (n_features,), float(init_keep_prob), device=self.device
             )
-        )
+        else:
+            # Array: use provided values
+            keep_prob_values = torch.tensor(
+                init_keep_prob, dtype=torch.float32, device=self.device
+            )
+            if keep_prob_values.shape != (n_features,):
+                raise ValueError(
+                    f"init_keep_prob array must have shape ({n_features},), "
+                    f"got {keep_prob_values.shape}"
+                )
+            # Clip to valid probability range to avoid log(0) or log(inf)
+            keep_prob_values = keep_prob_values.clamp(1e-6, 1 - 1e-6)
+
+        # Convert probabilities to log_alpha
+        mu = torch.log(keep_prob_values / (1 - keep_prob_values))
+        # Add jitter to break symmetry (if specified)
+        if self.log_alpha_jitter_sd > 0:
+            jitter = (
+                torch.randn(n_features, device=self.device)
+                * self.log_alpha_jitter_sd
+            )
+            self.log_alpha = nn.Parameter(mu + jitter)
+        else:
+            self.log_alpha = nn.Parameter(mu)
 
         # Cache for sparse tensor conversion
         self._cached_M_torch: torch.sparse.Tensor | None = None
@@ -312,8 +365,14 @@ class SparseCalibrationWeights(nn.Module):
             # No grouping - all targets weighted equally
             group_weights = torch.ones_like(y)
 
-        # Initialize weights
-        nn.init.normal_(self.log_weight, 0, 0.5)
+        # Add jitter to weights to break symmetry (if jitter_sd > 0)
+        if self.log_weight_jitter_sd > 0:
+            with torch.no_grad():
+                jitter = (
+                    torch.randn_like(self.log_weight)
+                    * self.log_weight_jitter_sd
+                )
+                self.log_weight.data += jitter
 
         # Setup optimizer
         optimizer = torch.optim.Adam([self.log_weight, self.log_alpha], lr=lr)

{l0_python-0.3.0.dist-info → l0_python-0.4.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: l0-python
-Version: 0.3.0
+Version: 0.4.0
 Summary: L0 regularization for sparse neural networks and intelligent sampling
 Author-email: PolicyEngine <hello@policyengine.org>
 License: MIT License
@@ -36,7 +36,7 @@ Classifier: License :: OSI Approved :: MIT License
 Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.13
 Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
-Requires-Python: >=3.13
+Requires-Python: >=3.11
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: torch>=2.0.0

{l0_python-0.3.0.dist-info → l0_python-0.4.0.dist-info}/RECORD

@@ -1,12 +1,12 @@
 l0/__init__.py,sha256=ezGVLGbdU_jAElAJSdLKjSWCnZRILAZi8qMBPjcHwQc,1192
-l0/calibration.py,sha256=H-5bo0SOCQXNBFuAQCzfBxw9-j4Ta0DcmJ66Hv-3Qtc,14860
+l0/calibration.py,sha256=EnIhQPkC6yvRp0-o1qG3SCVUp6pIqwkqkx3DX3uWdG0,17838
 l0/distributions.py,sha256=GvVoDasdxoDlPPdzt33lFzkuJgQsfkEn9e-abxru3nM,6325
 l0/gates.py,sha256=54OORMg9qSsdbIyzYIL9t9CCHWsKbumSKl_d4-YFuNU,11443
 l0/layers.py,sha256=j_DPEFdfUc37ojtoK2ufGzVYzZ_DwusqTcJXCVxZS68,13785
 l0/penalties.py,sha256=okqXgX2eGUmg5prcDclicTQALPPPF9FA5eQP09V4GTE,9569
 l0/sparse.py,sha256=UVpkrrJBnQ9iSNDoWK5F_sr43tQareWXzYHyMq_cKss,9726
-l0_python-0.3.0.dist-info/licenses/LICENSE,sha256=LoyFJ8qmpwWPp5RmxFGHhdBxsemvp_BqGyiCgNAZHlk,1069
-l0_python-0.3.0.dist-info/METADATA,sha256=sBDt2EahzEyNWSBsufl0geYWY5mwq9JMkSTVGppjPkg,6833
-l0_python-0.3.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-l0_python-0.3.0.dist-info/top_level.txt,sha256=brP2SoElgmGhNJk-qRFEQZ5UE-d1e87Tr60hSSh8n34,3
-l0_python-0.3.0.dist-info/RECORD,,
+l0_python-0.4.0.dist-info/licenses/LICENSE,sha256=LoyFJ8qmpwWPp5RmxFGHhdBxsemvp_BqGyiCgNAZHlk,1069
+l0_python-0.4.0.dist-info/METADATA,sha256=tVMaa0UIafPSIv9V9WWoxga7x2xtIVMhQ_egiX3C3y4,6833
+l0_python-0.4.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+l0_python-0.4.0.dist-info/top_level.txt,sha256=brP2SoElgmGhNJk-qRFEQZ5UE-d1e87Tr60hSSh8n34,3
+l0_python-0.4.0.dist-info/RECORD,,