DoRA: Decomposing Weights into Magnitude and Direction

DoRA (Weight-Decomposed Low-Rank Adaptation) represents a refined approach to fine-tuning that addresses a specific limitation of plain LoRA while preserving its efficiency. At its core, DoRA builds upon LoRA's foundational principles but introduces a decomposition step: it separates a pretrained weight matrix W into two distinct components — magnitude (m) and direction (V).

This separation is the key innovation: rather than plain LoRA's single, undifferentiated additive update (W_final = W + (alpha/r)(B × A), one combined adjustment applied uniformly), DoRA allows the fine-tuning process to target magnitude and direction INDEPENDENTLY. Intuitively, magnitude captures 'how big' a weight vector is, while direction captures 'which way it points' — these are genuinely different aspects of what a weight update can change, and plain LoRA's single additive update doesn't distinguish between adjusting one versus the other; it just moves the weight in some combined direction with some combined magnitude all at once.

By separating these two components and fine-tuning each on its own terms, DoRA improves parameter efficiency and performance relative to plain LoRA's simpler approach — the model gets a more expressive, more targeted way to adapt its weights within roughly the same parameter budget, rather than being restricted to whatever adjustments a single combined additive term can represent.

This makes DoRA a genuine refinement rather than an unrelated alternative — it's explicitly built on top of LoRA's low-rank adaptation principles (the efficiency of training small A/B-style matrices is preserved), while fixing a specific representational limitation in how the adaptation itself is structured. In practice, this makes DoRA worth reaching for specifically when you need the best possible accuracy within a LoRA-style parameter budget and can afford the modest added implementation complexity of the magnitude/direction split, compared to plain LoRA's simpler, more battle-tested single additive update.

import torch
import torch.nn as nn

def decompose_weight(W: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
    """Split a weight matrix into magnitude (per-column norm) and
    direction (unit vectors) — the two components DoRA fine-tunes
    independently, instead of one combined additive update."""
    magnitude = W.norm(dim=0, keepdim=True)  # "how big" — shape (1, k)
    direction = W / magnitude                 # "which way" — unit vectors
    return magnitude, direction

class DoRALayer(nn.Module):
    """Simplified DoRA: magnitude is directly trainable; direction is
    adapted via a LoRA-style low-rank update, then re-normalized."""
    def __init__(self, W_frozen: torch.Tensor, r: int, alpha: float):
        super().__init__()
        magnitude, direction = decompose_weight(W_frozen)
        self.direction_base = direction  # frozen base direction
        self.magnitude = nn.Parameter(magnitude.clone())  # trainable
        d, k = W_frozen.shape
        self.A = nn.Parameter(torch.randn(d, r) * 0.01)   # LoRA-style direction update
        self.B = nn.Parameter(torch.zeros(r, k))
        self.alpha, self.r = alpha, r

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        direction_delta = (self.alpha / self.r) * (self.A @ self.B)
        adapted_direction = self.direction_base + direction_delta
        adapted_direction = adapted_direction / adapted_direction.norm(dim=0, keepdim=True)
        return x @ (self.magnitude * adapted_direction)  # magnitude and direction, applied separately