adu2021 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.
# 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.
# 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.
# 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.
# 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.
# 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