DoRA

The motivating measurement comes first. Take a full fine-tune, decompose the weight change into a change in each column’s magnitude and a change in its direction, and plot the two against each other. Full fine-tuning shows a negative correlation — it makes large directional changes with small magnitude changes, or the reverse. LoRA shows a positive one: it moves both together, because a rank-$r$ additive update cannot easily do otherwise.

DoRA adopts the decomposition as the parameterisation. Write $W$ as a per-column magnitude vector $m$ times a column-normalised direction matrix, then adapt the direction with LoRA and let $m$ train freely.

$m$ is initialised to $\lVert W_0 \rVert_c$ and $B$ at zero, so training starts exactly at $W_0$ as it does for LoRA. The normalisation is what decouples the two: however far $BA$ pushes a column, the magnitude is set by $m$ alone.

The magnitude vector costs one scalar per output column — a rounding error against the low-rank factors, and the reason DoRA is compared against LoRA at equal rank rather than equal parameters.

Where it earns its keep is at rank 4 and 8, where LoRA is capacity-starved. By rank 64 the two converge, and the extra forward-pass cost stops paying for itself.

import math
import torch
from torch import Tensor, nn


class DoRALinear(nn.Module):
    def __init__(self, base: nn.Linear, r: int = 8, alpha: int = 16):
        super().__init__()
        self.base = base
        self.base.requires_grad_(False)

        self.A = nn.Parameter(torch.empty(r, base.in_features))
        self.B = nn.Parameter(torch.zeros(base.out_features, r))
        nn.init.kaiming_uniform_(self.A, a=math.sqrt(5))
        self.scale = alpha / r

        # one magnitude per output column, initialised from the base weight
        norm = base.weight.norm(dim=1, keepdim=True)
        self.magnitude = nn.Parameter(norm.detach().clone())

    def forward(self, x: Tensor) -> Tensor:
        weight = self.base.weight + (self.B @ self.A) * self.scale
        direction = weight / weight.norm(dim=1, keepdim=True).detach()
        return nn.functional.linear(x, self.magnitude * direction)

The .detach() on the norm is from the paper and is not optional: letting the gradient flow through the normaliser roughly doubles the backward-pass memory for no measured gain. Note that the merged weight must be materialised each forward pass to compute the column norms, which is precisely the cost LoRA avoids by keeping the factors separate.