l0-python 0.4.0__tar.gz → 0.4.2__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {l0_python-0.4.0 → l0_python-0.4.2}/PKG-INFO +108 -69
- l0_python-0.4.2/README.md +190 -0
- {l0_python-0.4.0 → l0_python-0.4.2}/l0/calibration.py +4 -4
- {l0_python-0.4.0 → l0_python-0.4.2}/l0_python.egg-info/PKG-INFO +108 -69
- {l0_python-0.4.0 → l0_python-0.4.2}/l0_python.egg-info/requires.txt +1 -1
- {l0_python-0.4.0 → l0_python-0.4.2}/pyproject.toml +2 -2
- {l0_python-0.4.0 → l0_python-0.4.2}/tests/test_calibration.py +226 -0
- l0_python-0.4.0/README.md +0 -151
- {l0_python-0.4.0 → l0_python-0.4.2}/LICENSE +0 -0
- {l0_python-0.4.0 → l0_python-0.4.2}/l0/__init__.py +0 -0
- {l0_python-0.4.0 → l0_python-0.4.2}/l0/distributions.py +0 -0
- {l0_python-0.4.0 → l0_python-0.4.2}/l0/gates.py +0 -0
- {l0_python-0.4.0 → l0_python-0.4.2}/l0/layers.py +0 -0
- {l0_python-0.4.0 → l0_python-0.4.2}/l0/penalties.py +0 -0
- {l0_python-0.4.0 → l0_python-0.4.2}/l0/sparse.py +0 -0
- {l0_python-0.4.0 → l0_python-0.4.2}/l0_python.egg-info/SOURCES.txt +0 -0
- {l0_python-0.4.0 → l0_python-0.4.2}/l0_python.egg-info/dependency_links.txt +0 -0
- {l0_python-0.4.0 → l0_python-0.4.2}/l0_python.egg-info/top_level.txt +0 -0
- {l0_python-0.4.0 → l0_python-0.4.2}/setup.cfg +0 -0
- {l0_python-0.4.0 → l0_python-0.4.2}/tests/test_distributions.py +0 -0
- {l0_python-0.4.0 → l0_python-0.4.2}/tests/test_gates.py +0 -0
- {l0_python-0.4.0 → l0_python-0.4.2}/tests/test_layers.py +0 -0
- {l0_python-0.4.0 → l0_python-0.4.2}/tests/test_penalties.py +0 -0
- {l0_python-0.4.0 → l0_python-0.4.2}/tests/test_sparse.py +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: l0-python
|
|
3
|
-
Version: 0.4.
|
|
3
|
+
Version: 0.4.2
|
|
4
4
|
Summary: L0 regularization for sparse neural networks and intelligent sampling
|
|
5
5
|
Author-email: PolicyEngine <hello@policyengine.org>
|
|
6
6
|
License: MIT License
|
|
@@ -51,7 +51,7 @@ Requires-Dist: scikit-learn>=1.0; extra == "dev"
|
|
|
51
51
|
Requires-Dist: yaml-changelog>=0.3.0; extra == "dev"
|
|
52
52
|
Requires-Dist: twine>=4.0.0; extra == "dev"
|
|
53
53
|
Provides-Extra: docs
|
|
54
|
-
Requires-Dist:
|
|
54
|
+
Requires-Dist: jupyter-book>=1.0; extra == "docs"
|
|
55
55
|
Requires-Dist: jupyter>=1.0; extra == "docs"
|
|
56
56
|
Requires-Dist: notebook>=7.0; extra == "docs"
|
|
57
57
|
Requires-Dist: ipykernel>=6.0; extra == "docs"
|
|
@@ -65,21 +65,29 @@ Dynamic: license-file
|
|
|
65
65
|
|
|
66
66
|
# L0 Regularization
|
|
67
67
|
|
|
68
|
-
|
|
68
|
+
[](https://pypi.org/project/l0-python/)
|
|
69
|
+
[](https://github.com/PolicyEngine/L0/actions)
|
|
69
70
|
|
|
70
|
-
|
|
71
|
+
A PyTorch implementation of L0 regularization based on [Louizos, Welling, & Kingma (2017)](https://arxiv.org/abs/1712.01312), with a **critical bug fix** to the original authors' implementation.
|
|
71
72
|
|
|
72
|
-
|
|
73
|
-
|
|
74
|
-
-
|
|
75
|
-
|
|
76
|
-
|
|
77
|
-
|
|
73
|
+
## Why This Package?
|
|
74
|
+
|
|
75
|
+
The original L0 implementation from AMLab-Amsterdam contains a bug in test-time gate computation where the temperature parameter is incorrectly omitted. This prevents proper sparsity and degrades performance. Our implementation corrects this error:
|
|
76
|
+
|
|
77
|
+
```python
|
|
78
|
+
# Original (buggy): gates never fully close
|
|
79
|
+
pi = sigmoid(log_alpha)
|
|
80
|
+
|
|
81
|
+
# Corrected: temperature division required for proper sparsity
|
|
82
|
+
pi = sigmoid(log_alpha / temperature)
|
|
83
|
+
```
|
|
84
|
+
|
|
85
|
+
This fix enables gates to achieve true 0/1 values, producing exact sparsity as intended by the L0 formulation.
|
|
78
86
|
|
|
79
87
|
## Installation
|
|
80
88
|
|
|
81
89
|
```bash
|
|
82
|
-
pip install l0
|
|
90
|
+
pip install l0-python
|
|
83
91
|
```
|
|
84
92
|
|
|
85
93
|
For development:
|
|
@@ -89,21 +97,76 @@ cd L0
|
|
|
89
97
|
pip install -e .[dev]
|
|
90
98
|
```
|
|
91
99
|
|
|
92
|
-
##
|
|
100
|
+
## Primary Use Case: Survey Calibration
|
|
101
|
+
|
|
102
|
+
This package was developed for PolicyEngine's survey calibration system, where we select a sparse subset of survey households while matching population targets.
|
|
103
|
+
|
|
104
|
+
```python
|
|
105
|
+
import numpy as np
|
|
106
|
+
from scipy import sparse as sp
|
|
107
|
+
from l0.calibration import SparseCalibrationWeights
|
|
108
|
+
|
|
109
|
+
# Setup: Q targets, N households
|
|
110
|
+
Q, N = 200, 10000
|
|
111
|
+
M = sp.random(Q, N, density=0.3, format="csr") # Household characteristics
|
|
112
|
+
y = np.random.uniform(1e6, 1e8, size=Q) # Population targets
|
|
113
|
+
|
|
114
|
+
# Initialize model
|
|
115
|
+
model = SparseCalibrationWeights(
|
|
116
|
+
n_features=N,
|
|
117
|
+
beta=0.35,
|
|
118
|
+
gamma=-0.1,
|
|
119
|
+
zeta=1.1,
|
|
120
|
+
init_keep_prob=0.5, # Start with 50% active probability
|
|
121
|
+
init_weights=1.0, # Or pass array of initial weights
|
|
122
|
+
log_weight_jitter_sd=0.05,
|
|
123
|
+
device="cuda", # GPU acceleration
|
|
124
|
+
)
|
|
125
|
+
|
|
126
|
+
# Train with L0+L2 regularization
|
|
127
|
+
model.fit(
|
|
128
|
+
M=M,
|
|
129
|
+
y=y,
|
|
130
|
+
lambda_l0=1e-6, # Controls sparsity level
|
|
131
|
+
lambda_l2=1e-8, # Prevents weight explosion
|
|
132
|
+
lr=0.15,
|
|
133
|
+
epochs=2000,
|
|
134
|
+
loss_type="relative", # Scale-invariant loss
|
|
135
|
+
verbose=True,
|
|
136
|
+
verbose_freq=100,
|
|
137
|
+
)
|
|
138
|
+
|
|
139
|
+
# Get results
|
|
140
|
+
active = model.get_active_weights()
|
|
141
|
+
print(f"Selected {active['count']} of {N} households ({100*active['count']/N:.1f}%)")
|
|
142
|
+
print(f"Sparsity: {model.get_sparsity():.1%}")
|
|
143
|
+
|
|
144
|
+
# Predict calibrated totals
|
|
145
|
+
y_pred = model.predict(M)
|
|
146
|
+
```
|
|
147
|
+
|
|
148
|
+
### Key Features for Calibration
|
|
149
|
+
|
|
150
|
+
- **Non-negative weights**: All weights constrained to be positive via log-space parameterization
|
|
151
|
+
- **Sparse solutions**: L0 penalty directly minimizes the count of active weights
|
|
152
|
+
- **Relative loss**: Scale-invariant loss for targets spanning orders of magnitude
|
|
153
|
+
- **Group-wise averaging**: Balance loss contributions across target groups with different cardinalities
|
|
154
|
+
- **GPU support**: CUDA acceleration for large-scale problems
|
|
155
|
+
|
|
156
|
+
## Neural Network Sparsification
|
|
93
157
|
|
|
94
|
-
|
|
158
|
+
The package also supports traditional neural network pruning:
|
|
95
159
|
|
|
96
160
|
```python
|
|
97
161
|
import torch
|
|
98
162
|
from l0 import L0Linear, compute_l0l2_penalty, TemperatureScheduler, update_temperatures
|
|
99
163
|
|
|
100
|
-
# Create a sparse model
|
|
101
164
|
class SparseModel(torch.nn.Module):
|
|
102
165
|
def __init__(self):
|
|
103
166
|
super().__init__()
|
|
104
167
|
self.fc1 = L0Linear(784, 256, init_sparsity=0.5)
|
|
105
168
|
self.fc2 = L0Linear(256, 10, init_sparsity=0.7)
|
|
106
|
-
|
|
169
|
+
|
|
107
170
|
def forward(self, x):
|
|
108
171
|
x = torch.relu(self.fc1(x))
|
|
109
172
|
return self.fc2(x)
|
|
@@ -112,95 +175,71 @@ model = SparseModel()
|
|
|
112
175
|
optimizer = torch.optim.Adam(model.parameters())
|
|
113
176
|
scheduler = TemperatureScheduler(initial_temp=1.0, final_temp=0.1)
|
|
114
177
|
|
|
115
|
-
# Training loop
|
|
116
178
|
for epoch in range(100):
|
|
117
|
-
# Update temperature
|
|
118
179
|
temp = scheduler.get_temperature(epoch)
|
|
119
180
|
update_temperatures(model, temp)
|
|
120
|
-
|
|
121
|
-
# Forward pass
|
|
181
|
+
|
|
122
182
|
output = model(input_data)
|
|
123
183
|
ce_loss = criterion(output, target)
|
|
124
|
-
|
|
125
|
-
# Add L0L2 penalty
|
|
126
184
|
penalty = compute_l0l2_penalty(model, l0_lambda=1e-3, l2_lambda=1e-4)
|
|
127
185
|
loss = ce_loss + penalty
|
|
128
|
-
|
|
129
|
-
# Backward pass
|
|
186
|
+
|
|
130
187
|
optimizer.zero_grad()
|
|
131
188
|
loss.backward()
|
|
132
189
|
optimizer.step()
|
|
133
190
|
```
|
|
134
191
|
|
|
135
|
-
###
|
|
192
|
+
### Available Layers
|
|
193
|
+
|
|
194
|
+
- `L0Linear`: Fully connected layer with L0 gates on weights
|
|
195
|
+
- `L0Conv2d`: 2D convolution with channel-wise L0 gates
|
|
196
|
+
- `L0DepthwiseConv2d`: Depthwise convolution with L0 gates
|
|
197
|
+
- `SparseMLP`: Multi-layer perceptron with built-in L0 regularization
|
|
198
|
+
|
|
199
|
+
## Intelligent Sampling Gates
|
|
200
|
+
|
|
201
|
+
Standalone gates for sample/feature selection:
|
|
136
202
|
|
|
137
203
|
```python
|
|
138
|
-
from l0 import SampleGate, HybridGate
|
|
204
|
+
from l0 import SampleGate, FeatureGate, HybridGate
|
|
139
205
|
|
|
140
|
-
#
|
|
206
|
+
# Select samples via learned gates
|
|
141
207
|
gate = SampleGate(n_samples=10000, target_samples=1000)
|
|
142
208
|
selected_data, indices = gate.select_samples(data)
|
|
143
209
|
|
|
144
|
-
#
|
|
210
|
+
# Select features
|
|
211
|
+
gate = FeatureGate(n_features=1000, max_features=50)
|
|
212
|
+
selected_data, indices = gate.select_features(data)
|
|
213
|
+
|
|
214
|
+
# Hybrid: combine L0 selection with random sampling
|
|
145
215
|
hybrid = HybridGate(
|
|
146
216
|
n_items=10000,
|
|
147
|
-
l0_fraction=0.25,
|
|
148
|
-
random_fraction=0.75
|
|
149
|
-
target_items=1000
|
|
217
|
+
l0_fraction=0.25, # 25% via learned L0 gates
|
|
218
|
+
random_fraction=0.75 # 75% random for coverage
|
|
150
219
|
)
|
|
151
220
|
selected, indices, types = hybrid.select(data)
|
|
152
221
|
```
|
|
153
222
|
|
|
154
|
-
|
|
155
|
-
|
|
156
|
-
```python
|
|
157
|
-
from l0 import FeatureGate
|
|
158
|
-
|
|
159
|
-
# Select top features
|
|
160
|
-
gate = FeatureGate(n_features=1000, max_features=50)
|
|
161
|
-
selected_data, feature_indices = gate.select_features(data)
|
|
162
|
-
|
|
163
|
-
# Get feature importance
|
|
164
|
-
importance = gate.get_feature_importance()
|
|
165
|
-
```
|
|
166
|
-
|
|
167
|
-
## Integration with PolicyEngine
|
|
168
|
-
|
|
169
|
-
This package is designed to work with PolicyEngine's calibration system:
|
|
170
|
-
|
|
171
|
-
```python
|
|
172
|
-
# In policyengine-us-data or similar
|
|
173
|
-
from l0 import HardConcrete
|
|
174
|
-
|
|
175
|
-
# Use for household selection in CPS calibration
|
|
176
|
-
gates = HardConcrete(
|
|
177
|
-
len(household_weights),
|
|
178
|
-
temperature=0.25,
|
|
179
|
-
init_mean=0.999 # Start with most households
|
|
180
|
-
)
|
|
223
|
+
## How L0 Regularization Works
|
|
181
224
|
|
|
182
|
-
|
|
183
|
-
masked_weights = weights * gates()
|
|
184
|
-
```
|
|
225
|
+
Unlike post-hoc pruning (setting small weights to zero), L0 regularization integrates sparsity into the training objective:
|
|
185
226
|
|
|
186
|
-
|
|
227
|
+
1. **Stochastic gates**: Each weight has a learned gate parameter controlling activation probability
|
|
228
|
+
2. **Hard Concrete distribution**: Enables differentiable 0/1 gate values during training
|
|
229
|
+
3. **Expected L0 penalty**: Minimizes the expected number of active gates
|
|
230
|
+
4. **Temperature annealing**: Gradually sharpens gates from soft to hard decisions
|
|
187
231
|
|
|
188
|
-
|
|
232
|
+
The result: the network learns *which* weights should be zero as part of optimization, not as a post-processing step.
|
|
189
233
|
|
|
190
234
|
## Testing
|
|
191
235
|
|
|
192
|
-
Run tests with:
|
|
193
236
|
```bash
|
|
194
237
|
pytest tests/ -v --cov=l0
|
|
195
238
|
```
|
|
196
239
|
|
|
197
|
-
## Acknowledgments
|
|
198
|
-
|
|
199
|
-
This implementation is inspired by and builds upon the [original L0 regularization code](https://github.com/AMLab-Amsterdam/L0_regularization) by AMLab Amsterdam, which accompanied the paper by Louizos et al. (2018).
|
|
200
|
-
|
|
201
240
|
## Citation
|
|
202
241
|
|
|
203
|
-
If you use this package, please cite:
|
|
242
|
+
If you use this package, please cite the original paper:
|
|
204
243
|
|
|
205
244
|
```bibtex
|
|
206
245
|
@article{louizos2017learning,
|
|
@@ -213,4 +252,4 @@ If you use this package, please cite:
|
|
|
213
252
|
|
|
214
253
|
## License
|
|
215
254
|
|
|
216
|
-
MIT License - see [LICENSE](LICENSE)
|
|
255
|
+
MIT License - see [LICENSE](LICENSE) for details.
|
|
@@ -0,0 +1,190 @@
|
|
|
1
|
+
# L0 Regularization
|
|
2
|
+
|
|
3
|
+
[](https://pypi.org/project/l0-python/)
|
|
4
|
+
[](https://github.com/PolicyEngine/L0/actions)
|
|
5
|
+
|
|
6
|
+
A PyTorch implementation of L0 regularization based on [Louizos, Welling, & Kingma (2017)](https://arxiv.org/abs/1712.01312), with a **critical bug fix** to the original authors' implementation.
|
|
7
|
+
|
|
8
|
+
## Why This Package?
|
|
9
|
+
|
|
10
|
+
The original L0 implementation from AMLab-Amsterdam contains a bug in test-time gate computation where the temperature parameter is incorrectly omitted. This prevents proper sparsity and degrades performance. Our implementation corrects this error:
|
|
11
|
+
|
|
12
|
+
```python
|
|
13
|
+
# Original (buggy): gates never fully close
|
|
14
|
+
pi = sigmoid(log_alpha)
|
|
15
|
+
|
|
16
|
+
# Corrected: temperature division required for proper sparsity
|
|
17
|
+
pi = sigmoid(log_alpha / temperature)
|
|
18
|
+
```
|
|
19
|
+
|
|
20
|
+
This fix enables gates to achieve true 0/1 values, producing exact sparsity as intended by the L0 formulation.
|
|
21
|
+
|
|
22
|
+
## Installation
|
|
23
|
+
|
|
24
|
+
```bash
|
|
25
|
+
pip install l0-python
|
|
26
|
+
```
|
|
27
|
+
|
|
28
|
+
For development:
|
|
29
|
+
```bash
|
|
30
|
+
git clone https://github.com/PolicyEngine/L0.git
|
|
31
|
+
cd L0
|
|
32
|
+
pip install -e .[dev]
|
|
33
|
+
```
|
|
34
|
+
|
|
35
|
+
## Primary Use Case: Survey Calibration
|
|
36
|
+
|
|
37
|
+
This package was developed for PolicyEngine's survey calibration system, where we select a sparse subset of survey households while matching population targets.
|
|
38
|
+
|
|
39
|
+
```python
|
|
40
|
+
import numpy as np
|
|
41
|
+
from scipy import sparse as sp
|
|
42
|
+
from l0.calibration import SparseCalibrationWeights
|
|
43
|
+
|
|
44
|
+
# Setup: Q targets, N households
|
|
45
|
+
Q, N = 200, 10000
|
|
46
|
+
M = sp.random(Q, N, density=0.3, format="csr") # Household characteristics
|
|
47
|
+
y = np.random.uniform(1e6, 1e8, size=Q) # Population targets
|
|
48
|
+
|
|
49
|
+
# Initialize model
|
|
50
|
+
model = SparseCalibrationWeights(
|
|
51
|
+
n_features=N,
|
|
52
|
+
beta=0.35,
|
|
53
|
+
gamma=-0.1,
|
|
54
|
+
zeta=1.1,
|
|
55
|
+
init_keep_prob=0.5, # Start with 50% active probability
|
|
56
|
+
init_weights=1.0, # Or pass array of initial weights
|
|
57
|
+
log_weight_jitter_sd=0.05,
|
|
58
|
+
device="cuda", # GPU acceleration
|
|
59
|
+
)
|
|
60
|
+
|
|
61
|
+
# Train with L0+L2 regularization
|
|
62
|
+
model.fit(
|
|
63
|
+
M=M,
|
|
64
|
+
y=y,
|
|
65
|
+
lambda_l0=1e-6, # Controls sparsity level
|
|
66
|
+
lambda_l2=1e-8, # Prevents weight explosion
|
|
67
|
+
lr=0.15,
|
|
68
|
+
epochs=2000,
|
|
69
|
+
loss_type="relative", # Scale-invariant loss
|
|
70
|
+
verbose=True,
|
|
71
|
+
verbose_freq=100,
|
|
72
|
+
)
|
|
73
|
+
|
|
74
|
+
# Get results
|
|
75
|
+
active = model.get_active_weights()
|
|
76
|
+
print(f"Selected {active['count']} of {N} households ({100*active['count']/N:.1f}%)")
|
|
77
|
+
print(f"Sparsity: {model.get_sparsity():.1%}")
|
|
78
|
+
|
|
79
|
+
# Predict calibrated totals
|
|
80
|
+
y_pred = model.predict(M)
|
|
81
|
+
```
|
|
82
|
+
|
|
83
|
+
### Key Features for Calibration
|
|
84
|
+
|
|
85
|
+
- **Non-negative weights**: All weights constrained to be positive via log-space parameterization
|
|
86
|
+
- **Sparse solutions**: L0 penalty directly minimizes the count of active weights
|
|
87
|
+
- **Relative loss**: Scale-invariant loss for targets spanning orders of magnitude
|
|
88
|
+
- **Group-wise averaging**: Balance loss contributions across target groups with different cardinalities
|
|
89
|
+
- **GPU support**: CUDA acceleration for large-scale problems
|
|
90
|
+
|
|
91
|
+
## Neural Network Sparsification
|
|
92
|
+
|
|
93
|
+
The package also supports traditional neural network pruning:
|
|
94
|
+
|
|
95
|
+
```python
|
|
96
|
+
import torch
|
|
97
|
+
from l0 import L0Linear, compute_l0l2_penalty, TemperatureScheduler, update_temperatures
|
|
98
|
+
|
|
99
|
+
class SparseModel(torch.nn.Module):
|
|
100
|
+
def __init__(self):
|
|
101
|
+
super().__init__()
|
|
102
|
+
self.fc1 = L0Linear(784, 256, init_sparsity=0.5)
|
|
103
|
+
self.fc2 = L0Linear(256, 10, init_sparsity=0.7)
|
|
104
|
+
|
|
105
|
+
def forward(self, x):
|
|
106
|
+
x = torch.relu(self.fc1(x))
|
|
107
|
+
return self.fc2(x)
|
|
108
|
+
|
|
109
|
+
model = SparseModel()
|
|
110
|
+
optimizer = torch.optim.Adam(model.parameters())
|
|
111
|
+
scheduler = TemperatureScheduler(initial_temp=1.0, final_temp=0.1)
|
|
112
|
+
|
|
113
|
+
for epoch in range(100):
|
|
114
|
+
temp = scheduler.get_temperature(epoch)
|
|
115
|
+
update_temperatures(model, temp)
|
|
116
|
+
|
|
117
|
+
output = model(input_data)
|
|
118
|
+
ce_loss = criterion(output, target)
|
|
119
|
+
penalty = compute_l0l2_penalty(model, l0_lambda=1e-3, l2_lambda=1e-4)
|
|
120
|
+
loss = ce_loss + penalty
|
|
121
|
+
|
|
122
|
+
optimizer.zero_grad()
|
|
123
|
+
loss.backward()
|
|
124
|
+
optimizer.step()
|
|
125
|
+
```
|
|
126
|
+
|
|
127
|
+
### Available Layers
|
|
128
|
+
|
|
129
|
+
- `L0Linear`: Fully connected layer with L0 gates on weights
|
|
130
|
+
- `L0Conv2d`: 2D convolution with channel-wise L0 gates
|
|
131
|
+
- `L0DepthwiseConv2d`: Depthwise convolution with L0 gates
|
|
132
|
+
- `SparseMLP`: Multi-layer perceptron with built-in L0 regularization
|
|
133
|
+
|
|
134
|
+
## Intelligent Sampling Gates
|
|
135
|
+
|
|
136
|
+
Standalone gates for sample/feature selection:
|
|
137
|
+
|
|
138
|
+
```python
|
|
139
|
+
from l0 import SampleGate, FeatureGate, HybridGate
|
|
140
|
+
|
|
141
|
+
# Select samples via learned gates
|
|
142
|
+
gate = SampleGate(n_samples=10000, target_samples=1000)
|
|
143
|
+
selected_data, indices = gate.select_samples(data)
|
|
144
|
+
|
|
145
|
+
# Select features
|
|
146
|
+
gate = FeatureGate(n_features=1000, max_features=50)
|
|
147
|
+
selected_data, indices = gate.select_features(data)
|
|
148
|
+
|
|
149
|
+
# Hybrid: combine L0 selection with random sampling
|
|
150
|
+
hybrid = HybridGate(
|
|
151
|
+
n_items=10000,
|
|
152
|
+
l0_fraction=0.25, # 25% via learned L0 gates
|
|
153
|
+
random_fraction=0.75 # 75% random for coverage
|
|
154
|
+
)
|
|
155
|
+
selected, indices, types = hybrid.select(data)
|
|
156
|
+
```
|
|
157
|
+
|
|
158
|
+
## How L0 Regularization Works
|
|
159
|
+
|
|
160
|
+
Unlike post-hoc pruning (setting small weights to zero), L0 regularization integrates sparsity into the training objective:
|
|
161
|
+
|
|
162
|
+
1. **Stochastic gates**: Each weight has a learned gate parameter controlling activation probability
|
|
163
|
+
2. **Hard Concrete distribution**: Enables differentiable 0/1 gate values during training
|
|
164
|
+
3. **Expected L0 penalty**: Minimizes the expected number of active gates
|
|
165
|
+
4. **Temperature annealing**: Gradually sharpens gates from soft to hard decisions
|
|
166
|
+
|
|
167
|
+
The result: the network learns *which* weights should be zero as part of optimization, not as a post-processing step.
|
|
168
|
+
|
|
169
|
+
## Testing
|
|
170
|
+
|
|
171
|
+
```bash
|
|
172
|
+
pytest tests/ -v --cov=l0
|
|
173
|
+
```
|
|
174
|
+
|
|
175
|
+
## Citation
|
|
176
|
+
|
|
177
|
+
If you use this package, please cite the original paper:
|
|
178
|
+
|
|
179
|
+
```bibtex
|
|
180
|
+
@article{louizos2017learning,
|
|
181
|
+
title={Learning Sparse Neural Networks through L0 Regularization},
|
|
182
|
+
author={Louizos, Christos and Welling, Max and Kingma, Diederik P},
|
|
183
|
+
journal={arXiv preprint arXiv:1712.01312},
|
|
184
|
+
year={2017}
|
|
185
|
+
}
|
|
186
|
+
```
|
|
187
|
+
|
|
188
|
+
## License
|
|
189
|
+
|
|
190
|
+
MIT License - see [LICENSE](LICENSE) for details.
|
|
@@ -444,9 +444,9 @@ class SparseCalibrationWeights(nn.Module):
|
|
|
444
444
|
rel_errors[group_mask].mean().item()
|
|
445
445
|
)
|
|
446
446
|
group_losses.append(group_mean_err)
|
|
447
|
-
|
|
447
|
+
mean_group_mare = np.mean(group_losses)
|
|
448
448
|
else:
|
|
449
|
-
|
|
449
|
+
mean_group_mare = mean_rel_err
|
|
450
450
|
|
|
451
451
|
# Calculate sparsity percentage
|
|
452
452
|
sparsity_pct = 100 * (
|
|
@@ -461,7 +461,7 @@ class SparseCalibrationWeights(nn.Module):
|
|
|
461
461
|
if target_groups is not None:
|
|
462
462
|
print(
|
|
463
463
|
f"Epoch {epoch+1:4d}: "
|
|
464
|
-
f"
|
|
464
|
+
f"mean_group_mare={mean_group_mare:.4%}, "
|
|
465
465
|
f"max_error={max_rel_err:.1%}, "
|
|
466
466
|
f"total_loss={actual_total_loss:.3f}, "
|
|
467
467
|
f"active={active_info['count']:4d}/{self.n_features} ({sparsity_pct:.1f}% sparse)"
|
|
@@ -469,7 +469,7 @@ class SparseCalibrationWeights(nn.Module):
|
|
|
469
469
|
else:
|
|
470
470
|
print(
|
|
471
471
|
f"Epoch {epoch+1:4d}: "
|
|
472
|
-
f"mean_error={mean_rel_err:.
|
|
472
|
+
f"mean_error={mean_rel_err:.4%}, "
|
|
473
473
|
f"max_error={max_rel_err:.1%}, "
|
|
474
474
|
f"total_loss={actual_total_loss:.3f}, "
|
|
475
475
|
f"active={active_info['count']:4d}/{self.n_features} ({sparsity_pct:.1f}% sparse)"
|
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: l0-python
|
|
3
|
-
Version: 0.4.
|
|
3
|
+
Version: 0.4.2
|
|
4
4
|
Summary: L0 regularization for sparse neural networks and intelligent sampling
|
|
5
5
|
Author-email: PolicyEngine <hello@policyengine.org>
|
|
6
6
|
License: MIT License
|
|
@@ -51,7 +51,7 @@ Requires-Dist: scikit-learn>=1.0; extra == "dev"
|
|
|
51
51
|
Requires-Dist: yaml-changelog>=0.3.0; extra == "dev"
|
|
52
52
|
Requires-Dist: twine>=4.0.0; extra == "dev"
|
|
53
53
|
Provides-Extra: docs
|
|
54
|
-
Requires-Dist:
|
|
54
|
+
Requires-Dist: jupyter-book>=1.0; extra == "docs"
|
|
55
55
|
Requires-Dist: jupyter>=1.0; extra == "docs"
|
|
56
56
|
Requires-Dist: notebook>=7.0; extra == "docs"
|
|
57
57
|
Requires-Dist: ipykernel>=6.0; extra == "docs"
|
|
@@ -65,21 +65,29 @@ Dynamic: license-file
|
|
|
65
65
|
|
|
66
66
|
# L0 Regularization
|
|
67
67
|
|
|
68
|
-
|
|
68
|
+
[](https://pypi.org/project/l0-python/)
|
|
69
|
+
[](https://github.com/PolicyEngine/L0/actions)
|
|
69
70
|
|
|
70
|
-
|
|
71
|
+
A PyTorch implementation of L0 regularization based on [Louizos, Welling, & Kingma (2017)](https://arxiv.org/abs/1712.01312), with a **critical bug fix** to the original authors' implementation.
|
|
71
72
|
|
|
72
|
-
|
|
73
|
-
|
|
74
|
-
-
|
|
75
|
-
|
|
76
|
-
|
|
77
|
-
|
|
73
|
+
## Why This Package?
|
|
74
|
+
|
|
75
|
+
The original L0 implementation from AMLab-Amsterdam contains a bug in test-time gate computation where the temperature parameter is incorrectly omitted. This prevents proper sparsity and degrades performance. Our implementation corrects this error:
|
|
76
|
+
|
|
77
|
+
```python
|
|
78
|
+
# Original (buggy): gates never fully close
|
|
79
|
+
pi = sigmoid(log_alpha)
|
|
80
|
+
|
|
81
|
+
# Corrected: temperature division required for proper sparsity
|
|
82
|
+
pi = sigmoid(log_alpha / temperature)
|
|
83
|
+
```
|
|
84
|
+
|
|
85
|
+
This fix enables gates to achieve true 0/1 values, producing exact sparsity as intended by the L0 formulation.
|
|
78
86
|
|
|
79
87
|
## Installation
|
|
80
88
|
|
|
81
89
|
```bash
|
|
82
|
-
pip install l0
|
|
90
|
+
pip install l0-python
|
|
83
91
|
```
|
|
84
92
|
|
|
85
93
|
For development:
|
|
@@ -89,21 +97,76 @@ cd L0
|
|
|
89
97
|
pip install -e .[dev]
|
|
90
98
|
```
|
|
91
99
|
|
|
92
|
-
##
|
|
100
|
+
## Primary Use Case: Survey Calibration
|
|
101
|
+
|
|
102
|
+
This package was developed for PolicyEngine's survey calibration system, where we select a sparse subset of survey households while matching population targets.
|
|
103
|
+
|
|
104
|
+
```python
|
|
105
|
+
import numpy as np
|
|
106
|
+
from scipy import sparse as sp
|
|
107
|
+
from l0.calibration import SparseCalibrationWeights
|
|
108
|
+
|
|
109
|
+
# Setup: Q targets, N households
|
|
110
|
+
Q, N = 200, 10000
|
|
111
|
+
M = sp.random(Q, N, density=0.3, format="csr") # Household characteristics
|
|
112
|
+
y = np.random.uniform(1e6, 1e8, size=Q) # Population targets
|
|
113
|
+
|
|
114
|
+
# Initialize model
|
|
115
|
+
model = SparseCalibrationWeights(
|
|
116
|
+
n_features=N,
|
|
117
|
+
beta=0.35,
|
|
118
|
+
gamma=-0.1,
|
|
119
|
+
zeta=1.1,
|
|
120
|
+
init_keep_prob=0.5, # Start with 50% active probability
|
|
121
|
+
init_weights=1.0, # Or pass array of initial weights
|
|
122
|
+
log_weight_jitter_sd=0.05,
|
|
123
|
+
device="cuda", # GPU acceleration
|
|
124
|
+
)
|
|
125
|
+
|
|
126
|
+
# Train with L0+L2 regularization
|
|
127
|
+
model.fit(
|
|
128
|
+
M=M,
|
|
129
|
+
y=y,
|
|
130
|
+
lambda_l0=1e-6, # Controls sparsity level
|
|
131
|
+
lambda_l2=1e-8, # Prevents weight explosion
|
|
132
|
+
lr=0.15,
|
|
133
|
+
epochs=2000,
|
|
134
|
+
loss_type="relative", # Scale-invariant loss
|
|
135
|
+
verbose=True,
|
|
136
|
+
verbose_freq=100,
|
|
137
|
+
)
|
|
138
|
+
|
|
139
|
+
# Get results
|
|
140
|
+
active = model.get_active_weights()
|
|
141
|
+
print(f"Selected {active['count']} of {N} households ({100*active['count']/N:.1f}%)")
|
|
142
|
+
print(f"Sparsity: {model.get_sparsity():.1%}")
|
|
143
|
+
|
|
144
|
+
# Predict calibrated totals
|
|
145
|
+
y_pred = model.predict(M)
|
|
146
|
+
```
|
|
147
|
+
|
|
148
|
+
### Key Features for Calibration
|
|
149
|
+
|
|
150
|
+
- **Non-negative weights**: All weights constrained to be positive via log-space parameterization
|
|
151
|
+
- **Sparse solutions**: L0 penalty directly minimizes the count of active weights
|
|
152
|
+
- **Relative loss**: Scale-invariant loss for targets spanning orders of magnitude
|
|
153
|
+
- **Group-wise averaging**: Balance loss contributions across target groups with different cardinalities
|
|
154
|
+
- **GPU support**: CUDA acceleration for large-scale problems
|
|
155
|
+
|
|
156
|
+
## Neural Network Sparsification
|
|
93
157
|
|
|
94
|
-
|
|
158
|
+
The package also supports traditional neural network pruning:
|
|
95
159
|
|
|
96
160
|
```python
|
|
97
161
|
import torch
|
|
98
162
|
from l0 import L0Linear, compute_l0l2_penalty, TemperatureScheduler, update_temperatures
|
|
99
163
|
|
|
100
|
-
# Create a sparse model
|
|
101
164
|
class SparseModel(torch.nn.Module):
|
|
102
165
|
def __init__(self):
|
|
103
166
|
super().__init__()
|
|
104
167
|
self.fc1 = L0Linear(784, 256, init_sparsity=0.5)
|
|
105
168
|
self.fc2 = L0Linear(256, 10, init_sparsity=0.7)
|
|
106
|
-
|
|
169
|
+
|
|
107
170
|
def forward(self, x):
|
|
108
171
|
x = torch.relu(self.fc1(x))
|
|
109
172
|
return self.fc2(x)
|
|
@@ -112,95 +175,71 @@ model = SparseModel()
|
|
|
112
175
|
optimizer = torch.optim.Adam(model.parameters())
|
|
113
176
|
scheduler = TemperatureScheduler(initial_temp=1.0, final_temp=0.1)
|
|
114
177
|
|
|
115
|
-
# Training loop
|
|
116
178
|
for epoch in range(100):
|
|
117
|
-
# Update temperature
|
|
118
179
|
temp = scheduler.get_temperature(epoch)
|
|
119
180
|
update_temperatures(model, temp)
|
|
120
|
-
|
|
121
|
-
# Forward pass
|
|
181
|
+
|
|
122
182
|
output = model(input_data)
|
|
123
183
|
ce_loss = criterion(output, target)
|
|
124
|
-
|
|
125
|
-
# Add L0L2 penalty
|
|
126
184
|
penalty = compute_l0l2_penalty(model, l0_lambda=1e-3, l2_lambda=1e-4)
|
|
127
185
|
loss = ce_loss + penalty
|
|
128
|
-
|
|
129
|
-
# Backward pass
|
|
186
|
+
|
|
130
187
|
optimizer.zero_grad()
|
|
131
188
|
loss.backward()
|
|
132
189
|
optimizer.step()
|
|
133
190
|
```
|
|
134
191
|
|
|
135
|
-
###
|
|
192
|
+
### Available Layers
|
|
193
|
+
|
|
194
|
+
- `L0Linear`: Fully connected layer with L0 gates on weights
|
|
195
|
+
- `L0Conv2d`: 2D convolution with channel-wise L0 gates
|
|
196
|
+
- `L0DepthwiseConv2d`: Depthwise convolution with L0 gates
|
|
197
|
+
- `SparseMLP`: Multi-layer perceptron with built-in L0 regularization
|
|
198
|
+
|
|
199
|
+
## Intelligent Sampling Gates
|
|
200
|
+
|
|
201
|
+
Standalone gates for sample/feature selection:
|
|
136
202
|
|
|
137
203
|
```python
|
|
138
|
-
from l0 import SampleGate, HybridGate
|
|
204
|
+
from l0 import SampleGate, FeatureGate, HybridGate
|
|
139
205
|
|
|
140
|
-
#
|
|
206
|
+
# Select samples via learned gates
|
|
141
207
|
gate = SampleGate(n_samples=10000, target_samples=1000)
|
|
142
208
|
selected_data, indices = gate.select_samples(data)
|
|
143
209
|
|
|
144
|
-
#
|
|
210
|
+
# Select features
|
|
211
|
+
gate = FeatureGate(n_features=1000, max_features=50)
|
|
212
|
+
selected_data, indices = gate.select_features(data)
|
|
213
|
+
|
|
214
|
+
# Hybrid: combine L0 selection with random sampling
|
|
145
215
|
hybrid = HybridGate(
|
|
146
216
|
n_items=10000,
|
|
147
|
-
l0_fraction=0.25,
|
|
148
|
-
random_fraction=0.75
|
|
149
|
-
target_items=1000
|
|
217
|
+
l0_fraction=0.25, # 25% via learned L0 gates
|
|
218
|
+
random_fraction=0.75 # 75% random for coverage
|
|
150
219
|
)
|
|
151
220
|
selected, indices, types = hybrid.select(data)
|
|
152
221
|
```
|
|
153
222
|
|
|
154
|
-
|
|
155
|
-
|
|
156
|
-
```python
|
|
157
|
-
from l0 import FeatureGate
|
|
158
|
-
|
|
159
|
-
# Select top features
|
|
160
|
-
gate = FeatureGate(n_features=1000, max_features=50)
|
|
161
|
-
selected_data, feature_indices = gate.select_features(data)
|
|
162
|
-
|
|
163
|
-
# Get feature importance
|
|
164
|
-
importance = gate.get_feature_importance()
|
|
165
|
-
```
|
|
166
|
-
|
|
167
|
-
## Integration with PolicyEngine
|
|
168
|
-
|
|
169
|
-
This package is designed to work with PolicyEngine's calibration system:
|
|
170
|
-
|
|
171
|
-
```python
|
|
172
|
-
# In policyengine-us-data or similar
|
|
173
|
-
from l0 import HardConcrete
|
|
174
|
-
|
|
175
|
-
# Use for household selection in CPS calibration
|
|
176
|
-
gates = HardConcrete(
|
|
177
|
-
len(household_weights),
|
|
178
|
-
temperature=0.25,
|
|
179
|
-
init_mean=0.999 # Start with most households
|
|
180
|
-
)
|
|
223
|
+
## How L0 Regularization Works
|
|
181
224
|
|
|
182
|
-
|
|
183
|
-
masked_weights = weights * gates()
|
|
184
|
-
```
|
|
225
|
+
Unlike post-hoc pruning (setting small weights to zero), L0 regularization integrates sparsity into the training objective:
|
|
185
226
|
|
|
186
|
-
|
|
227
|
+
1. **Stochastic gates**: Each weight has a learned gate parameter controlling activation probability
|
|
228
|
+
2. **Hard Concrete distribution**: Enables differentiable 0/1 gate values during training
|
|
229
|
+
3. **Expected L0 penalty**: Minimizes the expected number of active gates
|
|
230
|
+
4. **Temperature annealing**: Gradually sharpens gates from soft to hard decisions
|
|
187
231
|
|
|
188
|
-
|
|
232
|
+
The result: the network learns *which* weights should be zero as part of optimization, not as a post-processing step.
|
|
189
233
|
|
|
190
234
|
## Testing
|
|
191
235
|
|
|
192
|
-
Run tests with:
|
|
193
236
|
```bash
|
|
194
237
|
pytest tests/ -v --cov=l0
|
|
195
238
|
```
|
|
196
239
|
|
|
197
|
-
## Acknowledgments
|
|
198
|
-
|
|
199
|
-
This implementation is inspired by and builds upon the [original L0 regularization code](https://github.com/AMLab-Amsterdam/L0_regularization) by AMLab Amsterdam, which accompanied the paper by Louizos et al. (2018).
|
|
200
|
-
|
|
201
240
|
## Citation
|
|
202
241
|
|
|
203
|
-
If you use this package, please cite:
|
|
242
|
+
If you use this package, please cite the original paper:
|
|
204
243
|
|
|
205
244
|
```bibtex
|
|
206
245
|
@article{louizos2017learning,
|
|
@@ -213,4 +252,4 @@ If you use this package, please cite:
|
|
|
213
252
|
|
|
214
253
|
## License
|
|
215
254
|
|
|
216
|
-
MIT License - see [LICENSE](LICENSE)
|
|
255
|
+
MIT License - see [LICENSE](LICENSE) for details.
|
|
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
|
|
|
4
4
|
|
|
5
5
|
[project]
|
|
6
6
|
name = "l0-python"
|
|
7
|
-
version = "0.4.
|
|
7
|
+
version = "0.4.2"
|
|
8
8
|
description = "L0 regularization for sparse neural networks and intelligent sampling"
|
|
9
9
|
readme = "README.md"
|
|
10
10
|
authors = [
|
|
@@ -38,7 +38,7 @@ dev = [
|
|
|
38
38
|
"twine>=4.0.0",
|
|
39
39
|
]
|
|
40
40
|
docs = [
|
|
41
|
-
"
|
|
41
|
+
"jupyter-book>=1.0",
|
|
42
42
|
"jupyter>=1.0",
|
|
43
43
|
"notebook>=7.0",
|
|
44
44
|
"ipykernel>=6.0",
|
|
@@ -247,6 +247,232 @@ class TestSparseCalibrationWeights:
|
|
|
247
247
|
weights_with_l2.max() <= weights_no_l2.max() * 2.0
|
|
248
248
|
), "L2 should prevent extreme weights"
|
|
249
249
|
|
|
250
|
+
def test_pure_l2_penalty(self):
|
|
251
|
+
"""Test pure L2 regularization without L0."""
|
|
252
|
+
N = 50
|
|
253
|
+
Q = 30
|
|
254
|
+
|
|
255
|
+
np.random.seed(42)
|
|
256
|
+
torch.manual_seed(42)
|
|
257
|
+
|
|
258
|
+
# Create underdetermined problem where L2 helps regularize
|
|
259
|
+
M = sp.random(Q, N, density=0.5, format="csr")
|
|
260
|
+
M.data = np.abs(M.data) * 2 # Scale up to create larger weights
|
|
261
|
+
y = np.ones(Q) * 10 # Simple target
|
|
262
|
+
|
|
263
|
+
# Train without any regularization
|
|
264
|
+
model_no_reg = SparseCalibrationWeights(
|
|
265
|
+
n_features=N,
|
|
266
|
+
init_weights=2.0, # Start with larger weights
|
|
267
|
+
init_keep_prob=0.999, # Keep all weights active
|
|
268
|
+
)
|
|
269
|
+
model_no_reg.fit(
|
|
270
|
+
M,
|
|
271
|
+
y,
|
|
272
|
+
lambda_l0=0.0, # No L0
|
|
273
|
+
lambda_l2=0.0, # No L2
|
|
274
|
+
epochs=1000,
|
|
275
|
+
lr=0.02,
|
|
276
|
+
verbose=False,
|
|
277
|
+
)
|
|
278
|
+
|
|
279
|
+
# Train with strong L2 regularization
|
|
280
|
+
model_l2_only = SparseCalibrationWeights(
|
|
281
|
+
n_features=N,
|
|
282
|
+
init_weights=2.0, # Same starting point
|
|
283
|
+
init_keep_prob=0.999, # Keep all weights active
|
|
284
|
+
)
|
|
285
|
+
model_l2_only.fit(
|
|
286
|
+
M,
|
|
287
|
+
y,
|
|
288
|
+
lambda_l0=0.0, # No L0
|
|
289
|
+
lambda_l2=10.0, # Very strong L2 to see clear effect
|
|
290
|
+
epochs=1000,
|
|
291
|
+
lr=0.02,
|
|
292
|
+
verbose=False,
|
|
293
|
+
)
|
|
294
|
+
|
|
295
|
+
with torch.no_grad():
|
|
296
|
+
weights_no_reg = model_no_reg.get_weights(deterministic=True)
|
|
297
|
+
weights_l2 = model_l2_only.get_weights(deterministic=True)
|
|
298
|
+
|
|
299
|
+
# Get non-zero weights
|
|
300
|
+
active_no_reg = weights_no_reg[weights_no_reg > 1e-8]
|
|
301
|
+
active_l2 = weights_l2[weights_l2 > 1e-8]
|
|
302
|
+
|
|
303
|
+
# L2 should keep most weights active (no sparsity from L2)
|
|
304
|
+
assert (
|
|
305
|
+
len(active_l2) >= N * 0.9
|
|
306
|
+
), f"Pure L2 should not induce sparsity, got {len(active_l2)}/{N}"
|
|
307
|
+
|
|
308
|
+
# L2 norm should be smaller with regularization
|
|
309
|
+
l2_norm_no_reg = (weights_no_reg**2).sum().sqrt()
|
|
310
|
+
l2_norm_with_l2 = (weights_l2**2).sum().sqrt()
|
|
311
|
+
assert (
|
|
312
|
+
l2_norm_with_l2 < l2_norm_no_reg
|
|
313
|
+
), f"L2 regularization should reduce L2 norm: {l2_norm_with_l2:.2f} vs {l2_norm_no_reg:.2f}"
|
|
314
|
+
|
|
315
|
+
# Check coefficient of variation (CV = std/mean) is lower with L2
|
|
316
|
+
# This captures that L2 shrinks weights toward each other
|
|
317
|
+
if active_no_reg.mean() > 1e-6 and active_l2.mean() > 1e-6:
|
|
318
|
+
cv_no_reg = active_no_reg.std() / active_no_reg.mean()
|
|
319
|
+
cv_l2 = active_l2.std() / active_l2.mean()
|
|
320
|
+
assert (
|
|
321
|
+
cv_l2 < cv_no_reg * 1.2
|
|
322
|
+
), f"L2 should reduce relative variation: CV {cv_l2:.2f} vs {cv_no_reg:.2f}"
|
|
323
|
+
|
|
324
|
+
# L2 should prevent extreme weights
|
|
325
|
+
assert (
|
|
326
|
+
weights_l2.max() < weights_no_reg.max() * 1.5
|
|
327
|
+
), f"L2 should limit max weights: {weights_l2.max():.2f} vs {weights_no_reg.max():.2f}"
|
|
328
|
+
|
|
329
|
+
# Both should still fit reasonably well
|
|
330
|
+
y_pred_no_reg = model_no_reg.predict(M).cpu().numpy()
|
|
331
|
+
y_pred_l2 = model_l2_only.predict(M).cpu().numpy()
|
|
332
|
+
|
|
333
|
+
error_no_reg = np.abs((y - y_pred_no_reg) / y).mean()
|
|
334
|
+
error_l2 = np.abs((y - y_pred_l2) / y).mean()
|
|
335
|
+
|
|
336
|
+
# L2 model may have slightly worse fit due to regularization
|
|
337
|
+
assert (
|
|
338
|
+
error_l2 < 0.5
|
|
339
|
+
), f"L2 model should still fit reasonably: {error_l2:.3f}"
|
|
340
|
+
# But the trade-off is worth it for regularization
|
|
341
|
+
|
|
342
|
+
def test_l0_l2_combination(self):
|
|
343
|
+
"""Test that combining L0 and L2 gives both sparsity and regularization."""
|
|
344
|
+
N = 100 # features
|
|
345
|
+
Q = 50 # targets
|
|
346
|
+
|
|
347
|
+
np.random.seed(42)
|
|
348
|
+
torch.manual_seed(42)
|
|
349
|
+
|
|
350
|
+
# Create problem with potential for overfitting
|
|
351
|
+
M = sp.random(Q, N, density=0.4, format="csr")
|
|
352
|
+
M.data = np.abs(M.data) * 3
|
|
353
|
+
y = np.random.uniform(5, 20, size=Q)
|
|
354
|
+
|
|
355
|
+
# Model 1: Only L0 (sparsity without weight regularization)
|
|
356
|
+
model_l0_only = SparseCalibrationWeights(
|
|
357
|
+
n_features=N,
|
|
358
|
+
init_weights=2.0,
|
|
359
|
+
init_keep_prob=0.5, # Start with 50% probability
|
|
360
|
+
)
|
|
361
|
+
model_l0_only.fit(
|
|
362
|
+
M,
|
|
363
|
+
y,
|
|
364
|
+
lambda_l0=0.01, # Stronger L0 for sparsity
|
|
365
|
+
lambda_l2=0.0, # No L2
|
|
366
|
+
epochs=2000,
|
|
367
|
+
lr=0.02,
|
|
368
|
+
verbose=False,
|
|
369
|
+
)
|
|
370
|
+
|
|
371
|
+
# Model 2: Only L2 (weight regularization without sparsity)
|
|
372
|
+
model_l2_only = SparseCalibrationWeights(
|
|
373
|
+
n_features=N,
|
|
374
|
+
init_weights=2.0,
|
|
375
|
+
init_keep_prob=0.999, # Keep all weights active
|
|
376
|
+
)
|
|
377
|
+
model_l2_only.fit(
|
|
378
|
+
M,
|
|
379
|
+
y,
|
|
380
|
+
lambda_l0=0.0, # No L0
|
|
381
|
+
lambda_l2=0.1, # Moderate L2
|
|
382
|
+
epochs=1500,
|
|
383
|
+
lr=0.02,
|
|
384
|
+
verbose=False,
|
|
385
|
+
)
|
|
386
|
+
|
|
387
|
+
# Model 3: Combined L0+L2 (both sparsity and weight regularization)
|
|
388
|
+
model_l0_l2 = SparseCalibrationWeights(
|
|
389
|
+
n_features=N,
|
|
390
|
+
init_weights=2.0,
|
|
391
|
+
init_keep_prob=0.5, # Same starting point as L0-only
|
|
392
|
+
)
|
|
393
|
+
model_l0_l2.fit(
|
|
394
|
+
M,
|
|
395
|
+
y,
|
|
396
|
+
lambda_l0=0.01, # Same L0 as model 1
|
|
397
|
+
lambda_l2=0.1, # Add L2 regularization
|
|
398
|
+
epochs=2000,
|
|
399
|
+
lr=0.02,
|
|
400
|
+
verbose=False,
|
|
401
|
+
)
|
|
402
|
+
|
|
403
|
+
with torch.no_grad():
|
|
404
|
+
# Get weights and stats for all models
|
|
405
|
+
weights_l0_only = model_l0_only.get_weights(deterministic=True)
|
|
406
|
+
weights_l2_only = model_l2_only.get_weights(deterministic=True)
|
|
407
|
+
weights_l0_l2 = model_l0_l2.get_weights(deterministic=True)
|
|
408
|
+
|
|
409
|
+
# Count active weights
|
|
410
|
+
active_l0_only = (weights_l0_only > 1e-6).sum().item()
|
|
411
|
+
active_l2_only = (weights_l2_only > 1e-6).sum().item()
|
|
412
|
+
active_l0_l2 = (weights_l0_l2 > 1e-6).sum().item()
|
|
413
|
+
|
|
414
|
+
# L0-only should have sparsity
|
|
415
|
+
assert (
|
|
416
|
+
active_l0_only < N * 0.8
|
|
417
|
+
), f"L0 should induce sparsity: {active_l0_only}/{N} active"
|
|
418
|
+
|
|
419
|
+
# L2-only should have no/little sparsity
|
|
420
|
+
assert (
|
|
421
|
+
active_l2_only > N * 0.9
|
|
422
|
+
), f"L2 alone shouldn't induce sparsity: {active_l2_only}/{N} active"
|
|
423
|
+
|
|
424
|
+
# L0+L2 should have sparsity (from L0)
|
|
425
|
+
assert (
|
|
426
|
+
active_l0_l2 < N * 0.8
|
|
427
|
+
), f"L0+L2 should have sparsity: {active_l0_l2}/{N} active"
|
|
428
|
+
|
|
429
|
+
# Among active weights, L0+L2 should have smaller magnitudes than L0-only (from L2)
|
|
430
|
+
active_mask_l0_only = weights_l0_only > 1e-6
|
|
431
|
+
active_mask_l0_l2 = weights_l0_l2 > 1e-6
|
|
432
|
+
|
|
433
|
+
if active_mask_l0_only.any() and active_mask_l0_l2.any():
|
|
434
|
+
# Compare L2 norms of active weights
|
|
435
|
+
l2_norm_l0_only = (
|
|
436
|
+
(weights_l0_only[active_mask_l0_only] ** 2).sum().sqrt()
|
|
437
|
+
)
|
|
438
|
+
l2_norm_l0_l2 = (
|
|
439
|
+
(weights_l0_l2[active_mask_l0_l2] ** 2).sum().sqrt()
|
|
440
|
+
)
|
|
441
|
+
|
|
442
|
+
# L0+L2 should have smaller weight norms than L0-only
|
|
443
|
+
# (L2 regularization effect on top of sparsity)
|
|
444
|
+
assert (
|
|
445
|
+
l2_norm_l0_l2 < l2_norm_l0_only * 1.2
|
|
446
|
+
), f"L0+L2 should have controlled weights: {l2_norm_l0_l2:.2f} vs L0-only {l2_norm_l0_only:.2f}"
|
|
447
|
+
|
|
448
|
+
# Check prediction quality for all models
|
|
449
|
+
y_pred_l0_only = model_l0_only.predict(M).cpu().numpy()
|
|
450
|
+
y_pred_l2_only = model_l2_only.predict(M).cpu().numpy()
|
|
451
|
+
y_pred_l0_l2 = model_l0_l2.predict(M).cpu().numpy()
|
|
452
|
+
|
|
453
|
+
error_l0_only = np.abs((y - y_pred_l0_only) / (y + 1)).mean()
|
|
454
|
+
error_l2_only = np.abs((y - y_pred_l2_only) / (y + 1)).mean()
|
|
455
|
+
error_l0_l2 = np.abs((y - y_pred_l0_l2) / (y + 1)).mean()
|
|
456
|
+
|
|
457
|
+
# All should fit reasonably well (L2-only may have slightly worse fit due to regularization)
|
|
458
|
+
assert (
|
|
459
|
+
error_l0_only < 0.35
|
|
460
|
+
), f"L0-only should fit well: {error_l0_only:.3f}"
|
|
461
|
+
assert (
|
|
462
|
+
error_l2_only < 0.35
|
|
463
|
+
), f"L2-only should fit well: {error_l2_only:.3f}"
|
|
464
|
+
assert (
|
|
465
|
+
error_l0_l2 < 0.35
|
|
466
|
+
), f"L0+L2 should fit well: {error_l0_l2:.3f}"
|
|
467
|
+
|
|
468
|
+
print(
|
|
469
|
+
f"\nL0-only: {active_l0_only}/{N} active, error={error_l0_only:.3f}"
|
|
470
|
+
)
|
|
471
|
+
print(
|
|
472
|
+
f"L2-only: {active_l2_only}/{N} active, error={error_l2_only:.3f}"
|
|
473
|
+
)
|
|
474
|
+
print(f"L0+L2: {active_l0_l2}/{N} active, error={error_l0_l2:.3f}")
|
|
475
|
+
|
|
250
476
|
def test_group_wise_averaging(self):
|
|
251
477
|
"""Test that group-wise averaging balances loss contributions."""
|
|
252
478
|
N = 100 # features (households)
|
l0_python-0.4.0/README.md
DELETED
|
@@ -1,151 +0,0 @@
|
|
|
1
|
-
# L0 Regularization
|
|
2
|
-
|
|
3
|
-
A PyTorch implementation of L0 regularization for neural network sparsification and intelligent sampling, based on [Louizos, Welling, & Kingma (2017)](https://arxiv.org/abs/1712.01312).
|
|
4
|
-
|
|
5
|
-
## Features
|
|
6
|
-
|
|
7
|
-
- **Hard Concrete Distribution**: Differentiable approximation of L0 norm
|
|
8
|
-
- **Sparse Neural Network Layers**: L0Linear, L0Conv2d with automatic pruning
|
|
9
|
-
- **Intelligent Sampling**: Sample/feature selection gates for calibration
|
|
10
|
-
- **L0L2 Combined Penalty**: Recommended approach to prevent overfitting
|
|
11
|
-
- **Temperature Scheduling**: Annealing for improved convergence
|
|
12
|
-
- **TDD Development**: Comprehensive test coverage
|
|
13
|
-
|
|
14
|
-
## Installation
|
|
15
|
-
|
|
16
|
-
```bash
|
|
17
|
-
pip install l0
|
|
18
|
-
```
|
|
19
|
-
|
|
20
|
-
For development:
|
|
21
|
-
```bash
|
|
22
|
-
git clone https://github.com/PolicyEngine/L0.git
|
|
23
|
-
cd L0
|
|
24
|
-
pip install -e .[dev]
|
|
25
|
-
```
|
|
26
|
-
|
|
27
|
-
## Quick Start
|
|
28
|
-
|
|
29
|
-
### Neural Network Sparsification
|
|
30
|
-
|
|
31
|
-
```python
|
|
32
|
-
import torch
|
|
33
|
-
from l0 import L0Linear, compute_l0l2_penalty, TemperatureScheduler, update_temperatures
|
|
34
|
-
|
|
35
|
-
# Create a sparse model
|
|
36
|
-
class SparseModel(torch.nn.Module):
|
|
37
|
-
def __init__(self):
|
|
38
|
-
super().__init__()
|
|
39
|
-
self.fc1 = L0Linear(784, 256, init_sparsity=0.5)
|
|
40
|
-
self.fc2 = L0Linear(256, 10, init_sparsity=0.7)
|
|
41
|
-
|
|
42
|
-
def forward(self, x):
|
|
43
|
-
x = torch.relu(self.fc1(x))
|
|
44
|
-
return self.fc2(x)
|
|
45
|
-
|
|
46
|
-
model = SparseModel()
|
|
47
|
-
optimizer = torch.optim.Adam(model.parameters())
|
|
48
|
-
scheduler = TemperatureScheduler(initial_temp=1.0, final_temp=0.1)
|
|
49
|
-
|
|
50
|
-
# Training loop
|
|
51
|
-
for epoch in range(100):
|
|
52
|
-
# Update temperature
|
|
53
|
-
temp = scheduler.get_temperature(epoch)
|
|
54
|
-
update_temperatures(model, temp)
|
|
55
|
-
|
|
56
|
-
# Forward pass
|
|
57
|
-
output = model(input_data)
|
|
58
|
-
ce_loss = criterion(output, target)
|
|
59
|
-
|
|
60
|
-
# Add L0L2 penalty
|
|
61
|
-
penalty = compute_l0l2_penalty(model, l0_lambda=1e-3, l2_lambda=1e-4)
|
|
62
|
-
loss = ce_loss + penalty
|
|
63
|
-
|
|
64
|
-
# Backward pass
|
|
65
|
-
optimizer.zero_grad()
|
|
66
|
-
loss.backward()
|
|
67
|
-
optimizer.step()
|
|
68
|
-
```
|
|
69
|
-
|
|
70
|
-
### Intelligent Sample Selection
|
|
71
|
-
|
|
72
|
-
```python
|
|
73
|
-
from l0 import SampleGate, HybridGate
|
|
74
|
-
|
|
75
|
-
# Pure L0 selection
|
|
76
|
-
gate = SampleGate(n_samples=10000, target_samples=1000)
|
|
77
|
-
selected_data, indices = gate.select_samples(data)
|
|
78
|
-
|
|
79
|
-
# Hybrid selection (25% L0, 75% random)
|
|
80
|
-
hybrid = HybridGate(
|
|
81
|
-
n_items=10000,
|
|
82
|
-
l0_fraction=0.25,
|
|
83
|
-
random_fraction=0.75,
|
|
84
|
-
target_items=1000
|
|
85
|
-
)
|
|
86
|
-
selected, indices, types = hybrid.select(data)
|
|
87
|
-
```
|
|
88
|
-
|
|
89
|
-
### Feature Selection
|
|
90
|
-
|
|
91
|
-
```python
|
|
92
|
-
from l0 import FeatureGate
|
|
93
|
-
|
|
94
|
-
# Select top features
|
|
95
|
-
gate = FeatureGate(n_features=1000, max_features=50)
|
|
96
|
-
selected_data, feature_indices = gate.select_features(data)
|
|
97
|
-
|
|
98
|
-
# Get feature importance
|
|
99
|
-
importance = gate.get_feature_importance()
|
|
100
|
-
```
|
|
101
|
-
|
|
102
|
-
## Integration with PolicyEngine
|
|
103
|
-
|
|
104
|
-
This package is designed to work with PolicyEngine's calibration system:
|
|
105
|
-
|
|
106
|
-
```python
|
|
107
|
-
# In policyengine-us-data or similar
|
|
108
|
-
from l0 import HardConcrete
|
|
109
|
-
|
|
110
|
-
# Use for household selection in CPS calibration
|
|
111
|
-
gates = HardConcrete(
|
|
112
|
-
len(household_weights),
|
|
113
|
-
temperature=0.25,
|
|
114
|
-
init_mean=0.999 # Start with most households
|
|
115
|
-
)
|
|
116
|
-
|
|
117
|
-
# Apply gates during reweighting
|
|
118
|
-
masked_weights = weights * gates()
|
|
119
|
-
```
|
|
120
|
-
|
|
121
|
-
## Documentation
|
|
122
|
-
|
|
123
|
-
Full documentation available at: https://policyengine.github.io/L0/
|
|
124
|
-
|
|
125
|
-
## Testing
|
|
126
|
-
|
|
127
|
-
Run tests with:
|
|
128
|
-
```bash
|
|
129
|
-
pytest tests/ -v --cov=l0
|
|
130
|
-
```
|
|
131
|
-
|
|
132
|
-
## Acknowledgments
|
|
133
|
-
|
|
134
|
-
This implementation is inspired by and builds upon the [original L0 regularization code](https://github.com/AMLab-Amsterdam/L0_regularization) by AMLab Amsterdam, which accompanied the paper by Louizos et al. (2018).
|
|
135
|
-
|
|
136
|
-
## Citation
|
|
137
|
-
|
|
138
|
-
If you use this package, please cite:
|
|
139
|
-
|
|
140
|
-
```bibtex
|
|
141
|
-
@article{louizos2017learning,
|
|
142
|
-
title={Learning Sparse Neural Networks through L0 Regularization},
|
|
143
|
-
author={Louizos, Christos and Welling, Max and Kingma, Diederik P},
|
|
144
|
-
journal={arXiv preprint arXiv:1712.01312},
|
|
145
|
-
year={2017}
|
|
146
|
-
}
|
|
147
|
-
```
|
|
148
|
-
|
|
149
|
-
## License
|
|
150
|
-
|
|
151
|
-
MIT License - see [LICENSE](LICENSE) file for details.
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|