SLORR: Simple and Efficient In-Training Low-Rank Regularization

Overview§

SLORR introduces lightweight regularizers that encourage low-rank weight matrices during training, enabling better post-training factorization without architectural modifications.

Core Methodology§

Two variants based on differentiable surrogates for rank:

• SLORR-Hoyer: Uses the Hoyer sparsity metric $\text{Hoyer}(W) = \frac{\|W\|_1}{\|W\|_F}$, which is scale-invariant and smooth. The regularizer is $\mathcal{R}_{\text{Hoyer}} = \lambda \cdot \text{Hoyer}(W)$.
• SLORR-Nuclear: Approximates the nuclear norm $\|W\|_$ via a differentiable surrogate: $\|W\|_ \approx \|W\|_F^2 / \|W\|_2$ (using power iteration for spectral norm). The regularizer is $\mathcal{R}_{\text{nuc}} = \lambda \cdot \|W\|_*$.

Both are applied to original weight matrices (no extra parameters) and require only forward/backward passes that are GPU-friendly. The paper provides approximation guarantees and shows that the Hoyer variant works best in practice.

Implementation Details§

The regularizer is added to the standard training loss: $\mathcal{L}_{\text{total}} = \mathcal{L}_{\text{task}} + \alpha \cdot \mathcal{R}(W)$. The backward pass uses autograd; no manual gradient computation needed. The spectral norm for nuclear norm is estimated with 1-3 power iterations, which is cheap.

PyTorch Code Snippet§

class SLORRHoyer(torch.nn.Module):
    def __init__(self, lambd=1e-4):
        super().__init__()
        self.lambd = lambd

    def forward(self, W):
        l1 = W.norm(p=1)
        fro = W.norm(p='fro')
        hoyer = l1 / (fro + 1e-12)
        return self.lambd * hoyer

# Usage in training loop
model = YourModel()
regularizer = SLORRHoyer(lambd=1e-4)
optimizer = torch.optim.Adam(model.parameters())

for data, target in dataloader:
    optimizer.zero_grad()
    output = model(data)
    loss = nn.CrossEntropyLoss()(output, target)
    
    # Add regularization for all weight matrices
    reg_loss = 0
    for name, param in model.named_parameters():
        if 'weight' in name and param.dim() >= 2:
            reg_loss += regularizer(param)
    total_loss = loss + reg_loss
    total_loss.backward()
    optimizer.step()

Key Results§

• On ImageNet-1K (short-horizon continued training): ResNet-50, ViT-B/16, ViT-L/16 achieve higher post-SVD accuracy with <8% overhead.
• On LLM pretraining (135M, 560M): SLORR-Hoyer preserves perplexity after SVD (e.g., 50% compression) with <1% overhead.

Conclusion§

SLORR is a practical, plug-and-play regularizer for any neural network that can be trained with standard optimizers, improving compressibility without architectural changes.