Make lora initializers configurable.

Author: fchollet

keras/layers/core/dense.py

@@ -146,7 +146,9 @@ def compute_output_shape(self, input_shape):
         output_shape[-1] = self.units
         return tuple(output_shape)
 
-    def enable_lora(self, rank):
+    def enable_lora(
+        self, rank, a_initializer="he_uniform", b_initializer="zeros"
+    ):
         if self.kernel_constraint:
             raise ValueError(
                 "Lora is incompatible with kernel constraints. "
@@ -166,11 +168,11 @@ def enable_lora(self, rank):
         self.lora_kernel_a = self.add_weight(
             name="lora_kernel_a",
             shape=(self.kernel.shape[0], rank),
-            initializer="zeros",
+            initializer=initializers.get(a_initializer),
             regularizer=self.kernel_regularizer,
         )
         self.lora_kernel_b = self.add_weight(
-            name="lora_kernel_b",
+            name=initializers.get(b_initializer),
             shape=(rank, self.kernel.shape[1]),
             initializer="zeros",
             regularizer=self.kernel_regularizer,

keras/layers/core/einsum_dense.py

@@ -222,7 +222,9 @@ def call(self, inputs):
             x = self.activation(x)
         return x
 
-    def enable_lora(self, rank):
+    def enable_lora(
+        self, rank, a_initializer="he_uniform", b_initializer="zeros"
+    ):
         if self.kernel_constraint:
             raise ValueError(
                 "Lora is incompatible with kernel constraints. "
@@ -243,13 +245,13 @@ def enable_lora(self, rank):
         self.lora_kernel_a = self.add_weight(
             name="lora_kernel_a",
             shape=(self.kernel.shape[:-1] + (rank,)),
-            initializer="zeros",
+            initializer=initializers.get(a_initializer),
             regularizer=self.kernel_regularizer,
         )
         self.lora_kernel_b = self.add_weight(
             name="lora_kernel_b",
             shape=(rank, self.kernel.shape[-1]),
-            initializer="zeros",
+            initializer=initializers.get(b_initializer),
             regularizer=self.kernel_regularizer,
         )
         self.kernel.trainable = False

keras/layers/core/embedding.py

@@ -125,7 +125,9 @@ def compute_mask(self, inputs, mask=None):
     def compute_output_shape(self, input_shape):
         return input_shape + (self.output_dim,)
 
-    def enable_lora(self, rank):
+    def enable_lora(
+        self, rank, a_initializer="he_uniform", b_initializer="zeros"
+    ):
         if self.embeddings_constraint:
             raise ValueError(
                 "Lora is incompatible with embedding constraints. "
@@ -145,13 +147,13 @@ def enable_lora(self, rank):
         self.lora_embeddings_a = self.add_weight(
             name="lora_embeddings_a",
             shape=(self.embeddings.shape[0], rank),
-            initializer="zeros",
+            initializer=initializers.get(a_initializer),
             regularizer=self.embeddings_regularizer,
         )
         self.lora_embeddings_b = self.add_weight(
             name="lora_embeddings_b",
             shape=(rank, self.embeddings.shape[1]),
-            initializer="zeros",
+            initializer=initializers.get(b_initializer),
             regularizer=self.embeddings_regularizer,
         )
         self.embeddings.trainable = False