Fix TTQ initialization to match paper exactly

Paper (Zhu et al., ICLR 2017) specifies:
- Wp = Wn = E[|W|] (mean of absolute weights)
- delta = 0.7 * E[|W|]

Our bug: Used std(W) for delta and 1.0 for wp/wn

Impact: Wrong initialization scale affects convergence

Also add comprehensive TTQ_VERIFICATION.md documenting:
- Paper algorithm reference
- Implementation decisions (softplus, activation quantization)
- Verification checklist
- Bugs fixed and lessons learned

Author: dariocazzani

TTQ_VERIFICATION.md

@@ -0,0 +1,105 @@
+# TTQ Implementation Verification
+
+## Paper Reference
+**Trained Ternary Quantization** (Zhu et al., ICLR 2017)
+arXiv: https://arxiv.org/abs/1612.01064
+
+## Algorithm Summary (from paper)
+
+### Forward Pass
+
+**Quantization (Eq. 1):**
+```
+W_t = { +Wp   if W > delta
+      { -Wn   if W < -delta
+      {  0    otherwise
+```
+
+**Initialization (Eq. 2, Section 3.1):**
+- Threshold: `delta = 0.7 * E[|W|]`
+- Positive scale: `Wp = E[|W|]`
+- Negative scale: `Wn = E[|W|]`
+
+Where `E[|W|]` = mean of absolute weight values
+
+### Backward Pass
+Straight-through estimator (STE): gradients flow as if no quantization
+
+### Key Innovation
+Unlike fixed ternary {-1, 0, +1}, TTQ learns three FP32 parameters per layer:
+- `Wp` (positive scale)
+- `Wn` (negative scale)
+- `delta` (threshold)
+
+## Our Implementation
+
+### Files
+- `bitnet/nn/ttq_quantization.py` - Core quantization function
+- `bitnet/nn/ttq_linear.py` - Linear layer with TTQ
+- `bitnet/nn/ttq_conv2d.py` - Conv2d layer with TTQ
+- `tests/test_ttq_layers.py` - Test suite
+
+### Key Design Decisions
+
+**1. Positivity Constraint:**
+- Paper assumes Wp, Wn, delta > 0 but doesn't specify enforcement
+- We use `F.softplus()` to ensure positivity while maintaining gradients
+- Returns tuple `(quantized, wp_pos, wn_pos)` for consistent scaling
+
+**2. Activation Quantization:**
+- TTQ paper only specifies weight quantization
+- We use BitNet's activation quantization (`quantize_activations` + `dequantize`)
+- This allows fair comparison: both methods quantize weights AND activations
+- Beta for dequantization: `beta = (wp_pos + wn_pos) / 2`
+
+**3. Initialization:**
+- Wp, Wn = `mean(abs(weight))`  ✓ Matches paper
+- delta = `0.7 * mean(abs(weight))`  ✓ Matches paper
+
+## Verification Checklist
+
+- [x] Quantization logic matches Eq. 1
+- [x] Threshold comparison: `W > delta` and `W < -delta`
+- [x] Three learnable parameters: wp, wn, delta
+- [x] Initialization: Wp = Wn = E[|W|]
+- [x] Initialization: delta = 0.7 * E[|W|]
+- [x] Straight-through estimator for gradients
+- [x] Positivity enforcement (softplus)
+- [x] Consistent scale usage in quantization and dequantization
+- [x] Test suite covers shapes, gradients, initialization, stability
+
+## Differences from Pure TTQ
+
+1. **Activation Quantization:** We add BitNet-style activation quantization (8-bit)
+   - Reason: Fair comparison (both methods quantize weights + activations)
+   - Impact: More realistic for deployment
+
+2. **Positivity Enforcement:** We use softplus, paper doesn't specify
+   - Reason: Prevent training instability from negative scales
+   - Impact: Minor, gradients still flow
+
+## Bugs Fixed
+
+1. **Double softplus application:** quantization used softplus(wp), dequantization used softplus(softplus(wp))
+   - Fixed: Return wp_pos, wn_pos from ttq_quantize
+
+2. **Wrong initialization:** Used std(W) instead of mean(|W|) for delta
+   - Fixed: Both use `weight.abs().mean()`
+
+3. **NaN losses:** Parameters could go negative without constraints
+   - Fixed: Softplus enforcement
+
+## Expected Behavior
+
+- **Accuracy:** Should achieve ~0.5-1.5% better than BitNet+Recipe (based on literature)
+- **Complexity:** Requires 2 FP32 params per layer (vs BitNet+Recipe's 1 FP32 layer)
+- **Trade-off:** Better accuracy, more deployment complexity
+
+## Test Results
+
+All 9 tests pass:
+- Forward pass shapes
+- Gradient flow (wp, wn get gradients)
+- Correct initialization (E[|W|] and 0.7*E[|W|])
+- Numerical stability (no NaN in 10 training steps)
+- Various kernel sizes

bitnet/nn/ttq_conv2d.py

@@ -20,13 +20,12 @@ def __init__(self, *args, num_bits: int = 8, **kwargs):  # type: ignore[no-untyp
         super().__init__(*args, **kwargs)
         self.num_bits = num_bits
 
-        # Learnable positive/negative scales (init to 1.0)
-        self.register_parameter("wp", nn.Parameter(torch.ones(1)))
-        self.register_parameter("wn", nn.Parameter(torch.ones(1)))
-
-        # Learnable threshold (init to 0.7 * std as in paper)
-        weight_std = self.weight.data.std()
-        self.register_parameter("delta", nn.Parameter(torch.ones(1) * 0.7 * weight_std))
+        # Initialize scales and threshold per TTQ paper (Zhu et al., ICLR 2017)
+        # Eq. 2: threshold = 0.7 * E[|W|], scales = E[|W|]
+        weight_mean_abs = self.weight.data.abs().mean()
+        self.register_parameter("wp", nn.Parameter(torch.ones(1) * weight_mean_abs))
+        self.register_parameter("wn", nn.Parameter(torch.ones(1) * weight_mean_abs))
+        self.register_parameter("delta", nn.Parameter(torch.ones(1) * 0.7 * weight_mean_abs))
 
     def forward(self, x: Tensor) -> Tensor:
         # Activation quantization (same as BitNet)

bitnet/nn/ttq_linear.py

@@ -26,13 +26,12 @@ def __init__(
         super().__init__(in_features, out_features, bias)
         self.num_bits = num_bits
 
-        # Learnable positive/negative scales (init to 1.0)
-        self.register_parameter("wp", nn.Parameter(torch.ones(1)))
-        self.register_parameter("wn", nn.Parameter(torch.ones(1)))
-
-        # Learnable threshold (init to 0.7 * std as in paper)
-        weight_std = self.weight.data.std()
-        self.register_parameter("delta", nn.Parameter(torch.ones(1) * 0.7 * weight_std))
+        # Initialize scales and threshold per TTQ paper (Zhu et al., ICLR 2017)
+        # Eq. 2: threshold = 0.7 * E[|W|], scales = E[|W|]
+        weight_mean_abs = self.weight.data.abs().mean()
+        self.register_parameter("wp", nn.Parameter(torch.ones(1) * weight_mean_abs))
+        self.register_parameter("wn", nn.Parameter(torch.ones(1) * weight_mean_abs))
+        self.register_parameter("delta", nn.Parameter(torch.ones(1) * 0.7 * weight_mean_abs))
 
     def forward(self, x: Tensor) -> Tensor:
         # Activation quantization (same as BitNet)

tests/test_ttq_layers.py

@@ -33,13 +33,14 @@ def test_gradient_flows(self) -> None:
         # with classification loss, delta gets gradients through the loss.
 
     def test_parameters_initialized_properly(self) -> None:
-        """TTQ parameters should be initialized to reasonable values."""
+        """TTQ parameters should be initialized per paper (Eq. 2)."""
         layer = TTQLinear(64, 32)
-        # wp and wn should be initialized to 1.0
-        assert torch.allclose(layer.wp, torch.ones(1))
-        assert torch.allclose(layer.wn, torch.ones(1))
-        # delta should be initialized to 0.7 * weight.std()
-        assert layer.delta > 0
+        weight_mean_abs = layer.weight.data.abs().mean()
+        # wp and wn should be initialized to E[|W|]
+        assert torch.allclose(layer.wp, weight_mean_abs, rtol=1e-5)
+        assert torch.allclose(layer.wn, weight_mean_abs, rtol=1e-5)
+        # delta should be initialized to 0.7 * E[|W|]
+        assert torch.allclose(layer.delta, 0.7 * weight_mean_abs, rtol=1e-5)
 
     def test_numerical_stability_during_training(self) -> None:
         """Training should not produce NaN losses."""
@@ -94,13 +95,14 @@ def test_gradient_flows(self) -> None:
         # with classification loss, delta gets gradients through the loss.
 
     def test_parameters_initialized_properly(self) -> None:
-        """TTQ parameters should be initialized to reasonable values."""
+        """TTQ parameters should be initialized per paper (Eq. 2)."""
         layer = TTQConv2d(3, 16, kernel_size=3)
-        # wp and wn should be initialized to 1.0
-        assert torch.allclose(layer.wp, torch.ones(1))
-        assert torch.allclose(layer.wn, torch.ones(1))
-        # delta should be initialized to 0.7 * weight.std()
-        assert layer.delta > 0
+        weight_mean_abs = layer.weight.data.abs().mean()
+        # wp and wn should be initialized to E[|W|]
+        assert torch.allclose(layer.wp, weight_mean_abs, rtol=1e-5)
+        assert torch.allclose(layer.wn, weight_mean_abs, rtol=1e-5)
+        # delta should be initialized to 0.7 * E[|W|]
+        assert torch.allclose(layer.delta, 0.7 * weight_mean_abs, rtol=1e-5)
 
     def test_numerical_stability_during_training(self) -> None:
         """Training should not produce NaN losses."""