vlora-dev 0.2.0__py3-none-any.whl

vlora/ops.py

@@ -0,0 +1,229 @@
+"""Pure math operations for shared low-rank subspace computation.
+
+All functions are stateless and operate on raw tensors — no adapter
+or subspace concepts leak in here.
+"""
+
+from __future__ import annotations
+
+import torch
+from torch import Tensor
+
+
+def compute_svd(
+    data_matrix: Tensor,
+    num_components: int | None = None,
+    center: bool = True,
+) -> tuple[Tensor, Tensor, Tensor]:
+    """SVD on data matrix, returning top-k components.
+
+    Args:
+        data_matrix: (N, D) matrix where rows are observations.
+        num_components: Number of singular vectors to keep. If None, keep all.
+        center: Whether to mean-center rows before SVD.
+
+    Returns:
+        components: (k, D) right singular vectors (shared basis).
+        singular_values: (k,) corresponding singular values.
+        mean: (D,) row mean (zeros if center=False).
+    """
+    if center:
+        mean = data_matrix.mean(dim=0)
+        centered = data_matrix - mean
+    else:
+        mean = torch.zeros(data_matrix.shape[1], dtype=data_matrix.dtype, device=data_matrix.device)
+        centered = data_matrix
+
+    # full_matrices=False gives the economy SVD — U is (N, min(N,D))
+    U, S, Vh = torch.linalg.svd(centered, full_matrices=False)
+
+    if torch.isnan(S).any():
+        raise ValueError(
+            "SVD produced NaN singular values. Check input data for "
+            "NaN/Inf or try using float64 precision."
+        )
+
+    k = num_components if num_components is not None else len(S)
+    k = min(k, len(S))
+
+    return Vh[:k], S[:k], mean
+
+
+def project_onto_subspace(weights: Tensor, components: Tensor) -> Tensor:
+    """Project weight vectors onto the subspace basis.
+
+    Args:
+        weights: (D,) or (N, D) weight vectors to project.
+        components: (k, D) orthonormal basis vectors.
+
+    Returns:
+        loadings: (k,) or (N, k) projection coefficients.
+    """
+    # loadings = weights @ V^T  (since components = V^T, i.e. rows are basis vectors)
+    return weights @ components.T
+
+
+def reconstruct_from_subspace(components: Tensor, loadings: Tensor) -> Tensor:
+    """Reconstruct weight vectors from subspace loadings.
+
+    Args:
+        components: (k, D) orthonormal basis vectors.
+        loadings: (k,) or (N, k) projection coefficients.
+
+    Returns:
+        reconstructed: (D,) or (N, D) reconstructed weight vectors.
+    """
+    return loadings @ components
+
+
+def gram_schmidt(basis: Tensor, new_vectors: Tensor) -> Tensor:
+    """Orthogonalize new_vectors against an existing orthonormal basis.
+
+    Appends only those new directions that have non-trivial norm after
+    projection removal (threshold: 1e-6).
+
+    Args:
+        basis: (k, D) existing orthonormal basis.
+        new_vectors: (m, D) candidate vectors.
+
+    Returns:
+        expanded_basis: (k + n, D) where n <= m new orthogonal directions
+            were found.
+    """
+    vectors = list(basis)
+
+    for v in new_vectors:
+        v = v.clone()
+        # Remove components along every existing basis vector
+        for b in vectors:
+            v = v - (v @ b) * b
+        norm = v.norm()
+        if norm > 1e-6:
+            vectors.append(v / norm)
+
+    return torch.stack(vectors)
+
+
+def incremental_svd_update(
+    components: Tensor,
+    singular_values: Tensor,
+    mean: Tensor,
+    n_seen: int,
+    new_data: Tensor,
+    max_components: int | None = None,
+) -> tuple[Tensor, Tensor, Tensor, int]:
+    """Incrementally update SVD with new data points.
+
+    Uses the projection-residual approach: project new data onto existing
+    basis, compute residual, and if significant, expand the basis with
+    the residual direction via QR decomposition.
+
+    Args:
+        components: (k, D) current orthonormal basis.
+        singular_values: (k,) current singular values.
+        mean: (D,) current data mean.
+        n_seen: Number of data points seen so far.
+        new_data: (m, D) new data points to incorporate.
+        max_components: Cap on number of components. If None, allows growth.
+
+    Returns:
+        updated_components: (k', D) updated basis.
+        updated_singular_values: (k',) updated singular values.
+        updated_mean: (D,) updated mean.
+        new_n_seen: Updated count.
+    """
+    m = new_data.shape[0]
+    n_total = n_seen + m
+
+    # Update mean incrementally
+    new_mean = (mean * n_seen + new_data.sum(dim=0)) / n_total
+
+    # Center new data with updated mean
+    centered_new = new_data - new_mean
+
+    # Also adjust for mean shift on existing data:
+    # The old centered data had mean=0 relative to old_mean.
+    # Relative to new_mean, old data is shifted by (old_mean - new_mean).
+    mean_shift = mean - new_mean
+
+    # Project new centered data onto existing basis
+    projections = centered_new @ components.T  # (m, k)
+    residuals = centered_new - projections @ components  # (m, D)
+
+    # Find new orthogonal directions from residuals via QR
+    residual_norms = residuals.norm(dim=1)
+    significant = residual_norms > 1e-6
+    new_directions = []
+
+    if significant.any():
+        sig_residuals = residuals[significant]
+        # Orthogonalize residuals against existing basis and each other
+        expanded = gram_schmidt(components, sig_residuals)
+        new_directions_tensor = expanded[components.shape[0]:]
+        if new_directions_tensor.shape[0] > 0:
+            new_directions.append(new_directions_tensor)
+
+    # Build augmented system
+    if new_directions:
+        extra = torch.cat(new_directions, dim=0)
+        all_components = torch.cat([components, extra], dim=0)
+        # Approximate new singular values for expanded directions
+        extra_projections = centered_new @ extra.T  # (m, n_extra)
+        extra_svals = extra_projections.norm(dim=0)
+        all_svals = torch.cat([singular_values, extra_svals])
+    else:
+        all_components = components
+        # Update singular values to account for new data contribution
+        new_contributions = projections.norm(dim=0)
+        all_svals = torch.sqrt(singular_values ** 2 + new_contributions ** 2)
+
+    # Account for mean shift effect on singular values
+    shift_proj = mean_shift @ all_components.T
+    shift_contribution = shift_proj * (n_seen ** 0.5)
+    all_svals = torch.sqrt(all_svals ** 2 + shift_contribution ** 2)
+
+    # Sort by singular value magnitude (descending)
+    order = all_svals.argsort(descending=True)
+    all_components = all_components[order]
+    all_svals = all_svals[order]
+
+    # Cap components if needed
+    if max_components is not None and all_components.shape[0] > max_components:
+        all_components = all_components[:max_components]
+        all_svals = all_svals[:max_components]
+
+    return all_components, all_svals, new_mean, n_total
+
+
+def explained_variance_ratio(singular_values: Tensor) -> Tensor:
+    """Compute cumulative explained variance ratio from singular values.
+
+    Args:
+        singular_values: (k,) singular values in descending order.
+
+    Returns:
+        cumulative_ratio: (k,) cumulative fraction of total variance explained.
+    """
+    variance = singular_values ** 2
+    total = variance.sum()
+    if total == 0:
+        return torch.zeros_like(variance)
+    return torch.cumsum(variance, dim=0) / total
+
+
+def select_num_components(singular_values: Tensor, threshold: float = 0.6) -> int:
+    """Select number of components to explain at least `threshold` variance.
+
+    Args:
+        singular_values: (k,) singular values in descending order.
+        threshold: Minimum cumulative variance ratio (default 0.6 per paper).
+
+    Returns:
+        Number of components needed.
+    """
+    cumulative = explained_variance_ratio(singular_values)
+    # Find first index where cumulative ratio >= threshold
+    above = (cumulative >= threshold).nonzero(as_tuple=True)[0]
+    if len(above) == 0:
+        return len(singular_values)
+    return above[0].item() + 1

vlora/pipeline.py

@@ -0,0 +1,70 @@
+"""High-level convenience wrappers for the 3-step pipeline."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+from vlora.io import LoRAWeights, load_adapter, save_adapter
+from vlora.subspace import SharedSubspace
+
+
+def init_subspace(
+    adapter_paths: list[str | Path],
+    task_ids: list[str] | None = None,
+    variance_threshold: float = 0.6,
+    num_components: int | None = None,
+) -> SharedSubspace:
+    """Load adapters from disk and build a shared subspace in one call.
+
+    Args:
+        adapter_paths: Directories containing PEFT adapters.
+        task_ids: Names for each adapter.
+        variance_threshold: Variance threshold for auto component selection.
+        num_components: Explicit number of components (overrides threshold).
+
+    Returns:
+        Initialized SharedSubspace.
+    """
+    adapters = [load_adapter(p) for p in adapter_paths]
+    return SharedSubspace.from_adapters(
+        adapters,
+        task_ids=task_ids,
+        variance_threshold=variance_threshold,
+        num_components=num_components,
+    )
+
+
+def absorb_task(
+    subspace: SharedSubspace,
+    adapter_path: str | Path,
+    task_id: str,
+) -> None:
+    """Load a new adapter and absorb it into the subspace.
+
+    Args:
+        subspace: Existing shared subspace (modified in-place).
+        adapter_path: Directory containing the new PEFT adapter.
+        task_id: Name for the new task.
+    """
+    adapter = load_adapter(adapter_path)
+    subspace.absorb(adapter, task_id)
+
+
+def extract_adapter(
+    subspace: SharedSubspace,
+    task_id: str,
+    output_path: str | Path,
+) -> LoRAWeights:
+    """Reconstruct a task's adapter and save it to disk.
+
+    Args:
+        subspace: Shared subspace containing the task.
+        task_id: Task to reconstruct.
+        output_path: Directory to save the PEFT adapter to.
+
+    Returns:
+        The reconstructed LoRAWeights.
+    """
+    weights = subspace.reconstruct(task_id)
+    save_adapter(weights, output_path)
+    return weights

vlora/router.py

@@ -0,0 +1,173 @@
+"""TaskRouter — lightweight routing over adapter loadings.
+
+Routes inputs to a soft blend of task adapters by producing per-task
+weights. Since adapters are represented as small loading vectors in the
+shared subspace, blending is a cheap linear combination rather than
+reconstructing and merging full LoRA matrices.
+
+Usage:
+    subspace = SharedSubspace.load("shared_subspace/")
+    router = TaskRouter.from_subspace(subspace, hidden_dim=64)
+
+    # During inference:
+    model = VLoRAModel(base_model, subspace)
+    x = get_input_embedding(batch)  # (B, embed_dim)
+    blended = router.blend_loadings(x)  # per-input blended loadings
+"""
+
+from __future__ import annotations
+
+from typing import Literal
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+
+from vlora.subspace import SharedSubspace, TaskProjection
+
+
+class TaskRouter(nn.Module):
+    """Small MLP that produces soft task-blend weights from input features.
+
+    Given input embeddings (B, input_dim), outputs (B, num_tasks) blend
+    weights that sum to 1 (via softmax). These weights define a per-input
+    mixture of task loadings in the shared subspace.
+    """
+
+    def __init__(
+        self,
+        input_dim: int,
+        num_tasks: int,
+        hidden_dim: int = 64,
+        task_ids: list[str] | None = None,
+        temperature: float = 1.0,
+    ):
+        super().__init__()
+        self.input_dim = input_dim
+        self.num_tasks = num_tasks
+        self.hidden_dim = hidden_dim
+        self.task_ids = task_ids or [f"task_{i}" for i in range(num_tasks)]
+        self.temperature = temperature
+
+        self.net = nn.Sequential(
+            nn.Linear(input_dim, hidden_dim),
+            nn.ReLU(),
+            nn.Linear(hidden_dim, num_tasks),
+        )
+
+    def forward(self, x: Tensor) -> Tensor:
+        """Compute task blend weights.
+
+        Args:
+            x: (B, input_dim) input features.
+
+        Returns:
+            weights: (B, num_tasks) softmax blend weights.
+        """
+        logits = self.net(x) / self.temperature
+        return F.softmax(logits, dim=-1)
+
+    def blend_loadings(
+        self,
+        x: Tensor,
+        subspace: SharedSubspace,
+    ) -> TaskProjection:
+        """Produce blended loadings for a batch by mixing task loadings.
+
+        Computes a weighted average of all task loadings using the
+        router's output weights. Returns a single TaskProjection whose
+        loadings are the blend. For batched inference, uses the mean
+        blend across the batch.
+
+        Args:
+            x: (B, input_dim) input features.
+            subspace: SharedSubspace containing the tasks to blend.
+
+        Returns:
+            TaskProjection with blended loadings (one set for the batch).
+        """
+        weights = self.forward(x)  # (B, num_tasks)
+        # Average across batch for a single blended adapter
+        avg_weights = weights.mean(dim=0)  # (num_tasks,)
+
+        blended_a: dict[str, Tensor] = {}
+        blended_b: dict[str, Tensor] = {}
+
+        for layer in subspace.layer_names:
+            # Stack all task loadings: (num_tasks, k)
+            stack_a = torch.stack([
+                subspace.tasks[tid].loadings_a[layer]
+                for tid in self.task_ids
+            ])
+            stack_b = torch.stack([
+                subspace.tasks[tid].loadings_b[layer]
+                for tid in self.task_ids
+            ])
+
+            # Weighted combination: (k,)
+            blended_a[layer] = avg_weights @ stack_a
+            blended_b[layer] = avg_weights @ stack_b
+
+        return TaskProjection(
+            task_id="__routed__",
+            loadings_a=blended_a,
+            loadings_b=blended_b,
+        )
+
+    @classmethod
+    def from_subspace(
+        cls,
+        subspace: SharedSubspace,
+        input_dim: int,
+        hidden_dim: int = 64,
+        temperature: float = 1.0,
+        init_from_loadings: bool = True,
+    ) -> TaskRouter:
+        """Create a router matched to a subspace's task structure.
+
+        Optionally initializes the final linear layer's weights so that
+        the router starts biased toward separating tasks based on their
+        loading similarity (warm start for fine-tuning).
+
+        Args:
+            subspace: SharedSubspace with registered tasks.
+            input_dim: Dimension of input features the router will see.
+            hidden_dim: Router hidden layer size.
+            temperature: Softmax temperature (higher = softer blending).
+            init_from_loadings: If True, use task loading similarity to
+                bias the output layer (helps convergence).
+        """
+        task_ids = sorted(subspace.tasks.keys())
+        num_tasks = len(task_ids)
+
+        router = cls(
+            input_dim=input_dim,
+            num_tasks=num_tasks,
+            hidden_dim=hidden_dim,
+            task_ids=task_ids,
+            temperature=temperature,
+        )
+
+        if init_from_loadings and num_tasks > 1:
+            # Use task loading norms as output bias (tasks with larger
+            # loadings get slightly higher initial routing weight)
+            with torch.no_grad():
+                biases = []
+                for tid in task_ids:
+                    proj = subspace.tasks[tid]
+                    total_norm = sum(
+                        proj.loadings_a[l].norm() + proj.loadings_b[l].norm()
+                        for l in subspace.layer_names
+                    )
+                    biases.append(total_norm)
+                bias_tensor = torch.stack(biases)
+                bias_tensor = bias_tensor / (bias_tensor.max() + 1e-8)
+                router.net[-1].bias.data.copy_(bias_tensor)
+
+        return router
+
+    @property
+    def num_params(self) -> int:
+        """Total number of router parameters."""
+        return sum(p.numel() for p in self.parameters())