Create model.py

diff_lora/model.py

@@ -0,0 +1,83 @@
+import re
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class DiffLoRALinear(nn.Module):
+    """
+    Fused DiffLoRALinear implements a differential low-rank adapter:
+      Δy = (α/r) * [A_pos @ B_pos - τ * (A_neg @ B_neg)]
+    The fused version computes:
+      update = x_dropped @ concat(A_pos, A_neg) @ concat(B_pos, -τ * B_neg)
+    This version explicitly moves τ to the same device as the input.
+    """
+    def __init__(self, in_features: int, out_features: int, r: int = 8,
+                 lora_alpha: float = 16.0, dropout: float = 0.0,
+                 merge_weights: bool = False, init_method: str = "kaiming"):
+        super().__init__()
+        # Base linear layer (frozen)
+        self.linear = nn.Linear(in_features, out_features, bias=False)
+        self.linear.weight.requires_grad = False
+        self.in_features = in_features
+        self.out_features = out_features
+        self.r = r
+        self.scaling = lora_alpha / r
+        self.merge_weights = merge_weights
+        self.merged = False
+        self.lora_dropout = nn.Dropout(dropout) if dropout > 0.0 else nn.Identity()
+        # Low-rank parameters for positive and negative components
+        self.A_pos = nn.Parameter(torch.zeros(in_features, r))
+        self.B_pos = nn.Parameter(torch.zeros(r, out_features))
+        self.A_neg = nn.Parameter(torch.zeros(in_features, r))
+        self.B_neg = nn.Parameter(torch.zeros(r, out_features))
+        self.tau = nn.Parameter(torch.tensor(1.0))  # Scalar parameter
+        self.reset_parameters(init_method)
+
+    def reset_parameters(self, init_method: str = "kaiming"):
+        if init_method == "kaiming":
+            nn.init.kaiming_uniform_(self.A_pos, a=math.sqrt(5))
+            nn.init.zeros_(self.B_pos)
+            nn.init.kaiming_uniform_(self.A_neg, a=math.sqrt(5))
+            nn.init.zeros_(self.B_neg)
+        elif init_method == "xavier":
+            nn.init.xavier_uniform_(self.A_pos)
+            nn.init.zeros_(self.B_pos)
+            nn.init.xavier_uniform_(self.A_neg)
+            nn.init.zeros_(self.B_neg)
+        else:
+            raise ValueError(f"Unknown init_method: {init_method}")
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.merge_weights and self.merged:
+            return self.linear(x)
+        base_out = self.linear(x)
+        x_dropped = self.lora_dropout(x)
+        # Ensure tau is on the same device as x
+        tau = self.tau.to(x.device)
+        # Concatenate positive and negative parameters along the rank dimension
+        combined_A = torch.cat([self.A_pos, self.A_neg], dim=1)  # (in_features, 2*r)
+        combined_B = torch.cat([self.B_pos, -tau * self.B_neg], dim=0)  # (2*r, out_features)
+        update = x_dropped @ combined_A @ combined_B
+        delta = self.scaling * update
+        return base_out + delta
+
+def replace_linear_with_diff_lora(module: nn.Module, target_regex: str, r: int):
+    """
+    Recursively replace nn.Linear modules whose names match target_regex
+    with DiffLoRALinear modules using rank r.
+    """
+    for name, child in module.named_children():
+        if isinstance(child, nn.Linear) and re.search(target_regex, name, re.IGNORECASE):
+            new_layer = DiffLoRALinear(
+                in_features=child.in_features,
+                out_features=child.out_features,
+                r=r,
+                lora_alpha=16.0,
+                dropout=0.1,
+                merge_weights=False,
+            )
+            new_layer.linear.weight.data.copy_(child.weight.data)
+            setattr(module, name, new_layer)
+        else:
+            replace_linear_with_diff_lora(child, target_regex, r)