burn_core/module/param/lora.rs
1use super::Param;
2use burn_tensor::Tensor;
3
4/// A LoRA (Low-Rank Adaptation) adapter attached to a frozen weight [parameter](Param).
5///
6/// When present on a `Param<Tensor<2>>`, the parameter materializes its effective value as
7/// `base + scale * (a @ b)`, where `base` is the frozen (and optionally quantized) weight and
8/// `a`/`b` are the trainable low-rank factors. The frozen base is the stored value of the
9/// parameter; the adapter factors are surfaced to the optimizer, autodiff and record systems as
10/// regular parameters with their own [`ParamId`](super::ParamId)s through the module
11/// visitor/mapper traversal.
12#[derive(Clone, Debug)]
13pub struct LoraAdapter {
14 /// Down-projection factor with shape `[d_in, rank]` (trainable).
15 pub a: Param<Tensor<2>>,
16 /// Up-projection factor with shape `[rank, d_out]` (trainable).
17 pub b: Param<Tensor<2>>,
18 /// Scaling factor applied to the low-rank product, typically `alpha / rank`.
19 pub scale: f64,
20}
21
22impl LoraAdapter {
23 /// Compute the low-rank delta `scale * (a @ b)` with shape `[d_in, d_out]`.
24 ///
25 /// `a` and `b` are read through [`Param::val`], so the delta always reflects the current
26 /// (optimizer-updated) factors and keeps them as autodiff leaves for backpropagation.
27 pub fn delta(&self) -> Tensor<2> {
28 self.a.val().matmul(self.b.val()).mul_scalar(self.scale)
29 }
30}