[OFT] Chebyshev-Optimized Newton-Schulz (CANS): Faster and Better

modules/module/LoRAModule.py

@@ -578,15 +578,17 @@ class OFTModule(PeftBase):
     oft_block_size: int
     block_share: bool
     oft_scaled: bool
+    oft_cans: bool
     dropout_probability: float
     adjustment_info: tuple[int, int] | None # for reporting
 
-    def __init__(self, prefix: str, orig_module: nn.Module | None, oft_block_size: int, block_share: bool, oft_scaled: bool, **kwargs):
+    def __init__(self, prefix: str, orig_module: nn.Module | None, oft_block_size: int, block_share: bool, oft_scaled: bool, oft_cans: bool, **kwargs):
         super().__init__(prefix, orig_module)
         self.oft_block_size = oft_block_size
         self.rank = 0
         self.block_share = block_share
         self.oft_scaled = oft_scaled
+        self.oft_cans = oft_cans
         self.dropout_probability = kwargs.pop('dropout_probability', 0.0)
         self.oft_R = None
         self.adjustment_info = None
@@ -656,6 +658,7 @@ def initialize_weights(self):
             use_cayley_neumann=True,
             num_cayley_neumann_terms=5,
             dropout_probability=self.dropout_probability,
+            oft_cans=self.oft_cans,
         )
 
         nn.init.zeros_(self.oft_R.weight)
@@ -673,7 +676,7 @@ def forward(self, x, *args, **kwargs):
 
         # For Conv2d, we must rotate the weights, not the input, to preserve spatial information.
         orth_rotate = self.oft_R._cayley_batch(
-            effective_weight, self.oft_R.block_size, self.oft_R.use_cayley_neumann, self.oft_R.num_cayley_neumann_terms
+            effective_weight, self.oft_R.block_size, self.oft_R.use_cayley_neumann, self.oft_R.num_cayley_neumann_terms, self.oft_R.oft_cans
         )
         orth_rotate = self.oft_R.dropout(orth_rotate)
 
@@ -864,6 +867,7 @@ def __init__(
                 config.oft_block_size,
                 config.oft_block_share,
                 config.oft_scaled,
+                config.oft_cans,
             ]
             self.additional_kwargs = {
                 'dropout_probability': config.dropout_probability,

modules/module/oft_utils.py

@@ -47,6 +47,7 @@ def __init__(
         oft_scaled=False,
         use_cayley_neumann=True,
         num_cayley_neumann_terms=5,
+        oft_cans=False,
         dropout_probability=0.0,
     ):
         super().__init__()
@@ -62,14 +63,19 @@ def __init__(
             # allowing inference tools to automatically detect scaled oft.
             self.register_buffer("scaled_oft", torch.tensor(True))
         self.oft_scaled = oft_scaled
-        self.use_cayley_neumann = use_cayley_neumann
+        self.use_cayley_neumann = use_cayley_neumann and not oft_cans
         self.num_cayley_neumann_terms = num_cayley_neumann_terms
+        self.oft_cans = oft_cans
+        if oft_cans:
+            self.register_buffer("cans_oft", torch.tensor(True))
         # Create indices for upper triangle (excluding diagonal)
         rows, cols = torch.triu_indices(block_size, block_size, 1)
         self.register_buffer("rows", rows, persistent=False)
         self.register_buffer("cols", cols, persistent=False)
         self.dropout = MultiplicativeDropoutLayer(p=dropout_probability)
-
+        if not self.use_cayley_neumann:
+            id_mat = (torch.eye(block_size).unsqueeze(0).expand(r, block_size, block_size))
+            self.register_buffer("id_mat", id_mat, persistent=False)
 
     def _pytorch_skew_symmetric(self, vec, block_size):
         batch_size = vec.shape[0]
@@ -88,8 +94,57 @@ def _pytorch_skew_symmetric_inv(self, matrix, block_size):
         vec = matrix[:, self.rows, self.cols]
         return vec
 
+    def _cans_newton_schulz_iteration(
+        self,
+        G: torch.Tensor,
+        steps: int = 7,
+        eps: float = 1e-7,
+    ) -> torch.Tensor:
+        """
+        Chebyshev-Optimized Newton-Schulz iteration with a dynamically computed Chebyshev lower bound.
+        Optimized for G = I + Q (where Q is skew-symmetric).
+        """
+        original_dtype = G.dtype
+        X = G
+
+        # Max row sum is guaranteed to be >= the maximum singular value of X.
+        g_norm = X.abs().sum(dim=-1, keepdim=True).amax(dim=-2, keepdim=True).clamp_min(eps).detach()
+        X = X / g_norm
+
+        # Since min_singular_value(I + Q) >= 1, the min_singular_value of normalized X
+        # is guaranteed to be >= 1 / ||G||_F.
+        lower_bound = (1.0 / g_norm)
+        upper_bound = 1
+
+        for _ in range(steps):
+            lb, ub = lower_bound, upper_bound
+            lb_ub = lb * ub
+            e_sq = (lb**2 + lb_ub + ub**2) / 3.0
+            K = 2.0 * e_sq**1.5
+            L = lb_ub * (lb + ub)
+            denom = K + L
+            alpha = 6.0 / denom
+            c1 = alpha * e_sq
+            c3 = -alpha / 3.0
+
+            A = torch.bmm(X, X.mT)
+            X = c1 * X + c3 * torch.bmm(A, X)
+
+            # Dynamically update bounds for the next step
+            eps_val = (K - L) / denom
+
+            # bmm acts as an opaque boundary, forcing Inductor to
+            # materialize eps_val and severing the exponential AST tree.
+            # Shape is (B, 1, 1), so ones_like acts as an identity.
+            eps_val = torch.bmm(eps_val, torch.ones_like(eps_val))
+
+            lower_bound = 1 - eps_val
+            upper_bound = 1 + eps_val
+
+        return X.to(original_dtype)
+
     def _cayley_batch(
-        self, Q: torch.Tensor, block_size: int, use_cayley_neumann: bool = True, num_neumann_terms: int = 5
+        self, Q: torch.Tensor, block_size: int, use_cayley_neumann: bool = True, num_neumann_terms: int = 5, oft_cans: bool = False,
     ) -> torch.Tensor:
         """
         Perform the Cayley parametrization on a batch of skew-symmetric matrices.
@@ -113,13 +168,17 @@ def _cayley_batch(
                         R.add_(Q_power, alpha=2.0)
                     Q_power = torch.bmm(Q_power, Q_skew)
                     R.add_(Q_power)
+        elif oft_cans:
+            # Compute G = (I + Q)^2 = I + 2Q + Q^2
+            # Squaring the matrix doubles the rotation range and matches Cayley (I + 2Q).
+            Q_squared = torch.bmm(Q_skew, Q_skew)
+            G = self.id_mat + 2 * Q_skew + Q_squared
+            # Empirically, BF16 requires 5 steps to converge to ortho error ~1e-2 (its limit)
+            # While FP32 takes 7 steps to converge to ortho error ~1e-6
+            steps = 5 if G.dtype == torch.bfloat16 else 7
+            R = self._cans_newton_schulz_iteration(G=G, steps=steps)
         else:
-            id_mat = (
-                torch.eye(Q_skew.shape[-1], device=Q_skew.device)
-                .unsqueeze(0)
-                .expand(b, Q_skew.shape[-1], Q_skew.shape[-1])
-            )
-            R = torch.linalg.solve(id_mat + Q_skew, id_mat - Q_skew, left=False)
+            R = torch.linalg.solve(self.id_mat + Q_skew, self.id_mat - Q_skew, left=False)
 
         return R.to(previous_dtype)
 
@@ -134,7 +193,7 @@ def forward(self, x):
         effective_weight = self.weight / scaling_factor
 
         orth_rotate = self._cayley_batch(
-            effective_weight, self.block_size, self.use_cayley_neumann, self.num_cayley_neumann_terms
+            effective_weight, self.block_size, self.use_cayley_neumann, self.num_cayley_neumann_terms, self.oft_cans
         )
         orth_rotate = self.dropout(orth_rotate)
 

modules/ui/LoraTab.py

@@ -134,6 +134,11 @@ def setup_lora(self, peft_type: PeftType):
                              tooltip="Applies a scaling factor to the learned weights. This ensures that the effective learning rate remains consistent across different block sizes. Without this, different block sizes require significantly different learning rates.")
             components.switch(master, 2, 4, self.ui_state, "oft_scaled")
 
+            # CANS OFT
+            components.label(master, 4, 3, "Accelerated Newton-Schulz",
+                             tooltip="Replaces Cayley-Neumann with Chebyshev-Optimized Newton-Schulz (CANS) to improve orthogonalization stability and reduce error without the high computational cost of the exact solver.")
+            components.switch(master, 4, 4, self.ui_state, "oft_cans")
+
             # Dropout Percentage
             components.label(master, 2, 0, "Dropout Probability",
                             tooltip="Dropout probability. This percentage of the rotated adapter nodes that will be randomly restored to the base model initial statue. Helps with overfitting. 0 disables, 1 maximum.")

modules/util/config/TrainConfig.py

@@ -522,6 +522,7 @@ class TrainConfig(BaseConfig):
     oft_block_size: int
     oft_block_share: bool
     oft_scaled: bool
+    oft_cans: bool
 
     # lokr
     lokr_dim: int
@@ -1161,6 +1162,7 @@ def default_values() -> 'TrainConfig':
         data.append(("oft_block_size", 32, int, False))
         data.append(("oft_block_share", False, bool, False))
         data.append(("oft_scaled", False, bool, False))
+        data.append(("oft_cans", False, bool, False))
 
         # lokr
         data.append(("lokr_dim", 16, int, False))