Skip to main content Inicio Creadores adu2021 skillxiv msign-stable-rank-restoration
msign-stable-rank-restoration 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.
Ir a la instalación Skills Marketplace Descubre y explora habilidades de IA creadas por la comunidad.
Instalar con Codex o Claude Copia este prompt, pégalo en Codex, Claude u otro asistente, y deja que revise la página de la skill y la instale por ti.
Mostrar detalles del prompt Un comando directo omite el prompt de revisión. Revisa el origen antes de ejecutarlo.
npx skills add https://github.com/ADu2021/skillXiv --skill msign-stable-rank-restorationEl comando permanece en una sola línea. Desplázate horizontalmente para revisarlo antes de copiarlo.
¿Prefieres una copia local? Descarga los archivos que SkillsMP tiene disponibles ahora.
Descargar Zip Descargando... Ocupaciones relacionadas SOC
Basado en la clasificación ocupacional SOC
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 layersTheory : Prove causal chain from rank → Jacobian → gradient explosionSolution : Periodic SVD to compute W_sign(W) where all singular values → 1Integration : Drop-in modification to existing optimizersCost : <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.
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).
"""
2
1e-8
return
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.
"""
0
0
return
def
check_stability
self, model
"""
Monitor stable rank across all weight matrices.
Returns: is_at_risk, diagnostics
"""
self
1
if
self
self
0
return
False
'stable_ranks'
'alignments'
float
'inf'
for
in
if
len
2
self
'stable_ranks'
min
for
in
if
len
2
for
in
range
len
1
self
1
'alignments'
sum
'alignments'
len
'alignments'
if
'alignments'
else
0
0.5
and
0.7
return
Step 2: Implement matrix sign operation Compute W_sign via SVD to equalize singular values.
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
"""
U, S, Vh = torch.linalg.svd(weight_matrix, full_matrices=False )
S_new = torch.ones_like(S)
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).
"""
W_norm = torch.norm(weight_matrix)
Y = weight_matrix / (W_norm + 1e-8 )
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.
class MSignOptimizer (torch.optim.Adam):
def __init__ (
self,
params,
lr=1e-3 ,
betas=(0.9 , 0.999 1e-8 ,
rank_restore_interval=100 ,
apply_to_layers=None
):
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
super ().step()
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
if self .apply_to_layers is not None :
should_update = any (
layer_name in str (p)
for layer_name in self .apply_to_layers
)
if not should_update:
continue
with torch.no_grad():
W_sign = matrix_sign_svd(p.data, num_iterations=3 )
blend_ratio = 0.1
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.
class SelectiveMSignOptimizer (MSignOptimizer ):
def __init__ (self, model, *args, **kwargs ):
super ().__init__(model.parameters(), *args, **kwargs)
self .model = model
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():
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
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
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.
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):
input_ids = batch['input_ids' ].to(device)
labels = batch['labels' ].to(device)
outputs = model(input_ids)
logits = outputs.logits
loss = criterion(
logits.view(-1 , logits.shape[-1 ]),
labels.view(-1 )
)
optimizer.zero_grad()
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0 )
optimizer.step()
total_loss += loss.item()
num_batches += 1
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:.4 f} " )
print (f"Epoch {epoch + 1 } : Avg Loss={total_loss / num_batches:.4 f} \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.
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
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
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
