LoRA

class braintrace.nn.LoRA

A standalone LoRA layer.

LoRA (Low-Rank Adaptation) injects two low-rank factors into a layer so a large pre-trained model can be fine-tuned with far fewer parameters. This subclass preserves the upstream brainstate.nn.LoRA constructor and replaces only the forward pass so that the multiplication is routed through braintrace.lora_matmul() and therefore participates in eligibility-trace computation.

The layer adds a low-rank component \(\frac{1}{r} B A\) to the base weight, where \(B\) and \(A\) are learnable factors of rank \(r\):

\[W_{\mathrm{LoRA}} = W_{\text{orig}} + \frac{1}{r} B A\]

The scaling factor is fixed to 1 / lora_rank.

Parameters:

• in_features (int) – Number of input features.

• lora_rank (int) – Rank of the LoRA decomposition.

• out_features (int) – Number of output features.

• base_module (brainstate.nn.Module or None, optional) – Optional base layer that is called on x and added to the LoRA branch. Default None.

• kernel_init (Callable or ArrayLike, optional) – Initializer used for both lora_a (in×rank) and lora_b (rank×out). Default is LecunNormal(). To get the classic “LoRA-zero” initialisation use init.ZeroInit().

• param_type (type, optional) – ParamState subclass used to wrap the weights. Default is brainstate.ParamState.

• in_size (int or Sequence[int], optional) – Optional explicit input size override. Default None.

Variables:

• in_features (int) – Number of input features.

• out_features (int) – Number of output features.

• base_module (brainstate.nn.Module or None) – The optional base layer added to the LoRA branch.

• weight (ParamState) – ParamState whose value is a dict with two keys: 'lora_a' of shape (in_features, lora_rank) and 'lora_b' of shape (lora_rank, out_features).

Examples

>>> import brainstate
>>> import braintrace
>>>
>>> # Create a standalone LoRA layer
>>> brainstate.environ.set(precision=64)
>>> layer = braintrace.nn.LoRA(in_features=3, lora_rank=2, out_features=4)
>>> x = brainstate.random.randn(16, 3)
>>> y = layer(x)
>>> print(y.shape)
(16, 4)
>>>
>>> # Wrap around an existing linear layer
>>> linear = brainstate.nn.Linear(3, 4)
>>> wrapper = braintrace.nn.LoRA(3, 2, 4, base_module=linear)
>>> assert wrapper.base_module is linear
>>> y = wrapper(x)
>>> print(y.shape)
(16, 4)

update(x)

Apply the low-rank adaptation through the ETP lora_matmul.

Computes \(y = \frac{1}{r}\, x\, \mathbf{A}\, \mathbf{B}\) via braintrace.lora_matmul(), where \(\mathbf{A}\) is the input-facing factor lora_a of shape (in, rank) and \(\mathbf{B}\) is the output-facing factor lora_b of shape (rank, out) (so the LoRA factors participate in online-learning trace computation), and adds the optional base-module output.

Parameters:

x (ArrayLike) – Input array, of shape (..., in_features).

Returns:

ArrayLike – The adapted output, of shape (..., out_features).

LoRA.__init__(in_features, lora_rank, out_features, *, base_module=None, kernel_init=LecunNormal(   scale=1.0, mode='fan_in', in_axis=-2, out_axis=-1, distribution='truncated_normal', rng=RandomState(Array((), dtype=key<fry>) overlaying:   [2233333421 2029709265]), unit=Unit("1") ), param_type=<class 'brainstate.ParamState'>, in_size=None)