[OFT] Rework and Matrix Exponential Mode

modules/module/LoRAModule.py

@@ -578,15 +578,19 @@ class OFTModule(PeftBase):
     oft_block_size: int
     block_share: bool
     oft_scaled: bool
+    oft_clipped_norm: float | None
+    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_clipped_norm: float | None, 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_clipped_norm = oft_clipped_norm
+        self.oft_cans = oft_cans
         self.dropout_probability = kwargs.pop('dropout_probability', 0.0)
         self.oft_R = None
         self.adjustment_info = None
@@ -656,6 +660,8 @@ def initialize_weights(self):
             use_cayley_neumann=True,
             num_cayley_neumann_terms=5,
             dropout_probability=self.dropout_probability,
+            oft_clipped_norm=self.oft_clipped_norm,
+            oft_cans=self.oft_cans,
         )
 
         nn.init.zeros_(self.oft_R.weight)
@@ -672,8 +678,8 @@ def forward(self, x, *args, **kwargs):
         effective_weight = self.oft_R.weight / scaling_factor
 
         # 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
+        orth_rotate = self.oft_R._compute_orthogonal_matrix(
+            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 +870,8 @@ def __init__(
                 config.oft_block_size,
                 config.oft_block_share,
                 config.oft_scaled,
+                config.oft_clipped_norm,
+                config.oft_cans,
             ]
             self.additional_kwargs = {
                 'dropout_probability': config.dropout_probability,

modules/module/oft_utils.py

@@ -47,7 +47,9 @@ def __init__(
         oft_scaled=False,
         use_cayley_neumann=True,
         num_cayley_neumann_terms=5,
+        oft_cans=False,
         dropout_probability=0.0,
+        oft_clipped_norm: float | None = 0.95,
     ):
         super().__init__()
         self.r = r
@@ -62,14 +64,38 @@ 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_exp", torch.tensor([]))
         # 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)
+        if oft_clipped_norm == -1:
+            if oft_cans:
+                # 0.95 * pi (~3.11) avoids the gradient ambiguity/singularity
+                # at exactly 180 degrees (pi).
+                self.oft_clipped_norm = 0.95 * math.pi
+            elif use_cayley_neumann:
+                # Neumann series diverges if norm >= 1.0, so 0.95 is the safe max.
+                self.oft_clipped_norm = 0.95
+        else:
+            self.oft_clipped_norm = oft_clipped_norm
+        if oft_clipped_norm is not None:
+            self.register_buffer("clipped_oft", torch.tensor(self.oft_clipped_norm))
+            # Initialize states for Spectral Normalization via Power Iteration
+            u = torch.randn(r, block_size)
+            u = u / u.norm(dim=1, keepdim=True).clamp_min(1e-12)
+            self.register_buffer("u_state", u, persistent=False)
+            v = torch.randn(r, block_size)
+            v = v / v.norm(dim=1, keepdim=True).clamp_min(1e-12)
+            self.register_buffer("v_state", v, persistent=False)
 
     def _pytorch_skew_symmetric(self, vec, block_size):
         batch_size = vec.shape[0]
@@ -88,18 +114,131 @@ def _pytorch_skew_symmetric_inv(self, matrix, block_size):
         vec = matrix[:, self.rows, self.cols]
         return vec
 
-    def _cayley_batch(
-        self, Q: torch.Tensor, block_size: int, use_cayley_neumann: bool = True, num_neumann_terms: int = 5
+    def _break_inductor_graph(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Acts as an opaque boundary for TorchInductor. Forces materialization
+        of the tensor to sever deeply nested AST trees in torch.compile.
+        """
+        return torch.bmm(x, torch.ones_like(x))
+
+    @torch.no_grad()
+    def _spectral_norm(self, Q_skew):
+        u = self.u_state.unsqueeze(-1).to(Q_skew.dtype)
+        v = self.v_state.unsqueeze(-1).to(Q_skew.dtype)
+        # Update v (Right Singular Vector)
+        v_raw = torch.bmm(Q_skew.mT, u)
+        v_norm = torch.linalg.vector_norm(v_raw, dim=1, keepdim=True)
+        candidate_v = v_raw / v_norm.clamp_min(1e-8)
+        next_v = torch.where(v_norm >= 1e-6, candidate_v, v)
+        # Update u (Left Singular Vector)
+        u_raw = torch.bmm(Q_skew, next_v)
+        u_norm = torch.linalg.vector_norm(u_raw, dim=1, keepdim=True)
+        candidate_u = u_raw / u_norm.clamp_min(1e-8)
+        next_u = torch.where(u_norm >= 1e-6, candidate_u, u)
+        if self.training:
+            self.v_state.copy_(next_v.squeeze(-1))
+            self.u_state.copy_(next_u.squeeze(-1))
+        return next_v, next_u
+
+    def _cans_newton_schulz_iteration(
+        self,
+        G: torch.Tensor,
+        steps: int = 3,
+        eps: float = 1e-7,
+    ) -> torch.Tensor:
+        """
+        Chebyshev-Optimized Newton-Schulz iteration with a dynamically computed Chebyshev lower bound.
+        """
+        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
+
+        # The 4th-order Taylor expansion of exp(Q) has a minimum singular value of exactly 0.5
+        # (occurring at ||Q|| = sqrt(6) ~ 2.449). Therefore, the min_singular_value of normalized X
+        # is guaranteed to be >= 0.5 / g_norm.
+        lower_bound = 0.5 / 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
+            eps_val = self._break_inductor_graph(eps_val)
+
+            lower_bound = 1 - eps_val
+            upper_bound = 1 + eps_val
+
+        return X.to(original_dtype)
+
+    def _matrix_exp_cans(self, Q_skew: torch.Tensor) -> torch.Tensor:
+        """
+        Approximates the Matrix Exponential using a 4th-order Taylor expansion,
+        Scaling & Squaring, and Chebyshev-Optimized Newton-Schulz (CANS).
+        """
+        num_squarings = 2
+        id_mat = self.id_mat
+
+        # Scaling step
+        Q_scaled = Q_skew / (2 ** num_squarings)
+        Q_squared = torch.bmm(Q_scaled, Q_scaled)
+
+        # 4th-order Taylor expansion: exp(Q) ≈ I + Q + Q^2/2 + Q^3/6 + Q^4/24
+        # Factored to minimize matrix multiplications: (I + Q) + Q^2 * (0.5*I + 1/6*Q + 1/24*Q^2)
+        taylor_higher_order = 0.5 * id_mat + (1.0 / 6.0) * Q_scaled + (1.0 / 24.0) * Q_squared
+        G = torch.baddbmm(id_mat + Q_scaled, Q_squared, taylor_higher_order)
+
+        # Orthogonalize the approximation (CANS)
+        # Empirically, CANS requires 3 steps to converge
+        R = self._cans_newton_schulz_iteration(G=G, steps=3)
+
+        # Squaring step to recover full rotation
+        for _ in range(num_squarings):
+            R = torch.bmm(R, R)
+
+        # Final standard Newton-Schulz step to correct drift caused by squaring in lower precision
+        # R_new = R + 0.5 * R * (I - R^T R)
+        residual = torch.baddbmm(id_mat, R.mT, R, beta=1.0, alpha=-1.0)
+        R = torch.baddbmm(R, R, residual, beta=1.0, alpha=0.5)
+
+        return R
+
+    def _compute_orthogonal_matrix(
+        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.
+        Converts learned weights into a batch of orthogonal matrices.
         """
         b, _ = Q.shape
         previous_dtype = Q.dtype
 
         Q_skew = self._pytorch_skew_symmetric(Q, block_size)
 
-        if use_cayley_neumann:
+        # Spectral Normalization / Clipping
+        if self.oft_clipped_norm is not None:
+            v_norm, u_norm = self._spectral_norm(Q_skew)
+            u_raw_grad = torch.bmm(Q_skew, v_norm)
+            sigma = torch.sum(u_norm * u_raw_grad, dim=1, keepdim=True)
+            max_norm = self.oft_clipped_norm
+            Q_skew = Q_skew * (max_norm / torch.clamp(sigma, min=max_norm))
+
+        if oft_cans:
+            R = self._matrix_exp_cans(Q_skew)
+        elif use_cayley_neumann:
             R = torch.eye(block_size, device=Q.device, dtype=Q.dtype).repeat(b, 1, 1)
             if num_neumann_terms > 1:
                 R.add_(Q_skew, alpha=2.0)
@@ -114,12 +253,7 @@ def _cayley_batch(
                     Q_power = torch.bmm(Q_power, Q_skew)
                     R.add_(Q_power)
         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)
 
@@ -130,11 +264,18 @@ def forward(self, x):
 
         orig_shape = x.shape
 
-        scaling_factor = 2 * math.sqrt(self.block_size - 1) if self.oft_scaled else 1
-        effective_weight = self.weight / scaling_factor
+        if self.oft_scaled:
+            # Cayley has a 2x gradient multiplier (I + 2Q). Exp has a 1x multiplier (I + Q).
+            # We drop the 2 for Exp/CANS to maintain consistent effective learning rates.
+            is_cayley = self.use_cayley_neumann and not self.oft_cans
+            multiplier = 2.0 if is_cayley else 1.0
+            scaling_factor = multiplier * math.sqrt(self.block_size - 1)
+            effective_weight = self.weight / scaling_factor
+        else:
+            effective_weight = self.weight
 
-        orth_rotate = self._cayley_batch(
-            effective_weight, self.block_size, self.use_cayley_neumann, self.num_cayley_neumann_terms
+        orth_rotate = self._compute_orthogonal_matrix(
+            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,16 @@ 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, "Matrix Exponential CANS",
+                             tooltip="Replaces Cayley-Neumann with Matrix Exponential with Chebyshev-Optimized Newton-Schulz (CANS) to improve orthogonalization stability.")
+            components.switch(master, 4, 4, self.ui_state, "oft_cans")
+
+            # Clip OFT max norm
+            components.label(master, 5, 0, "Spectral Norm Clipping",
+                             tooltip="Strictly clips the spectral norm of the OFT matrix to guarantee convergence of the Cayley parametrization (requires norm <= 1.0). Smaller values constrain the learned rotation to stay near the identity matrix, limiting adaptation. Default: 1.0 (e.g. 0.8 = 80% of maximum expressiveness). Leave empty to disable.")
+            components.entry(master, 5, 1, self.ui_state, "oft_clipped_norm")
+
             # 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/ui/MuonAdamWindow.py

@@ -73,14 +73,13 @@ def create_adam_params_ui(self, master):
             'use_bias_correction': {'title': 'Bias Correction', 'tooltip': 'Turn on Adam\'s bias correction.', 'type': 'bool'},
             'weight_decay': {'title': 'Weight Decay', 'tooltip': 'Regularization to prevent overfitting.', 'type': 'float'},
             'use_orthograd': {'title': 'use_orthograd', 'tooltip': 'Use orthograd method', 'type': 'bool'},
-            'nnmf_factor': {'title': 'Factored Optimizer', 'tooltip': 'Enables a memory-efficient mode by applying fast low-rank factorization to the optimizers states. It combines factorization for magnitudes with 1-bit compression for signs, drastically reducing VRAM usage and allowing for larger models or batch sizes. This is an approximation which may slightly alter training dynamics.', 'type': 'bool'},
-            'orthogonal_gradient': {'title': 'OrthoGrad', 'tooltip': 'Reduces overfitting by removing the gradient component parallel to the weight, thus improving generalization.', 'type': 'bool'},
+            'orthogonal_gradient': {'title': 'OrthoGrad', 'tooltip': 'Reduces overfitting by removing the gradient component parallel to the weight, thus improving generalization. This has two modes: 1. flattened: Standard vectorized OrthoGrad. Fastest, but loses the structural properties of matrices. 2. iterative: Matrix-wise OrthoGrad, preserves structure by iteratively projecting rows and columns.', 'type': 'OrthoGrad'},
             'use_atan2': {'title': 'Atan2 Scaling', 'tooltip': 'A robust replacement for eps, which also incorporates gradient clipping, bounding and stabilizing the optimizer updates.', 'type': 'bool'},
-            'use_AdEMAMix': {'title': 'AdEMAMix EMA', 'tooltip': 'Adds a second, slow-moving EMA, which is combined with the primary momentum to stabilize updates, and accelerate the training.', 'type': 'bool'},
-            'beta3_ema': {'title': 'Beta3 EMA', 'tooltip': 'Coefficient for slow-moving EMA of AdEMAMix.', 'type': 'float'},
-            'Simplified_AdEMAMix': {'title': 'Simplified AdEMAMix', 'tooltip': "Enables a simplified, single-EMA variant of AdEMAMix. Instead of blending two moving averages (fast and slow momentum), this version combines the raw current gradient (controlled by 'Grad α') directly with a single theory-based momentum. This makes the optimizer highly responsive to recent gradient information, which can accelerate training in all batch size scenarios when tuned correctly.", 'type': 'bool'},
-            'alpha_grad': {'title': 'Grad α', 'tooltip': 'Controls the mixing coefficient between raw gradients and momentum gradients in Simplified AdEMAMix. Higher values (e.g., 10-100) emphasize recent gradients, suitable for small batch sizes to reduce noise. Lower values (e.g., 0-1) emphasize historical gradients, suitable for large batch sizes for stability. Setting to 0 uses only momentum gradients without raw gradient contribution.', 'type': 'float'},
             'kourkoutas_beta': {'title': 'Kourkoutas Beta', 'tooltip': 'Enables a layer-wise dynamic β₂ adaptation. This feature makes the optimizer more responsive to "spiky" gradients by lowering β₂ during periods of high variance, and more stable during calm periods by raising β₂ towards its maximum. It can significantly improve training stability and final loss.', 'type': 'bool'},
+            'nesterov_coef': {'title': 'Nesterov Coef', 'tooltip': 'Controls the mixing coefficient between momentum gradients and raw gradients in Nesterov momentum. For a factor of 0.8, the final update will be 80% of the momentum gradients and 20% raw gradient. Leaving it unset toggles the standard Nestrov behavior (where nesterov_coef = beta1 or momentum). Setting it to 0 cancels momentum contribution.', 'type': 'float'},
+            'factored_2nd': {'title': 'Factored 2nd', 'tooltip': 'Whether to keep the first moment uncompressed (dense), while only factorizing the second moment. This makes the optimizer highly robust to a wide range of LRs, mimicking high-order optimization.', 'type': 'bool'},
+            'fisher_wd': {'title': 'Fisher Weight Decay', 'tooltip': 'Applies adaptive, scale-invariant weight-decay regularization based on the Fisher Information Matrix (approximated by Adam\'s second moment). It reduces penalty for "important" high-curvature weights while accelerating decay for "useless" weights in flat regions. Leading to improved convergence and better final performance.', 'type': 'bool'},
+            'state_precision': {'title': 'state_precision', 'tooltip': """The quantization format used to store the optimizer states to save VRAM. Options include: 'auto': Stores the states in the original parameter's precision. 'factored': Enables a memory-efficient mode by applying fast low-rank factorization to the optimizers states. It combines factorization for magnitudes with 1-bit compression for signs, drastically reducing VRAM usage and allowing for larger models or batch sizes. 'fp32': Uses full FP32. 'bf16_sr': Uses BF16 with stochastic rounding for a balance of precision and memory. 'int8_sr': Uses 8-bit block-wise quantization with stochastic rounding.""", 'type': 'StatePrecision'},
         }
         # @formatter:on
 
@@ -103,5 +102,9 @@ def create_adam_params_ui(self, master):
 
             if param_type != 'bool':
                 components.entry(master, row, col + 1, self.adam_ui_state, key)
+            elif param_type == 'StatePrecision':
+                components.options(master, row, col + 1, ["auto", "factored", "fp32", "bf16_sr", "int8_sr"], self.adam_ui_state, key)
+            elif param_type == 'OrthoGrad':
+                components.options(master, row, col + 1, ["disabled", "flattened", "iterative"], self.adam_ui_state, key)
             else:
                 components.switch(master, row, col + 1, self.adam_ui_state, key)