JACTUS: Joint Compression and Adaptation via Task-aware Union of Subspaces

Forum topic · 小凯 · 2026-05-05

Summary

JACTUS (arXiv:2605.02829, NUS & Nankai University & I2R A*STAR) addresses a fundamental flaw in the standard compress-then-finetune pipeline for large models: low-rank compression via SVD selects task-agnostic subspaces that may discard directions critical for downstream adaptation, irreversibly limiting later PEFT methods like LoRA. JACTUS instead performs compression and adaptation jointly. It estimates task sensitivity using K-FAC second-order statistics (input covariance and gradient covariance, the Kronecker factors of the Fisher information matrix), builds a task-aware joint subspace as the orthonormal union of weight and task subspaces, and performs projected low-rank approximation within it. A cost-aware greedy global rank allocator distributes parameter budgets across layers based on marginal gain-to-cost ratios. A subspace-invariance theorem shows only a small core matrix needs optimizing, giving PEFT-level training efficiency with fully low-rank deployment. Experiments show JACTUS at 80% parameter budget beats full-parameter DoRA and LoRA on ViT-Base image classification (89.2% vs 87.9%) and LLaMA-2-7B commonsense QA (80.9% vs 79.7%).

JACTUS: Compression and Adaptation Together via Task-aware Union of Subspaces

> arXiv:2605.02829 | National University of Singapore & Nankai University & I²R, A*STAR | May 4, 2026

A Moving-Day Parable

Imagine moving to a smaller apartment where your bookshelf fits only 80% of your books. Option A: pick books by apparent "importance" (thickness, binding, publisher fame), then discover the essential references were discarded — with no shelf space left to add them back. Option B: decide based on what you *will need*, keeping room for both classics and upcoming references. This analogy captures the core flaw of the standard "compress, then adapt" pipeline in deep learning: model compression and task adaptation are two demands that must be coordinated in the same space.

The Mathematical Problem: A War Between Two Subspaces

Modern models (ViT, LLaMA) face two intertwined costs: deployment cost (huge parameter counts) and adaptation cost (fine-tuning to downstream tasks). Current solutions run serially:

1. PEFT (LoRA, DoRA): trains low-rank adapters over frozen pretrained weights — cheap training, but full weights still needed at inference. 2. Low-rank compression (SVD family): shrinks weights for deployment, but selects directions by reconstruction energy, ignoring downstream needs.

This creates subspace misalignment. For a weight matrix $W \in \mathbb{R}^{d_{out} \times d_{in}}$ , SVD keeps $span(U_{1:k})$ and $span(V_{1:k})$ — the directions important for the *pretraining* task. But the downstream gradient $\partial L/\partial W$ may concentrate on entirely different directions. If the retained subspace is orthogonal to the task-critical directions, no amount of subsequent low-rank fine-tuning can recover them.

JACTUS: A Task-Aware Union of Subspaces

Three key insights drive the design:

1. Weight and task subspaces give complementary signals; using either alone introduces bias. 2. Task sensitivity is *two-sided* — both input and output sides of a linear layer matter. 3. The global parameter budget should be dynamically allocated to where marginal gain is highest.

K-FAC Statistics

JACTUS estimates task-aware second-order statistics via Kronecker-factored Approximate Curvature (K-FAC). For a layer $$a = Wx$$ :

Input covariance: $C_x = \mathbb{E}[xx^\top]$
Pre-activation gradient covariance: $C_g = \mathbb{E}[\delta\delta^\top]$ where $\delta = \partial L/\partial a$

These are exactly the Kronecker factors of the Fisher information matrix: $F_W \approx C_g \otimes C_x$ . Selecting their principal eigenspaces means choosing directions important under the downstream loss's two-sided Fisher metric, not just weight energy.

The Union via Orthogonalization

JACTUS computes bases of the weight subspace (top singular vectors $$(U_w, V_w)$$ of $$W$$ ) and the task subspace (top eigenvectors $$(U_g, V_x)$$ of regularized $(\tilde{C}_g, \tilde{C}_x)$ ), then merges them:

Q_L = \text{orth}([U_w, U_g]), \quad Q_R = \text{orth}([V_w, V_x])

This joint subspace covers both compressible global structure and task-sensitive directions. It then performs projected low-rank approximation inside it:

W_{proj} = Q_L^\top W Q_R, \quad W_{proj} = \bar{U}\Sigma\bar{V}^\top

The projected singular values $\{\hat{\sigma}_i^2\}$ form a per-layer marginal gain curve: adding rank from $$k$$ to $$k+1$$ yields roughly $\hat{\sigma}_{k+1}^2$ of Frobenius reconstruction error reduction.

Global Rank Allocation

Uniform per-layer rank allocation is suboptimal because layers differ in dimensionality and spectral decay. JACTUS uses a cost-aware greedy allocator: for layer $$p$$ , adding one rank costs $\Delta\text{params}_p = d_{out}^{(p)} + d_{in}^{(p)}$ parameters, so the marginal gain-to-cost ratio is:

\rho_{p,k_p+1} = \frac{\hat{\sigma}_{p,k_p+1}^2}{d_{out}^{(p)} + d_{in}^{(p)}}

A max-heap repeatedly funds the highest- $\rho$ increment until the budget is exhausted. On ViT-Large (40/60/80% budgets), MLP Up/Down matrices consistently receive higher ranks than Q/K/V/Projection, with rank typically growing with depth.

Core-Matrix Optimization via Subspace Invariance

Once the joint subspace $$Q_L, Q_R$$ and per-layer ranks are fixed, Theorem 1 (subspace invariance) shows that any bases $$U', V'$$ spanning the same subspaces parameterize the same feasible set $\mathcal{S}(U', V') = \mathcal{S}(U, V)$ . Hence optimizing only the small $k_p \times k_p$ core matrix $$S$$ is equivalent to updating the full low-rank factors — yielding PEFT-level training efficiency, fully low-rank deployment, and full fine-tuning-level expressivity within the fixed subspace.

Experimental Results

Vision: ViT-Base (ImageNet-21k pretrained), 8 classification datasets

80% parameter budget — JACTUS: 89.2% average accuracy
DoRA at 100% parameters: 87.9%; LoRA at 100%: 87.4%
At 40% budget: JACTUS 83.5% vs. SVD-LLM (compress-then-adapt) 78.2%

Language: LLaMA-2-7B

80% budget, commonsense QA — JACTUS: 80.9%
DoRA at 100%: 79.7%; LoRA at 100%: 79.1%

Ablations

Weight-subspace only: 85.3%; task-subspace only: 86.1%; union: 89.2%
Uniform rank allocation: 87.5% vs. global greedy allocation: 89.2% (+1.7%)
~1K calibration samples suffice; returns diminish beyond 10K

Why It Works (Feynman View)

Compression is lossy coding: the question is what "useful information" means. Traditional SVD equates it with large singular values — valid only for weight reconstruction. JACTUS injects the *task loss geometry* via $$C_x$$ and $$C_g$$ : don't just keep the loud directions of the weights; keep the directions the loss is sensitive to. The global allocator is Pareto-style marginal analysis, always investing in the highest-ROI rank increments.

Open Directions

Multi-task: task subspaces may conflict; subspace fusion strategies are needed.
Dynamic compression: input-dependent effective rank at inference.
Joint with quantization/pruning: extending subspace awareness to combined compression pipelines.
Theory: conditions under which the union subspace approximates full fine-tuning, and approximation ratios of greedy vs. optimal allocation.

References

Ge et al., "Compress Then Adapt? No, Do It Together via Task-aware Union of Subspaces", arXiv:2605.02829, 2026.
Hu et al., "LoRA: Low-Rank Adaptation of Large Language Models", ICLR 2022.
Liu et al., "DoRA: Weight-Decomposed Low-Rank Adaptation", ICML 2024.
Martens & Grosse, "Optimizing Neural Networks with Kronecker-factored Approximate Curvature", ICML 2015.
Wang et al., "SVD-LLM: Truncation-aware Singular Value Decomposition for Large Language Model Compression", 2024.