xnn.common.finetune.lora.LoRAEquivariantLinear

class xnn.common.finetune.lora.LoRAEquivariantLinear(base, rank=4, alpha=1.0, *, init_std=0.001)

Bases: Module

Block-wise low-rank update of an e3nn.o3.Linear.

An o3.Linear holds one weight block W_p of shape (mul_in, mul_out) per path p (an input irrep block connected to an output block of the same irrep), applied as path_weight * W_p on the irrep’s channels. Each block gets its own rank-r update acting on the effective map: path_weight * W_p + scaling * A_p B_p, with A_p of shape (mul_in, r) Gaussian and B_p of shape (r, mul_out) zero at the start. Every block stays a scalar matrix on one irrep, so the adapted layer is equivariant. The base linear is evaluated with the adapted flat weight through its external weight argument. Bias terms (scalar outputs) are left unchanged.

Parameters:

• base (e3nn.o3.Linear) – The layer to adapt (internal, shared weights).

• rank (int, optional) – Rank of every block’s update, by default 4.

• alpha (float, optional) – Scaling numerator, by default 1.0.

• init_std (float, optional) – Standard deviation of the input factors at initialization, by default 1e-3.

Variables:

• base (e3nn.o3.Linear) – The adapted layer (frozen weight).

• lora_out (lora_in,) – The factors, one pair per weight path, in instruction order.

• paths (list of tuple) – (offset, numel, path_shape, path_weight) of each adapted path in the flat weight.

delta()

The unscaled update of the flat weight (every path, in layout order).

Return type:

Tensor

adapted_weight()

W + scaling * delta as a flat weight.

Return type:

Tensor

forward(x)

Evaluate the base linear with the adapted weight.

Parameters:

x (Tensor)

Return type:

Tensor

merge()

Fold the update into the base weight and return the base linear.

Return type:

Module

extra_repr()

Set the extra representation of the module.

To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.

Return type:

str