論文公開から数日でb1.58を早くも実装した方が現れる。

コードの変更だけを確認。

$ git diff 6198f9cd9d37ac45e37fa9f8689a289dc0f724e0 \
-- ':(exclude)pyproject.toml'\
':(exclude)poetry.lock' > bitnet_change.patch

binarized -> ternarized の変更で変更箇所が増えるので手直ししてみる。
binarizedではなくternarizedだと言いたい気持ちを感じる。

ここの変更も無しでいいかな。

def quantize_activations(self, x: torch.Tensor, b: int = 8):

def quantize_activations(self, x, b=8):

下記が一部調整したDIFF。
ほんとに変更少なく実装できたんですね。凄い。

diff --git a/bitnet_llama/module.py b/bitnet_llama/module.py
index 1b920f2..a2069da 100644
--- a/bitnet_llama/module.py
+++ b/bitnet_llama/module.py
@@ -13,9 +13,7 @@ class BitLinearNaive(nn.Linear):
         self.eps = 1e-5  # Small epsilon value to avoid division by zero and overflow
 
     def binarize_weights(self):
-        alpha = self.weight.mean()
-        binarized_weights = torch.sign(self.weight - alpha)
-        return binarized_weights
+        return BitLinear.quantize_tensor(self.weight, self.eps)
 
     def quantize_activations(self, x, b=8):
         Q_b = 2 ** (b - 1)
@@ -25,15 +23,6 @@ class BitLinearNaive(nn.Linear):
         )
         return quantized_x
 
-    def scale_activations(self, x, b=8):
-        Q_b = 2 ** (b - 1)
-        eta = x.min()
-        gamma = x.abs().max()
-        scaled_x = torch.clamp(
-            (x - eta) * Q_b / (gamma + self.eps), self.eps, Q_b - self.eps
-        )
-        return scaled_x
-
     def forward(self, input):
         # Binarize weights
         binarized_weights = self.binarize_weights()
@@ -44,10 +33,6 @@ class BitLinearNaive(nn.Linear):
         # Quantize activations (before non-linear functions like ReLU)
         output = self.quantize_activations(output)
 
-        # For the sake of demonstration, we'll also include the scaling step.
-        # In practice, this would be done before a non-linear function in a forward pass.
-        output = self.scale_activations(output)
-
         return output
 
 
@@ -57,9 +42,17 @@ class BitLinear(nn.Linear):
         self.num_groups = num_groups
         self.eps = 1e-5
 
+    @staticmethod
+    def quantize_tensor(x: torch.Tensor, eps: float):
+        return torch.clamp(
+            torch.round(x / (x.mean() + eps)),
+            -1,
+             1,
+        )
+
     def ste_binarize(self, x):
         # Apply the sign function for binarization
-        binarized_x = torch.sign(x)
+        binarized_x = BitLinear.quantize_tensor(x, self.eps)
         # Use STE: during backward pass, we bypass the binarization
         binarized_x = (binarized_x - x).detach() + x
         return binarized_x
@@ -74,11 +67,8 @@ class BitLinear(nn.Linear):
             end_idx = (g + 1) * group_size
             weight_group = self.weight[start_idx:end_idx]
 
-            # Binarize each group using STE
-            alpha_g = weight_group.mean()
-            binarized_weights[start_idx:end_idx] = self.ste_binarize(
-                weight_group - alpha_g
-            )
+            # Binarize weights
+            binarized_weights[start_idx:end_idx] = self.ste_binarize(weight_group)
 
         return binarized_weights
 
@@ -125,7 +115,7 @@ class BitLinearOptimized(nn.Linear):
 
         # Initialize 1-bit quantized weights and store them as int8
         self.register_buffer(
-            "quantized_weights", torch.sign(self.weight.data).to(torch.int8)
+            "quantized_weights", BitLinear.quantize_tensor(self.weight.data, self.eps).to(torch.int8)
         )
         # Clear the original weights to save memory
         del self.weight
@@ -138,7 +128,7 @@ class BitLinearOptimized(nn.Linear):
     @weight.setter
     def weight(self, value):
         # Update the quantized_weights when the weight property is set
-        self.quantized_weights.data = torch.sign(value).to(torch.int8)
+        self.quantized_weights.data = BitLinear.quantize_tensor(value, self.eps)
 
     def dequantize_weights(self):
         # Convert quantized_weights back to bfloat16 and compute alpha for the weights
@@ -148,7 +138,7 @@ class BitLinearOptimized(nn.Linear):
 
     def ste_binarize(self, x):
         # Apply the sign function for binarization
-        binarized_x = torch.sign(x)
+        binarized_x = BitLinear.quantize_tensor(x, self.eps)
         # Use STE: during backward pass, we bypass the binarization
         binarized_x = (binarized_x - x).detach() + x
         return binarized_x
@@ -167,10 +157,7 @@ class BitLinearOptimized(nn.Linear):
             weight_group = weights[start_idx:end_idx]
 
             # Binarize each group using STE
-            alpha_g = weight_group.mean()
-            binarized_weights[start_idx:end_idx] = self.ste_binarize(
-                weight_group - alpha_g
-            )
+            binarized_weights[start_idx:end_idx] = self.ste_binarize(weight_group)
 
         return binarized_weights