Add configurable lora_alpha parameter for LoRA in multiple Keras layers (#21139)

* feat: Add alpha parameter to enable_lora

Adds an alpha scaling parameter to LoRA layers, defaulting to rank for backward compatibility.

* feat: Add lora_alpha tests to Dense, Embedding, and EinsumDense layers

* fix: Fix LoRA test failures by using ops to do numpy conversion

* fix: remove .numpy() in LoRA tests

* docs: Apply backticks to keywords per review

Updated docstrings to enclose parameters like 'alpha' and 'rank' in backticks as requested in PR review.

Author: b05505027

keras/src/layers/convolutional/base_conv.py

@@ -80,6 +80,11 @@ class BaseConv(Layer):
             computation cost of fine-tuning large dense layers.
             You can also enable LoRA on an existing layer by calling
             `layer.enable_lora(rank)`.
+        lora_alpha: Optional integer. If set, this parameter scales the
+            low-rank adaptation delta (computed as the product of two lower-rank
+            trainable matrices) during the forward pass. The delta is scaled by
+            `lora_alpha / lora_rank`, allowing you to fine-tune the strength of
+            the LoRA adjustment independently of `lora_rank`.
     """
 
     def __init__(
@@ -102,6 +107,7 @@ def __init__(
         kernel_constraint=None,
         bias_constraint=None,
         lora_rank=None,
+        lora_alpha=None,
         **kwargs,
     ):
         super().__init__(activity_regularizer=activity_regularizer, **kwargs)
@@ -124,6 +130,7 @@ def __init__(
         self.kernel_constraint = constraints.get(kernel_constraint)
         self.bias_constraint = constraints.get(bias_constraint)
         self.lora_rank = lora_rank
+        self.lora_alpha = lora_alpha if lora_alpha is not None else lora_rank
         self.lora_enabled = False
         self.input_spec = InputSpec(min_ndim=self.rank + 2)
         self.data_format = self.data_format
@@ -215,7 +222,7 @@ def build(self, input_shape):
             self.bias = None
         self.built = True
         if self.lora_rank:
-            self.enable_lora(self.lora_rank)
+            self.enable_lora(self.lora_rank, lora_alpha=self.lora_alpha)
 
     @property
     def kernel(self):
@@ -224,9 +231,9 @@ def kernel(self):
                 "You must build the layer before accessing `kernel`."
             )
         if self.lora_enabled:
-            return self._kernel + ops.matmul(
-                self.lora_kernel_a, self.lora_kernel_b
-            )
+            return self._kernel + (
+                self.lora_alpha / self.lora_rank
+            ) * ops.matmul(self.lora_kernel_a, self.lora_kernel_b)
         return self._kernel
 
     def convolution_op(self, inputs, kernel):
@@ -268,7 +275,11 @@ def compute_output_shape(self, input_shape):
         )
 
     def enable_lora(
-        self, rank, a_initializer="he_uniform", b_initializer="zeros"
+        self,
+        rank,
+        lora_alpha=None,
+        a_initializer="he_uniform",
+        b_initializer="zeros",
     ):
         if self.kernel_constraint:
             raise ValueError(
@@ -301,6 +312,7 @@ def enable_lora(
         self._tracker.lock()
         self.lora_enabled = True
         self.lora_rank = rank
+        self.lora_alpha = lora_alpha if lora_alpha is not None else rank
 
     def save_own_variables(self, store):
         # Do nothing if the layer isn't yet built
@@ -363,6 +375,7 @@ def get_config(self):
         )
         if self.lora_rank:
             config["lora_rank"] = self.lora_rank
+            config["lora_alpha"] = self.lora_alpha
         return config
 
     def _check_load_own_variables(self, store):

keras/src/layers/convolutional/conv_test.py

@@ -9,6 +9,7 @@
 from keras.src import constraints
 from keras.src import layers
 from keras.src import models
+from keras.src import ops
 from keras.src import saving
 from keras.src import testing
 
@@ -735,6 +736,51 @@ def call(self, x):
             model.conv2d.lora_kernel_a.path, "mymodel/conv2d/lora_kernel_a"
         )
 
+    @pytest.mark.requires_trainable_backend
+    def test_enable_lora_with_alpha(self):
+        # Create a `Conv2D` layer with a small kernel for simplicity.
+        layer = layers.Conv2D(filters=3, kernel_size=(2, 2), padding="valid")
+        # Use a fixed input shape: batch size 1, height=4, width=4, channels=3.
+        input_shape = (1, 4, 4, 3)
+        layer.build(input_shape)
+
+        # Set the base kernel to known, deterministic values.
+        base_kernel = np.linspace(
+            0, 1, num=np.prod(layer.kernel.shape), dtype=np.float32
+        )
+        base_kernel = base_kernel.reshape(layer.kernel.shape)
+        layer.kernel.assign(base_kernel)
+
+        # Enable LoRA with `rank`=2 and a custom `lora_alpha` value (e.g. 3.0).
+        layer.enable_lora(rank=2, lora_alpha=3.0)
+        self.assertEqual(layer.lora_rank, 2)
+        self.assertEqual(layer.lora_alpha, 3.0)
+
+        # For `Conv2D`, assume the LoRA weights have shapes:
+        #   `lora_kernel_a`: (kernel_height, kernel_width, in_channels, rank)
+        #   `lora_kernel_b`: (rank, out_channels)
+        lora_a_shape = layer.lora_kernel_a.shape
+        lora_b_shape = layer.lora_kernel_b.shape
+
+        # Assign known constant values to LoRA weights.
+        lora_a = np.full(lora_a_shape, 0.1, dtype=np.float32)
+        lora_b = np.full(lora_b_shape, 0.2, dtype=np.float32)
+        layer.lora_kernel_a.assign(lora_a)
+        layer.lora_kernel_b.assign(lora_b)
+
+        # Compute the expected delta.
+        # Flatten `lora_kernel_a` to shape (-1, `rank`),
+        # multiply with `lora_kernel_b`,
+        # then reshape to the kernel's shape.
+        scaling = 3.0 / 2  # `lora_alpha / lora_rank`
+        delta = np.matmul(lora_a.reshape(-1, 2), lora_b)
+        delta = delta.reshape(base_kernel.shape)
+        expected_effective_kernel = base_kernel + scaling * delta
+
+        # Compare the effective kernel computed via the property.
+        actual_effective_kernel = ops.convert_to_numpy(layer.kernel)
+        self.assertAllClose(actual_effective_kernel, expected_effective_kernel)
+
     @pytest.mark.requires_trainable_backend
     def test_lora_rank_argument(self):
         self.run_layer_test(

keras/src/layers/core/dense.py

@@ -57,6 +57,11 @@ class Dense(Layer):
             computation cost of fine-tuning large dense layers.
             You can also enable LoRA on an existing
             `Dense` layer by calling `layer.enable_lora(rank)`.
+        lora_alpha: Optional integer. If set, this parameter scales the
+            low-rank adaptation delta (computed as the product of two lower-rank
+            trainable matrices) during the forward pass. The delta is scaled by
+            `lora_alpha / lora_rank`, allowing you to fine-tune the strength of
+            the LoRA adjustment independently of `lora_rank`.
 
     Input shape:
         N-D tensor with shape: `(batch_size, ..., input_dim)`.
@@ -82,6 +87,7 @@ def __init__(
         kernel_constraint=None,
         bias_constraint=None,
         lora_rank=None,
+        lora_alpha=None,
         **kwargs,
     ):
         super().__init__(activity_regularizer=activity_regularizer, **kwargs)
@@ -95,6 +101,7 @@ def __init__(
         self.kernel_constraint = constraints.get(kernel_constraint)
         self.bias_constraint = constraints.get(bias_constraint)
         self.lora_rank = lora_rank
+        self.lora_alpha = lora_alpha if lora_alpha is not None else lora_rank
         self.lora_enabled = False
         self.input_spec = InputSpec(min_ndim=2)
         self.supports_masking = True
@@ -135,9 +142,9 @@ def kernel(self):
                 "You must build the layer before accessing `kernel`."
             )
         if self.lora_enabled:
-            return self._kernel + ops.matmul(
-                self.lora_kernel_a, self.lora_kernel_b
-            )
+            return self._kernel + (
+                self.lora_alpha / self.lora_rank
+            ) * ops.matmul(self.lora_kernel_a, self.lora_kernel_b)
         return self._kernel
 
     def call(self, inputs, training=None):
@@ -154,7 +161,11 @@ def compute_output_shape(self, input_shape):
         return tuple(output_shape)
 
     def enable_lora(
-        self, rank, a_initializer="he_uniform", b_initializer="zeros"
+        self,
+        rank,
+        lora_alpha=None,
+        a_initializer="he_uniform",
+        b_initializer="zeros",
     ):
         if self.kernel_constraint:
             raise ValueError(
@@ -187,6 +198,7 @@ def enable_lora(
         self._tracker.lock()
         self.lora_enabled = True
         self.lora_rank = rank
+        self.lora_alpha = lora_alpha if lora_alpha is not None else rank
 
     def save_own_variables(self, store):
         # Do nothing if the layer isn't yet built
@@ -261,6 +273,7 @@ def get_config(self):
         }
         if self.lora_rank:
             config["lora_rank"] = self.lora_rank
+            config["lora_alpha"] = self.lora_alpha
         return {**base_config, **config}
 
     def _check_load_own_variables(self, store):
@@ -410,7 +423,7 @@ def grad_fn(*args, upstream=None):
         if self.lora_enabled:
             lora_x = ops.matmul(inputs, self.lora_kernel_a)
             lora_x = ops.matmul(lora_x, self.lora_kernel_b)
-            x = ops.add(x, lora_x)
+            x = ops.add(x, (self.lora_alpha / self.lora_rank) * lora_x)
         if self.bias is not None:
             x = ops.add(x, self.bias)
         if self.activation is not None:
@@ -544,7 +557,8 @@ def _get_kernel_with_merged_lora(self):
                 kernel_value = ops.divide(kernel_value, kernel_scale)
                 kernel_value = ops.add(
                     kernel_value,
-                    ops.matmul(self.lora_kernel_a, self.lora_kernel_b),
+                    (self.lora_alpha / self.lora_rank)
+                    * ops.matmul(self.lora_kernel_a, self.lora_kernel_b),
                 )
                 kernel_value, kernel_scale = quantizers.abs_max_quantize(
                     kernel_value, axis=0, to_numpy=True

keras/src/layers/core/dense_test.py

@@ -271,6 +271,32 @@ def test_enable_lora(self):
         model.load_weights(temp_filepath)
         self.assertAllClose(model.predict(x), new_model.predict(x))
 
+    @pytest.mark.requires_trainable_backend
+    def test_enable_lora_with_alpha(self):
+        # Create a `Dense` layer and build it.
+        layer = layers.Dense(units=8)
+        layer.build((None, 4))
+
+        # Enable LoRA with `rank`=2 and `lora_alpha`=3.0.
+        layer.enable_lora(2, lora_alpha=3.0)
+        self.assertEqual(layer.lora_rank, 2)
+        self.assertEqual(layer.lora_alpha, 3.0)
+
+        # Manually compute the expected effective kernel:
+        # `effective_kernel_expected` = `base_kernel` +
+        # `lora_alpha / lora_rank` * `lora_kernel_a @ lora_kernel_b`
+        base_kernel = ops.convert_to_numpy(layer._kernel)
+        lora_update = np.matmul(
+            ops.convert_to_numpy(layer.lora_kernel_a),
+            ops.convert_to_numpy(layer.lora_kernel_b),
+        )
+        effective_kernel_expected = base_kernel + (3.0 / 2) * lora_update
+
+        # Verify that the effective kernel matches expectation.
+        self.assertAllClose(
+            ops.convert_to_numpy(layer.kernel), effective_kernel_expected
+        )
+
     @pytest.mark.requires_trainable_backend
     def test_lora_weight_name(self):
         class MyModel(models.Model):

keras/src/layers/core/einsum_dense.py

@@ -58,6 +58,11 @@ class EinsumDense(Layer):
             computation cost of fine-tuning large dense layers.
             You can also enable LoRA on an existing
             `EinsumDense` layer by calling `layer.enable_lora(rank)`.
+         lora_alpha: Optional integer. If set, this parameter scales the
+            low-rank adaptation delta (computed as the product of two lower-rank
+            trainable matrices) during the forward pass. The delta is scaled by
+            `lora_alpha / lora_rank`, allowing you to fine-tune the strength of
+            the LoRA adjustment independently of `lora_rank`.
         **kwargs: Base layer keyword arguments, such as `name` and `dtype`.
 
     Examples:
@@ -125,6 +130,7 @@ def __init__(
         kernel_constraint=None,
         bias_constraint=None,
         lora_rank=None,
+        lora_alpha=None,
         **kwargs,
     ):
         super().__init__(**kwargs)
@@ -142,6 +148,7 @@ def __init__(
         self.kernel_constraint = constraints.get(kernel_constraint)
         self.bias_constraint = constraints.get(bias_constraint)
         self.lora_rank = lora_rank
+        self.lora_alpha = lora_alpha if lora_alpha is not None else lora_rank
         self.lora_enabled = False
 
     def build(self, input_shape):
@@ -184,7 +191,7 @@ def build(self, input_shape):
             self.bias = None
         self.built = True
         if self.lora_rank:
-            self.enable_lora(self.lora_rank)
+            self.enable_lora(self.lora_rank, lora_alpha=self.lora_alpha)
 
     @property
     def kernel(self):
@@ -193,9 +200,9 @@ def kernel(self):
                 "You must build the layer before accessing `kernel`."
             )
         if self.lora_enabled:
-            return self._kernel + ops.matmul(
-                self.lora_kernel_a, self.lora_kernel_b
-            )
+            return self._kernel + (
+                self.lora_alpha / self.lora_rank
+            ) * ops.matmul(self.lora_kernel_a, self.lora_kernel_b)
         return self._kernel
 
     def compute_output_shape(self, _):
@@ -210,7 +217,11 @@ def call(self, inputs, training=None):
         return x
 
     def enable_lora(
-        self, rank, a_initializer="he_uniform", b_initializer="zeros"
+        self,
+        rank,
+        lora_alpha=None,
+        a_initializer="he_uniform",
+        b_initializer="zeros",
     ):
         if self.kernel_constraint:
             raise ValueError(
@@ -243,6 +254,7 @@ def enable_lora(
         self._tracker.lock()
         self.lora_enabled = True
         self.lora_rank = rank
+        self.lora_alpha = lora_alpha if lora_alpha is not None else rank
 
     def save_own_variables(self, store):
         # Do nothing if the layer isn't yet built
@@ -321,6 +333,7 @@ def get_config(self):
         }
         if self.lora_rank:
             config["lora_rank"] = self.lora_rank
+            config["lora_alpha"] = self.lora_alpha
         return {**base_config, **config}
 
     def _check_load_own_variables(self, store):
@@ -525,7 +538,7 @@ def grad_fn(*args, upstream=None):
         if self.lora_enabled:
             lora_x = ops.einsum(self.equation, inputs, self.lora_kernel_a)
             lora_x = ops.matmul(lora_x, self.lora_kernel_b)
-            x = ops.add(x, lora_x)
+            x = ops.add(x, (self.lora_alpha / self.lora_rank) * lora_x)
         if self.bias is not None:
             x += self.bias
         if self.activation is not None:
@@ -700,7 +713,8 @@ def _argsort(seq):
                 kernel_value = ops.divide(kernel_value, kernel_scale)
                 kernel_value = ops.add(
                     kernel_value,
-                    ops.matmul(self.lora_kernel_a, self.lora_kernel_b),
+                    (self.lora_alpha / self.lora_rank)
+                    * ops.matmul(self.lora_kernel_a, self.lora_kernel_b),
                 )
                 kernel_value, kernel_scale = quantizers.abs_max_quantize(
                     kernel_value, axis=self._kernel_reduced_axes, to_numpy=True

keras/src/layers/core/einsum_dense_test.py

@@ -361,6 +361,49 @@ def test_enable_lora(self):
         model.load_weights(temp_filepath)
         self.assertAllClose(model.predict(x), new_model.predict(x))
 
+    @pytest.mark.requires_trainable_backend
+    def test_enable_lora_with_alpha(self):
+        # Use a simple equation that mimics a `Dense` layer behavior.
+        equation = "ab,bc->ac"
+        output_shape = 3  # This means the kernel shape will be (input_dim, 3).
+        bias_axes = None
+
+        # Create and build the `EinsumDense` layer
+        # with an input shape (None, 2).
+        layer = layers.EinsumDense(
+            equation=equation, output_shape=output_shape, bias_axes=bias_axes
+        )
+        # Build the layer with an input shape of (batch, 2).
+        layer.build((None, 2))
+
+        # Set the base kernel weights to a known value.
+        base_kernel = np.array(
+            [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], dtype=np.float32
+        )
+        layer._kernel.assign(base_kernel)
+
+        # Enable LoRA with `rank`=2 and a custom `lora_alpha`=3.0.
+        layer.enable_lora(rank=2, lora_alpha=3.0)
+        self.assertEqual(layer.lora_rank, 2)
+        self.assertEqual(layer.lora_alpha, 3.0)
+
+        # The expected shapes are:
+        #   `base_kernel`: (2, 3)
+        #   `lora_kernel_a`: (2, 2) and `lora_kernel_b`: (2, 3)
+        a_val = np.array([[0.1, 0.2], [0.3, 0.4]], dtype=np.float32)
+        b_val = np.array([[0.5, 0.6, 0.7], [0.8, 0.9, 1.0]], dtype=np.float32)
+        layer.lora_kernel_a.assign(a_val)
+        layer.lora_kernel_b.assign(b_val)
+
+        # Compute expected effective kernel.
+        # Scaling factor is `lora_alpha / lora_rank` = 3.0 / 2 = 1.5
+        expected_delta = 1.5 * np.matmul(a_val, b_val)
+        expected_kernel = base_kernel + expected_delta
+
+        # Verify that the effective kernel property returns the expected value.
+        actual_kernel = ops.convert_to_numpy(layer.kernel)
+        self.assertAllClose(actual_kernel, expected_kernel)
+
     @pytest.mark.requires_trainable_backend
     def test_lora_rank_argument(self):
         self.run_layer_test(

keras/src/layers/core/embedding.py

@@ -65,6 +65,11 @@ class Embedding(Layer):
             computation cost of fine-tuning large embedding layers.
             You can also enable LoRA on an existing
             `Embedding` layer by calling `layer.enable_lora(rank)`.
+        lora_alpha: Optional integer. If set, this parameter scales the
+            low-rank adaptation delta (computed as the product of two lower-rank
+            trainable matrices) during the forward pass. The delta is scaled by
+            `lora_alpha / lora_rank`, allowing you to fine-tune the strength of
+            the LoRA adjustment independently of `lora_rank`.
 
     Input shape:
         2D tensor with shape: `(batch_size, input_length)`.
@@ -83,6 +88,7 @@ def __init__(
         mask_zero=False,
         weights=None,
         lora_rank=None,
+        lora_alpha=None,
         **kwargs,
     ):
         input_length = kwargs.pop("input_length", None)
@@ -100,6 +106,7 @@ def __init__(
         self.supports_masking = mask_zero
         self.autocast = False
         self.lora_rank = lora_rank
+        self.lora_alpha = lora_alpha if lora_alpha is not None else lora_rank
         self.lora_enabled = False
 
         if weights is not None:
@@ -129,9 +136,10 @@ def build(self, input_shape=None):
     @property
     def embeddings(self):
         if self.lora_enabled:
-            return self._embeddings + ops.matmul(
-                self.lora_embeddings_a, self.lora_embeddings_b
-            )
+            return self._embeddings + (
+                self.lora_alpha / self.lora_rank
+            ) * ops.matmul(self.lora_embeddings_a, self.lora_embeddings_b)
+
         return self._embeddings
 
     def call(self, inputs):
@@ -149,7 +157,11 @@ def compute_output_shape(self, input_shape):
         return (*input_shape, self.output_dim)
 
     def enable_lora(
-        self, rank, a_initializer="he_uniform", b_initializer="zeros"
+        self,
+        rank,
+        lora_alpha=None,
+        a_initializer="he_uniform",
+        b_initializer="zeros",
     ):
         if self.embeddings_constraint:
             raise ValueError(
@@ -182,6 +194,7 @@ def enable_lora(
         self._tracker.lock()
         self.lora_enabled = True
         self.lora_rank = rank
+        self.lora_alpha = lora_alpha if lora_alpha is not None else rank
 
     def save_own_variables(self, store):
         # Do nothing if the layer isn't yet built
@@ -246,6 +259,7 @@ def get_config(self):
         }
         if self.lora_rank:
             config["lora_rank"] = self.lora_rank
+            config["lora_alpha"] = self.lora_alpha
         return {**base_config, **config}
 
     def _check_load_own_variables(self, store):
@@ -339,7 +353,9 @@ def _int8_call(self, inputs, training=None):
         if self.lora_enabled:
             lora_outputs = ops.take(self.lora_embeddings_a, inputs, axis=0)
             lora_outputs = ops.matmul(lora_outputs, self.lora_embeddings_b)
-            outputs = ops.add(outputs, lora_outputs)
+            outputs = ops.add(
+                outputs, (self.lora_alpha / self.lora_rank) * lora_outputs
+            )
         return outputs
 
     def quantize(self, mode, type_check=True):

keras/src/layers/core/embedding_test.py

@@ -181,6 +181,38 @@ def test_enable_lora(self):
         model.load_weights(temp_filepath)
         self.assertAllClose(model.predict(x), new_model.predict(x))
 
+    @pytest.mark.requires_trainable_backend
+    def test_enable_lora_with_alpha(self):
+        # Create an `Embedding` layer without specifying `lora_rank`
+        layer = layers.Embedding(input_dim=3, output_dim=2)
+        layer.build((None,))  # Build the layer
+
+        # Set the base embeddings to known values.
+        base_emb = np.array(
+            [[0.1, 0.2], [0.3, 0.4], [0.5, 0.6]], dtype=np.float32
+        )
+        layer.embeddings.assign(base_emb)
+
+        # Enable LoRA with a custom alpha: `rank`=2, `lora_alpha`=3.0.
+        layer.enable_lora(2, lora_alpha=3.0)
+        self.assertEqual(layer.lora_rank, 2)
+        self.assertEqual(layer.lora_alpha, 3.0)
+
+        # Manually assign known values to lora weights.
+        a_val = np.array([[0.1, 0.1], [0.2, 0.2], [0.3, 0.3]], dtype=np.float32)
+        b_val = np.array([[0.5, 0.5], [0.6, 0.6]], dtype=np.float32)
+        layer.lora_embeddings_a.assign(a_val)
+        layer.lora_embeddings_b.assign(b_val)
+
+        # Compute the expected delta.
+        # Scaling factor: (3.0 / 2) = 1.5
+        effective_delta = 1.5 * np.matmul(a_val, b_val)
+        expected_embeddings = base_emb + effective_delta
+
+        # Verify that the effective embeddings match expectation.
+        actual_embeddings = ops.convert_to_numpy(layer.embeddings)
+        self.assertAllClose(actual_embeddings, expected_embeddings)
+
     @pytest.mark.requires_trainable_backend
     def test_lora_rank_argument(self):
         self.run_layer_test(