""" LoRA (Low-Rank Adaptation) — lightweight adapter injection for ArcisVLM v2. Provides LoRA layers that can be dynamically injected into and removed from the MoE decoder's attention projections at inference time. The HyperMother generates LoRA weights per-camera/per-scene via the HyperNetwork, and this module handles the injection mechanics. Key design decisions: - Target the fused QKV projection, applying LoRA only to Q and V slices (NOT K, following LoRA best practices from Hu et al. 2021) - NOT applied to MoE FFN gating — the learned router was trained across Stages 2-3 and disrupting it with LoRA would break load-balanced routing - LoRA modifies what the model *attends to*; MoE controls how it *processes* References: - LoRA: Low-Rank Adaptation (Hu et al., 2021) - HypeLoRA: Calibrated adapters with uncertainty (arXiv: 2603.19278) - Doc-to-LoRA / Sakana AI (arXiv: 2602.15902) """ from __future__ import annotations import math from dataclasses import dataclass from typing import Optional import torch import torch.nn as nn @dataclass class LoRAConfig: """Configuration for LoRA adapter injection.""" rank: int = 16 # Low-rank dimension (r) alpha: float = 32.0 # Scaling factor (α) dropout: float = 0.05 # Dropout on LoRA path targets: tuple[str, ...] = ("q", "v") # Which projections to adapt (subset of q, k, v) class LoRALayer(nn.Module): """ Low-Rank Adaptation layer: W' = W + (α/r) * B @ A A is initialized with Kaiming uniform, B with zeros, so the LoRA contribution starts as zero (identity behavior). Args: in_dim: Input dimension of the original projection out_dim: Output dimension of the original projection rank: Low-rank bottleneck dimension alpha: Scaling factor dropout: Dropout rate on the LoRA path """ def __init__( self, in_dim: int, out_dim: int, rank: int = 16, alpha: float = 32.0, dropout: float = 0.05, ): super().__init__() self.rank = rank self.scaling = alpha / rank # Down-projection: in_dim → rank self.A = nn.Parameter(torch.empty(rank, in_dim)) # Up-projection: rank → out_dim self.B = nn.Parameter(torch.empty(out_dim, rank)) self.dropout = nn.Dropout(dropout) if dropout > 0.0 else nn.Identity() self.reset_parameters() def reset_parameters(self) -> None: """Initialize A with Kaiming uniform, B with zeros (start as identity).""" nn.init.kaiming_uniform_(self.A, a=math.sqrt(5)) nn.init.zeros_(self.B) def forward(self, x: torch.Tensor) -> torch.Tensor: """ Compute LoRA delta: (α/r) * x @ A^T @ B^T Args: x: [..., in_dim] — input tensor (any leading batch dims) Returns: [..., out_dim] — LoRA contribution to add to original output """ return (self.dropout(x) @ self.A.T @ self.B.T) * self.scaling @property def num_params(self) -> int: """Total number of parameters in this LoRA layer.""" return self.A.numel() + self.B.numel() @classmethod def from_flat_params( cls, flat_params: torch.Tensor, in_dim: int, out_dim: int, rank: int = 16, alpha: float = 32.0, dropout: float = 0.05, ) -> "LoRALayer": """ Create a LoRALayer from a flat parameter vector (output of HyperNetwork). Args: flat_params: [rank * in_dim + out_dim * rank] — concatenated A and B in_dim, out_dim, rank, alpha, dropout: LoRA configuration Returns: LoRALayer with parameters set from flat_params """ layer = cls(in_dim, out_dim, rank, alpha, dropout) a_size = rank * in_dim b_size = out_dim * rank assert flat_params.numel() == a_size + b_size, ( f"Expected {a_size + b_size} params, got {flat_params.numel()}" ) with torch.no_grad(): layer.A.copy_(flat_params[:a_size].reshape(rank, in_dim)) layer.B.copy_(flat_params[a_size:a_size + b_size].reshape(out_dim, rank)) return layer class LoRAInjector: """ Manages LoRA injection into the MoE decoder's attention layers. Handles the mapping from a flat parameter vector (HyperNetwork output) to individual LoRA layers across all decoder transformer blocks, targeting Q and V projections in the fused QKV linear. The fused QKV layout is: [Q_params | K_params | V_params] where each segment is embed_dim wide. LoRA is applied as a residual on the Q and V slices of the output. """ def __init__(self, config: LoRAConfig, num_blocks: int, embed_dim: int): self.config = config self.num_blocks = num_blocks self.embed_dim = embed_dim # Calculate total LoRA params needed self._params_per_layer = self._calc_params_per_layer() self._num_lora_layers = num_blocks * len(config.targets) self.total_params = self._params_per_layer * self._num_lora_layers def _calc_params_per_layer(self) -> int: """Params for one LoRA layer: rank * in_dim + out_dim * rank.""" r = self.config.rank d = self.embed_dim return r * d + d * r # A: [r, d] + B: [d, r] def create_lora_layers( self, flat_params: Optional[torch.Tensor] = None, device: Optional[torch.device] = None, ) -> list[dict[str, LoRALayer]]: """ Create LoRA layers for all decoder blocks. Args: flat_params: Optional flat parameter vector from HyperNetwork. If None, creates randomly initialized LoRA layers. device: Target device for the layers Returns: List of dicts, one per block, mapping target name → LoRALayer e.g., [{"q": LoRALayer, "v": LoRALayer}, ...] """ layers = [] offset = 0 for block_idx in range(self.num_blocks): block_layers = {} for target in self.config.targets: if flat_params is not None: chunk = flat_params[offset:offset + self._params_per_layer] lora = LoRALayer.from_flat_params( chunk, in_dim=self.embed_dim, out_dim=self.embed_dim, rank=self.config.rank, alpha=self.config.alpha, dropout=self.config.dropout, ) offset += self._params_per_layer else: lora = LoRALayer( in_dim=self.embed_dim, out_dim=self.embed_dim, rank=self.config.rank, alpha=self.config.alpha, dropout=self.config.dropout, ) if device is not None: lora = lora.to(device) block_layers[target] = lora layers.append(block_layers) return layers def compute_total_lora_params( num_blocks: int, embed_dim: int, rank: int = 16, targets: tuple[str, ...] = ("q", "v"), ) -> int: """ Calculate the total number of LoRA parameters for the decoder. Useful for sizing the HyperNetwork output layer. Args: num_blocks: Number of transformer blocks in decoder embed_dim: Model dimension rank: LoRA rank targets: Which projections get LoRA (default: Q and V) Returns: Total parameter count """ params_per_layer = rank * embed_dim + embed_dim * rank # A + B return params_per_layer * num_blocks * len(targets)