メニュー
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.
デフォルトでは、最初にソースを確認する Prompt が選択されています。直接コマンドに切り替えるか、ローカルコピーをダウンロードすることもできます。
インストールを決める前に、SKILL.md と SkillsMP に表示されている付属ファイルをお読みください。
Codex または Claude でインストール この Prompt をコピーして Codex、Claude、または他のアシスタントに貼り付けると、Skill ページを確認してインストールできます。
直接コマンドでは確認用 Prompt が省略されます。実行前にソースを確認してください。
npx skills add https://github.com/ADu2021/skillXiv --skill msign-stable-rank-restorationコマンドは1行のまま表示されます。コピー前に横へスクロールして全体を確認してください。
ローカルで確認しますか？SkillsMP が現在取得できるファイルをダウンロードできます。
SKILL.md を表示中
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. |
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.
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.
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).
"""
gram = weight_matrix.T @ weight_matrix
trace_gram = torch.trace(gram)
gram_squared = gram @ gram
trace_gram2 = torch.trace(gram_squared)
stable_rank = (trace_gram ** ) / (trace_gram2 + )
stable_rank.item()
():
jac_i_flat = weight_matrix_i.flatten()
jac_j_flat = weight_matrix_j.flatten()
alignment = torch.nn.functional.cosine_similarity(
jac_i_flat.unsqueeze(),
jac_j_flat.unsqueeze()
).item()
alignment
():
.step +=
.step % .check_interval != :
, {}
diagnostics = {: {}, : []}
min_rank = ()
name, param model.named_parameters():
(param.shape) >= :
sr = .compute_stable_rank(param.data)
diagnostics[][name] = sr
min_rank = (min_rank, sr)
param_list = [p p model.parameters() (p.shape) >= ]
i ((param_list) - ):
alignment = .compute_jacobian_alignment(
param_list[i], param_list[i + ]
)
diagnostics[].append(alignment)
avg_alignment = (diagnostics[]) / (
diagnostics[]
) diagnostics[]
is_at_risk = (min_rank < ) (avg_alignment > )
is_at_risk, diagnostics
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
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
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
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
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:
Key empirical findings:
Common pitfalls:
Validation: Recommended to combine with gradient clipping (norm=1.0) and learning rate warmup for maximum stability.
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