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

Forum topic · 小凯 · 2026-05-05

Summary

JACTUS is a method that unifies model compression and task adaptation, which are traditionally performed sequentially (compress via SVD first, then adapt with PEFT methods like LoRA). The standard pipeline suffers from subspace misalignment: compression keeps task-agnostic directions based on weight energy, potentially discarding directions critical for downstream tasks. JACTUS builds a task-aware union of subspaces by combining the weight subspace (from SVD of the weight matrix) with a task subspace derived from K-FAC second-order statistics (input and gradient covariances), merging them via orthogonalization. Within this joint subspace, it performs projected low-rank approximation and uses a cost-aware greedy rank allocator that distributes a global parameter budget by marginal gain per unit parameter. A subspace-invariance theorem shows only a small core matrix needs optimization, giving PEFT-level training efficiency with compact low-rank deployment. Experiments show JACTUS at an 80% parameter budget outperforms full-parameter DoRA and LoRA on ViT-Base image classification (89.2% vs 87.9%) and on LLaMA-2-7B commonsense QA (80.9% vs 79.7%), with ablations confirming the value of both the union subspace and global rank allocation.

JACTUS: Compression and Adaptation, Together — When Model Slimming No Longer Sacrifices the Soul

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

The Problem: Two Wars in One Space

Modern pipelines treat model compression and task adaptation as a serial process: first compress the model with SVD-style truncation, then adapt it with parameter-efficient fine-tuning (PEFT) methods like LoRA/DoRA. This hides a fatal subspace misalignment problem.

Compression picks directions that minimize reconstruction error — important for *pretraining*, not the downstream task.
Adaptation needs to update directions where the downstream loss gradient is sensitive.

If the retained subspace is orthogonal to the task's critical gradient directions, no amount of later fine-tuning can recover the discarded task-relevant information. Mathematically, SVD keeps the span of top singular vectors of $W \in \mathbb{R}^{d_{out} \times d_{in}}$ , but the important directions of $\partial L/\partial W$ may be entirely different.

Core Idea: A Task-aware Union of Subspaces

JACTUS rests on three insights: 1. Weight and task subspaces provide complementary signals. 2. Task sensitivity is two-sided: both input-side and output-side second-order statistics matter for a linear layer $$y = Wx$$ . 3. The parameter budget should be dynamically allocated across layers.

K-FAC Statistics

JACTUS estimates task-aware second-order statistics via K-FAC:

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 the Kronecker factors of the Fisher information matrix: $F_W \approx C_g \otimes C_x$ . Selecting their principal eigenspaces amounts to choosing important directions under the downstream-loss-induced two-sided Fisher metric.

The Union

Bases are extracted separately and merged by orthogonalization:

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

where $$(U_w, V_w)$$ come from SVD of $$W$$ and $$(U_g, V_x)$$ from eigendecomposition of regularized $(\tilde{C}_g, \tilde{C}_x)$ . This union covers both weight-important and task-sensitive directions.

Within the joint subspace, JACTUS computes a projected low-rank approximation $W_{proj} = Q_L^\top W Q_R$ , whose squared singular values form a per-layer marginal-gain curve.

Global Rank Allocation

Rather than uniform per-layer budgets, JACTUS allocates rank greedily by marginal gain per parameter cost:

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

A max-heap picks the best gain/cost increment until the budget is exhausted. On ViT-Large, MLP Up/Down matrices consistently receive higher rank than Q/K/V/Projection, and rank tends to grow with depth.

Core Matrix Optimization

Theorem 1 (subspace invariance): any full-rank bases spanning the same subspaces reparameterize the same achievable solution set. So after fixing $$Q_L, Q_R$$ and per-layer ranks, only a small $k_p \times k_p$ core matrix $$S$$ needs optimizing. This yields:

PEFT-level training efficiency (training only $$S$$ )
Compact low-rank deployment (no frozen full weights kept)
Full fine-tuning-level expressivity within the fixed subspace

Results

ViT-Base (ImageNet-21k pretrain), 8 datasets, 80% budget:

JACTUS: 89.2% average accuracy
DoRA (100% params): 87.9%; LoRA (100% params): 87.4%
At 40% budget: JACTUS 83.5% vs SVD-LLM-style compress-then-adapt 78.2%

LLaMA-2-7B, commonsense QA, 80% budget:

JACTUS: 80.9% vs DoRA 79.7%, LoRA 79.1% (both 100% params)

Ablations:

Union subspace (89.2%) > task-only (86.1%) > weight-only SVD (85.3%)
Global greedy allocation beats uniform allocation by 1.7 points (89.2% vs 87.5%)
~1K calibration samples suffice for stable covariance estimates

Why It Works

Traditional SVD assumes "useful information" equals large singular values — true only for reconstructing the weights. For downstream tasks, useful information should be defined by the geometry of the task loss. JACTUS injects that geometry ( $$C_x$$ , $$C_g$$ ) into compression: keep not only the "loud" directions of the weights, but also the directions the loss is "sensitive" to. The global allocator is essentially a marginal-analysis / Pareto-optimal resource allocation.

Open Directions

Multi-task: conflicting task subspaces need more sophisticated fusion
Dynamic compression: input-dependent effective rank at inference
Joint with quantization/pruning: extending subspace-awareness to combined compression pipelines
Theory: approximation guarantees for the union subspace and greedy allocation

References

JACTUS: Ge et al., "Compress Then Adapt? No, Do It Together via Task-aware Union of Subspaces", arXiv:2605.02829, 2026
LoRA: Hu et al., ICLR 2022
DoRA: Liu et al., ICML 2024
K-FAC: Martens & Grosse, ICML 2015
SVD-LLM: Wang et al., 2024