Msign Stable Rank Restoration

---
name: msign-stable-rank-restoration
title: "MSign: An Optimizer Preventing Training Instability via Stable Rank Restoration"
version: 0.0.2
engine: skillxiv-v0.0.2-claude-opus-4.6
license: MIT
url: "https://arxiv.org/abs/2602.01734"
keywords: [Optimizer, Training Stability, Gradient Explosions, Matrix Decomposition, Numerical Stability]
description: "Prevent unrecoverable gradient explosions in LLM training by periodically restoring weight matrix stable rank through SVD-based matrix sign operations, eliminating sudden training failures without computational burden."
---

# MSign: An Optimizer Preventing Training Instability via Stable Rank Restoration

## Problem Context

LLM pretraining exhibits sudden, unrecoverable gradient explosions that waste significant computational resources. These failures lack obvious early warning and occur after thousands of training steps. Understanding their root cause requires analyzing weight matrix properties. Low stable rank (concentration of singular values) combined with high layer Jacobian alignment creates conditions for exponential gradient growth.

## Core Concept

MSign identifies [stable rank collapse, Jacobian alignment, causal mechanisms] as precursors to training failures. The optimizer applies [periodic matrix sign operations, SVD-based restoration, lightweight updates] to restore stable rank by equalizing non-zero singular values to 1, breaking the chain that leads to gradient explosions.

## Architecture Overview

- **Diagnosis**: Monitor stable rank of weight matrices and Jacobian alignment between layers
- **Theory**: Prove causal chain from rank → Jacobian → gradient explosion
- **Solution**: Periodic SVD to compute W_sign(W) where all singular values → 1
- **Integration**: Drop-in modification to existing optimizers
- **Cost**: <7% throughput reduction; can be applied selectively to attention layers only

## Implementation

### Step 1: Analyze stable rank and diagnose failure risk

Compute stable rank metrics to identify when training is at risk.

```python
# Stable rank diagnostics
class StableRankMonitor:
    def __init__(self, check_interval=100):
        self.check_interval = check_interval
        self.step = 0
        self.stable_ranks = []
        self.jacobian_alignments = []

    def compute_stable_rank(self, weight_matrix):
        """
        Compute stable rank: (trace(M^T M))^2 / trace((M^T M)^2)
        Measures concentration of singular values.
        High stable rank = evenly distributed singular values (good).
        Low stable rank = concentrated singular values (bad).
        """
        # Compute Gram matrix M^T M
        gram = weight_matrix.T @ weight_matrix

        # Trace of Gram
        trace_gram = torch.trace(gram)

        # Trace of Gram^2
        gram_squared = gram @ gram
        trace_gram2 = torch.trace(gram_squared)

        # Stable rank
        stable_rank = (trace_gram ** 2) / (trace_gram2 + 1e-8)

        return stable_rank.item()

    def compute_jacobian_alignment(self, weight_matrix_i, weight_matrix_j):
        """
        Compute cosine similarity between Jacobians of adjacent layers.
        Measures how aligned the gradients are across layers.
        """
        # Jacobians are weight matrices themselves (simplified)
        # In practice, compute Hessian eigenvectors
        jac_i_flat = weight_matrix_i.flatten()
        jac_j_flat = weight_matrix_j.flatten()

        # Cosine similarity
        alignment = torch.nn.functional.cosine_similarity(
            jac_i_flat.unsqueeze(0),
            jac_j_flat.unsqueeze(0)
        ).item()

        return alignment

    def check_stability(self, model):
        """
        Monitor stable rank across all weight matrices.
        Returns: is_at_risk, diagnostics
        """
        self.step += 1

        if self.step % self.check_interval != 0:
            return False, {}

        diagnostics = {'stable_ranks': {}, 'alignments': []}
        min_rank = float('inf')

        # Check all weight matrices
        for name, param in model.named_parameters():
            if len(param.shape) >= 2:  # Only matrices
                sr = self.compute_stable_rank(param.data)
                diagnostics['stable_ranks'][name] = sr
                min_rank = min(min_rank, sr)

        # Check Jacobian alignments (simplified: use weight matrix correlation)
        param_list = [p for p in model.parameters() if len(p.shape) >= 2]
        for i in range(len(param_list) - 1):
            alignment = self.compute_jacobian_alignment(
                param_list[i], param_list[i + 1]
            )
            diagnostics['alignments'].append(alignment)

        # Risk detection: low rank + high alignment
        avg_alignment = sum(diagnostics['alignments']) / len(
            diagnostics['alignments']
        ) if diagnostics['alignments'] else 0

        is_at_risk = (min_rank < 0.5) and (avg_alignment > 0.7)

        return is_at_risk, diagnostics
```

### Step 2: Implement matrix sign operation

Compute W_sign via SVD to equalize singular values.

```python
# Matrix sign via SVD
def matrix_sign_svd(weight_matrix, num_iterations=5):
    """
    Compute sign(W) = W @ (W^T W)^{-1/2}
    Equalizes all singular values to 1, restoring stable rank.

    Args:
        weight_matrix: Tensor of shape (out_features, in_features)
        num_iterations: Newton-Schulz iterations for inverse square root

    Returns:
        W_sign: Matrix with singular values = 1
    """
    # SVD decomposition: W = U @ S @ V^T
    U, S, Vh = torch.linalg.svd(weight_matrix, full_matrices=False)

    # Singular values > 0 by definition
    # Create new singular values: all ones
    S_new = torch.ones_like(S)

    # Reconstruct: W_sign = U @ I @ V^T = U @ V^T
    W_sign = U @ Vh

    return W_sign

def matrix_sign_newton_schulz(weight_matrix, num_iterations=3):
    """
    Compute matrix sign via Newton-Schulz iteration (more efficient for large matrices).
    """
    # Initialize: Y_0 = W / ||W||
    W_norm = torch.norm(weight_matrix)
    Y = weight_matrix / (W_norm + 1e-8)

    # Newton-Schulz iteration: Y_{n+1} = Y_n @ (I + Y_n^T Y_n)^{-1} / 2
    I = torch.eye(weight_matrix.shape[0], device=weight_matrix.device)

    for _ in range(num_iterations):
        Y_T_Y = Y.T @ Y
        inv_term = torch.linalg.inv(I + Y_T_Y)
        Y = 0.5 * Y @ inv_term

    return Y
```

### Step 3: Integrate stable rank restoration into optimizer step

Periodically apply matrix sign operation to maintain stable rank.

```python
# MSign optimizer
class MSignOptimizer(torch.optim.Adam):
    def __init__(
        self,
        params,
        lr=1e-3,
        betas=(0.9, 0.999),
        eps=1e-8,
        rank_restore_interval=100,
        apply_to_layers=None  # None = all layers, else specific layer names
    ):
        super().__init__(params, lr=lr, betas=betas, eps=eps)
        self.rank_restore_interval = rank_restore_interval
        self.step_count = 0
        self.apply_to_layers = apply_to_layers
        self.monitor = StableRankMonitor(check_interval=rank_restore_interval)

    def step(self, closure=None):
        """
        Single optimization step with periodic stable rank restoration.
        """
        loss = None
        if closure is not None:
            loss = closure()

        self.step_count += 1

        # Standard Adam step
        super().step()

        # Periodic stable rank restoration
        if self.step_count % self.rank_restore_interval == 0:
            self._restore_stable_rank()

        return loss

    def _restore_stable_rank(self):
        """
        Apply matrix sign operation to restore stable rank.
        """
        for group in self.param_groups:
            for p in group['params']:
                if len(p.shape) < 2:
                    continue  # Skip vectors

                # Check if this layer should be updated
                if self.apply_to_layers is not None:
                    # Layer filtering logic (simplified)
                    should_update = any(
                        layer_name in str(p)
                        for layer_name in self.apply_to_layers
                    )
                    if not should_update:
                        continue

                # Apply matrix sign operation
                with torch.no_grad():
                    W_sign = matrix_sign_svd(p.data, num_iterations=3)

                    # Blend with original: soft update to avoid disruption
                    blend_ratio = 0.1  # Conservative 10% update
                    p.data = (1 - blend_ratio) * p.data + blend_ratio * W_sign
```

### Step 4: Apply selectively to attention layers

For efficiency, apply restoration only to critical layers.

```python
# Selective MSign for attention layers
class SelectiveMSignOptimizer(MSignOptimizer):
    def __init__(self, model, *args, **kwargs):
        super().__init__(model.parameters(), *args, **kwargs)
        self.model = model

        # Identify attention layer names
        self.attention_layers = self._identify_attention_layers(model)

    def _identify_attention_layers(self, model):
        """Identify which layers are attention layers."""
        attention_layers = []

        for name, module in model.named_modules():
            # Check for attention patterns in name
            if any(
                pattern in name.lower()
                for pattern in ['attention', 'attn', 'self_attn', 'query', 'key', 'value']
            ):
                attention_layers.append(name)

        return attention_layers

    def _restore_stable_rank(self):
        """
        Apply matrix sign operation only to attention layers.
        """
        for group in self.param_groups:
            for p in group['params']:
                if len(p.shape) < 2:
                    continue

                # Check if parameter belongs to attention layer
                param_name = None
                for name, param in self.model.named_parameters():
                    if param is p:
                        param_name = name
                        break

                if param_name is None:
                    continue

                is_attention_param = any(
                    layer_name in param_name
                    for layer_name in self.attention_layers
                )

                if not is_attention_param:
                    continue

                # Apply restoration
                with torch.no_grad():
                    W_sign = matrix_sign_svd(p.data, num_iterations=2)
                    blend_ratio = 0.1
                    p.data = (1 - blend_ratio) * p.data + blend_ratio * W_sign
```

### Step 5: Training with MSign optimizer

Complete training loop using MSign.

```python
# Training with MSign
def train_with_msign(
    model, train_loader, device='cuda',
    rank_restore_interval=100, apply_to_attention_only=True
):
    """
    Train LLM using MSign optimizer for stability.
    """
    if apply_to_attention_only:
        optimizer = SelectiveMSignOptimizer(
            model,
            lr=1e-3,
            rank_restore_interval=rank_restore_interval
        )
    else:
        optimizer = MSignOptimizer(
            model.parameters(),
            lr=1e-3,
            rank_restore_interval=rank_restore_interval
        )

    model = model.to(device)
    criterion = torch.nn.CrossEntropyLoss()

    num_epochs = 3
    for epoch in range(num_epochs):
        total_loss = 0.0
        num_batches = 0

        for batch_idx, batch in enumerate(train_loader):
            # Move batch to device
            input_ids = batch['input_ids'].to(device)
            labels = batch['labels'].to(device)

            # Forward pass
            outputs = model(input_ids)
            logits = outputs.logits

            # Loss computation
            loss = criterion(
                logits.view(-1, logits.shape[-1]),
                labels.view(-1)
            )

            # Backward and optimization step
            optimizer.zero_grad()
            loss.backward()

            # Gradient clipping for extra safety
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)

            optimizer.step()

            total_loss += loss.item()
            num_batches += 1

            # Periodic diagnostics
            if (batch_idx + 1) % 500 == 0:
                avg_loss = total_loss / num_batches
                print(f"Epoch {epoch + 1}, Batch {batch_idx + 1}: "
                      f"Loss={avg_loss:.4f}")

        print(f"Epoch {epoch + 1}: Avg Loss={total_loss / num_batches:.4f}\n")

    return model
```

## Practical Guidance

**When to use**: Large-scale LLM training (1B+) prone to sudden gradient explosions. Most effective with dense, well-initialized models where failure is infrequent but catastrophic.

**Hyperparameters**:
- **rank_restore_interval**: 50-200 steps
  - Every 100 typical
  - More frequent for unstable training
  - Less frequent for stable training to reduce overhead
- **blend_ratio**: 0.05-0.15 (how much of the sign matrix to use)
  - Conservative 0.1 recommended
- **apply_to_attention_only**: True for 2-3% throughput cost, False for 7% cost

**Key empirical findings**:
- Prevents gradient explosions completely (100% success rate in tested scenarios)
- Throughput reduction: <7% overhead
- Works across dense and MoE models
- Selective attention-layer application provides good cost-benefit

**Common pitfalls**:
- blend_ratio too high → disrupts learned weights
- rank_restore_interval too frequent → computational waste
- Applying to all layers on large models → 10%+ overhead
- Not combining with gradient clipping → less robust

**Validation**: Recommended to combine with gradient clipping (norm=1.0) and learning rate warmup for maximum stability.

## Reference

Paper: https://arxiv.org/abs/2602.01734
Code: Available at author's repository
Theoretical analysis: Stable rank, Jacobian alignment, gradient explosion chains
Metrics: Training curves, failure-free iterations, convergence speed