adu2021 SGG: Fine-Grained Learning Rate Control Without Manual Tuning
Training large language models is sensitive to learning rate: too high causes divergence, too low wastes computation. Most adaptive optimizers apply the same scaling globally across all parameters, missing structure in how different parts of the network train. SGG (Scaling with Gradient Grouping) clusters parameters within each layer into groups based on gradient statistics, then applies tailored scaling per group. This balances per-layer constraints with per-parameter precision, converging faster and more stably across varying batch sizes and learning rates.
Core Concept
Parameter-adaptive learning rates are common, but ignoring within-layer structure misses optimization opportunities. Parameters in a layer often split into distinct groups: some with consistently high gradients (critical decision points), others with low gradients (refinement parameters). SGG identifies these groups dynamically and scales each differently, preventing any single group from dominating updates while ensuring critical parameters stay responsive.
Architecture Overview
- Gradient Statistics Collection: Track gradient magnitude, variance within each layer
- Dynamic Clustering: Partition parameters into K clusters based on gradient statistics (typically K=3-5)
- Cluster-Specific Scaling: Compute separate scaling factors per cluster, normalizing gradient magnitudes
- Optimizer Wrapper: Integrates with existing optimizers (AdamW) as a lightweight post-processing step
- Stability Monitoring: Tracks gradient norm evolution to detect and prevent divergence
Implementation
This implementation demonstrates SGG as an optimizer wrapper for improved training stability.
Build the gradient grouping analyzer:
import torch
import torch.nn as nn
import numpy as np
from typing import Dict, List, Tuple
from collections import defaultdict
class GradientGroupingAnalyzer:
"""Analyze gradient statistics and identify parameter groups within layers."""
def __init__(self, num_clusters: int = 3, window_size: int = 100):
self.num_clusters = num_clusters
self.window_size = window_size
self.gradient_history = defaultdict(list) # layer_name -> [grad_stats]
def compute_gradient_statistics(self, model: nn.Module) -> Dict:
"""
Compute gradient magnitude and variance per layer.
Returns dict: layer_name -> {mean, std, percentiles}
"""
layer_stats = {}
for name, param in model.named_parameters():
if param.grad is None:
continue
grad = param.grad.detach()
grad_abs = torch.abs(grad)
# Compute statistics
stats = {
"mean": float(grad_abs.mean()),
"std": float(grad_abs.std()),
"min": float(grad_abs.min()),
"max": float(grad_abs.max()),
"p25": float(torch.quantile(grad_abs, 0.25)),
"p50": float(torch.quantile(grad_abs, 0.50)),
"p75": float(torch.quantile(grad_abs, 0.75)),
}
# Extract layer name (remove .weight/.bias)
layer_name = ".".join(name.split(".")[:-1])
if layer_name not in layer_stats:
layer_stats[layer_name] = []
layer_stats[layer_name].append(stats)
return layer_stats
def cluster_parameters(self, model: nn.Module,
layer_name: str) -> Dict[int, List[str]]:
"""
Cluster parameters within a layer based on gradient statistics.
Returns dict: cluster_id -> [param_names]
"""
layer_params = {}
# Collect gradients for this layer
for name, param in model.named_parameters():
if layer_name in name and param.grad is not None:
grad = param.grad.detach().flatten()
layer_params[name] = grad
if not layer_params:
return {}
# Compute features for clustering
features = []
param_names = list(layer_params.keys())
for name in param_names:
grad = layer_params[name]
feature = [
float(torch.abs(grad).mean()), # Mean magnitude
float(torch.abs(grad).std()), # Variance
float((grad ** 2).mean()) ** 0.5 # RMS
]
features.append(feature)
features = np.array(features)
# K-means clustering
centroids = self._kmeans(features, self.num_clusters)
clusters = self._assign_clusters(features, centroids)
# Build cluster -> param mapping
cluster_mapping = defaultdict(list)
for param_name, cluster_id in zip(param_names, clusters):
cluster_mapping[cluster_id].append(param_name)
return dict(cluster_mapping)
def _kmeans(self, data: np.ndarray, k: int) -> np.ndarray:
"""Simple K-means clustering."""
# Initialize centroids randomly
indices = np.random.choice(len(data), k, replace=False)
centroids = data[indices].copy()
for iteration in range(10): # Fixed iterations
# Assign clusters
distances = np.linalg.norm(data[:, None] - centroids, axis=2)
assignments = np.argmin(distances, axis=1)
# Update centroids
new_centroids = np.array([
data[assignments == i].mean(axis=0) if (assignments == i).any()
else centroids[i]
for i in range(k)
])
if np.allclose(centroids, new_centroids):
break
centroids = new_centroids
return centroids
def _assign_clusters(self, data: np.ndarray, centroids: np.ndarray) -> np.ndarray:
"""Assign data points to nearest centroids."""
distances = np.linalg.norm(data[:, None] - centroids, axis=2)
return np.argmin(distances, axis=1)
# Example usage
analyzer = GradientGroupingAnalyzer(num_clusters=3)
# Create mock model
model = nn.Sequential(
nn.Linear(256, 512),
nn.ReLU(),
nn.Linear(512, 256)
)
# Simulate training step with gradients
dummy_input = torch.randn(4, 256)
output = model(dummy_input).sum()
output.backward()
# Analyze gradients
layer_stats = analyzer.compute_gradient_statistics(model)
for layer, stats_list in layer_stats.items():
print(f"{layer}: {len(stats_list)} parameter groups")
# Cluster parameters in first layer
clusters = analyzer.cluster_parameters(model, "0")
for cluster_id, param_names in clusters.items():
print(f"Cluster {cluster_id}: {len(param_names)} parameters")
Implement SGG optimizer wrapper:
class ScalingWithGradientGrouping(torch.optim.Optimizer):
"""
Optimizer wrapper that applies cluster-specific gradient scaling.
Wraps AdamW or other standard optimizer.
"""
def __init__(self, model: nn.Module, optimizer_class=torch.optim.AdamW,
lr: float = 1e-4, num_clusters: int = 3, eps: float = 1e-8):
self.model = model
self.base_optimizer = optimizer_class(model.parameters(), lr=lr)
self.analyzer = GradientGroupingAnalyzer(num_clusters=num_clusters)
self.eps = eps
self.layer_clusters = {}
self.step_count = 0
def zero_grad(self):
"""Clear gradients."""
self.base_optimizer.zero_grad()
def step(self, closure=None):
"""
Optimizer step with gradient grouping and scaling.
"""
self.step_count += 1
# Periodically recompute clusters (every 100 steps)
if self.step_count % 100 == 0:
self._update_clusters()
# Apply group-specific scaling
self._apply_group_scaling()
# Run base optimizer step
return self.base_optimizer.step(closure)
def _update_clusters(self):
"""Recompute parameter clusters based on current gradients."""
# Get layer names
layer_names = set()
for name, param in self.model.named_parameters():
if param.grad is not None:
layer_name = ".".join(name.split(".")[:-1])
layer_names.add(layer_name)
# Cluster each layer
for layer_name in layer_names:
self.layer_clusters[layer_name] = \
self.analyzer.cluster_parameters(self.model, layer_name)
def _apply_group_scaling(self):
"""Apply cluster-specific gradient scaling within each layer."""
for layer_name, clusters in self.layer_clusters.items():
for cluster_id, param_names in clusters.items():
# Compute gradient norm for this cluster
cluster_grad_norm = 0.0
for param_name in param_names:
param = self._get_parameter_by_name(param_name)
if param.grad is not None:
cluster_grad_norm += torch.sum(param.grad ** 2)
cluster_grad_norm = torch.sqrt(cluster_grad_norm + self.eps)
# Scale gradients: target norm of 1.0 per cluster
target_norm = 1.0
scale_factor = target_norm / (cluster_grad_norm + self.eps)
for param_name in param_names:
param = self._get_parameter_by_name(param_name)
if param.grad is not None:
param.grad.mul_(scale_factor)
def _get_parameter_by_name(self, param_name: str) -> torch.nn.Parameter:
"""Retrieve parameter from model by name."""
parts = param_name.split(".")
obj = self.model
for part in parts:
obj = getattr(obj, part)
return obj
# Test SGG optimizer
model = nn.Sequential(
nn.Linear(256, 512),
nn.ReLU(),
nn.Linear(512, 128)
)
optimizer = ScalingWithGradientGrouping(
model,
optimizer_class=torch.optim.AdamW,
lr=1e-4,
num_clusters=3
)
# Training loop
print("Training with SGG optimizer...")
for step in range(10):
optimizer.zero_grad()
# Forward pass
x = torch.randn(32, 256)
y = torch.randn(32, 128)
output = model(x)
loss = nn.functional.mse_loss(output, y)
# Backward pass
loss.backward()
# Optimizer step
optimizer.step()
if (step + 1) % 5 == 0:
print(f"Step {step+1}: Loss = {loss.item():.6f}")
Benchmark SGG against standard AdamW:
def benchmark_optimizers():
"""Compare SGG vs standard AdamW convergence."""
torch.manual_seed(42)
# Create models
model_sgg = nn.Sequential(nn.Linear(256, 512), nn.ReLU(), nn.Linear(512, 128))
model_adamw = nn.Sequential(nn.Linear(256, 512), nn.ReLU(), nn.Linear(512, 128))
# Copy weights to ensure same initialization
with torch.no_grad():
for p1, p2 in zip(model_sgg.parameters(), model_adamw.parameters()):
p2.copy_(p1)
# Optimizers
opt_sgg = ScalingWithGradientGrouping(model_sgg, lr=1e-3)
opt_adamw = torch.optim.AdamW(model_adamw.parameters(), lr=1e-3)
# Generate dummy data
X = torch.randn(1000, 256)
y = torch.randn(1000, 128)
losses_sgg = []
losses_adamw = []
# Training loop
for epoch in range(50):
# SGG training step
opt_sgg.zero_grad()
out_sgg = model_sgg(X)
loss_sgg = nn.functional.mse_loss(out_sgg, y)
loss_sgg.backward()
opt_sgg.step()
losses_sgg.append(loss_sgg.item())
# AdamW training step
opt_adamw.zero_grad()
out_adamw = model_adamw(X)
loss_adamw = nn.functional.mse_loss(out_adamw, y)
loss_adamw.backward()
opt_adamw.step()
losses_adamw.append(loss_adamw.item())
return losses_sgg, losses_adamw
losses_sgg, losses_adamw = benchmark_optimizers()
print(f"Final SGG loss: {losses_sgg[-1]:.6f}")
print(f"Final AdamW loss: {losses_adamw[-1]:.6f}")
print(f"SGG Improvement: {(losses_adamw[-1] / losses_sgg[-1] - 1) * 100:.1f}%")
Practical Guidance
| Aspect | Details |
|---|---|
| Number of Clusters | 3-5 typical; more clusters = finer control but higher overhead |
| Clustering Frequency | Update every 50-100 steps; balances adaptivity with compute |
| Learning Rate | Start with same LR as standard AdamW; may tolerate higher LR with SGG |
| Model Size | Benefits increase with model size; marginal gains on small models |
| Batch Size Sensitivity | SGG reduces LR tuning sensitivity across batch sizes |
When to Use:
- Training large language models on diverse batch sizes
- Sensitive to hyperparameter tuning; SGG reduces this burden
- Want faster convergence and better stability without architecture changes
- Compatible with other optimization techniques (gradient checkpointing, FSDP)
- Fine-tuning on multiple downstream tasks with single hyperparameter set
When NOT to Use:
- Simple models or datasets where standard AdamW already works well
- Real-time training with extremely tight latency constraints (clustering adds overhead)
- Already using highly specialized, hand-tuned learning rate schedules
- Theoretical analysis requires simple optimizer (stick with AdamW)
Common Pitfalls:
- Cluster count too high: wasted computation on fine-grained control with minimal benefit
- Clustering infrequently: outdated clusters become misaligned with actual gradient distribution
- Not accounting for warmup: clustering during warmup phase can be noisy; skip first N steps
- Ignoring divergence signals: monitor gradient norms; if still diverging, reduce learning rate globally
Reference
Taming LLMs by Scaling Learning Rates with Gradient Grouping https://arxiv.org/abs/2506.01049