xnn.common.finetune.lora#

Low-rank adaptation (LoRA) of the linear layers of any xnn potential.

LoRA (Hu et al., arXiv:2106.09685) freezes a pretrained weight W and learns a low-rank update W + (alpha / r) B A with A of shape (r, d_in) and B of shape (d_out, r), B = 0 at the start so the adapted model equals the pretrained one. For the equivariant linear layers of MACE and NequIP the update must keep its block structure: an o3.Linear is a scalar matrix per irrep block that never mixes irreps, and the low-rank update is applied block by block with the same rank (Tompa et al., arXiv:2606.12704, eq. 2; the path-wise decomposition of Wang et al., ELoRA, ICML 2025, restricted to linear layers), so the adapted layer stays an equivariant linear layer.

Two adapters cover every linear layer in xnn:

  • LoRAAdapter for any module with a dense 2-D weight (torch.nn.Linear, the layers of e3nn.nn.FullyConnectedNet, PhysNet’s dense layers), whichever axis is the input;

  • LoRAEquivariantLinear for e3nn.o3.Linear.

Both compute the adapted weight and evaluate the original module with it (torch.func.functional_call / e3nn’s external weight argument), so the adapted layer is the base layer’s own code with a different weight: one code path for training and inference, and merge_lora() folds the update into the base weights for deployment at no inference cost.

Enable from a config with model.extra["lora"]:

model:
  pretrained: mace-mp-0-small
  lora: {rank: 4, alpha: 1.0}        # or lora: 4, or lora: true (rank 4)

Module Attributes

DEFAULT_RANK

Default rank of the low-rank update.

DEFAULT_ALPHA

Default LoRA scaling numerator (the update is scaled by alpha / rank).

DEFAULT_INIT_STD

Default standard deviation of the Gaussian input factor at initialization.

Functions

has_lora(model)

Whether model holds any LoRA adapter.

inject_lora(model[, rank, alpha, init_std, ...])

Replace every eligible linear layer of model by a LoRA adapter, in place.

lora_modules(model)

Every LoRA adapter inside model.

lora_options(spec)

Normalize a model.extra["lora"] entry to inject_lora() keywords.

lora_parameters(model)

The LoRA factors of model (the parameters LoRA trains).

merge_lora(model)

Fold every LoRA update into its base layer, in place.

register_lora_target(cls[, weight, in_axis])

Make instances of cls eligible for inject_lora().

Classes

LoRAAdapter(base[, rank, alpha, weight, ...])

Low-rank update of one dense weight matrix of any module.

LoRAEquivariantLinear(base[, rank, alpha, ...])

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