梯度下降算法全解：从 SGD 到 Adam 的优化之路

import numpy as np

# 批量梯度下降
def batch_gd(X, y, lr=0.01, epochs=100):
    """使用全部数据计算梯度"""
    m = X.shape[0]
    theta = np.zeros(X.shape[1])
    for _ in range(epochs):
        gradient = (1/m) * X.T @ (X @ theta - y)
        theta -= lr * gradient
    return theta

# 随机梯度下降
def sgd(X, y, lr=0.01, epochs=10):
    """每次只用一个样本"""
    m = X.shape[0]
    theta = np.zeros(X.shape[1])
    for _ in range(epochs):
        for i in range(m):
            gradient = X[i:i+1].T @ (X[i:i+1] @ theta - y[i:i+1])
            theta -= lr * gradient
    return theta

# 小批量梯度下降
def mini_batch_gd(X, y, lr=0.01, epochs=50, batch_size=32):
    m = X.shape[0]
    theta = np.zeros(X.shape[1])
    for _ in range(epochs):
        indices = np.random.permutation(m)
        X_shuffled, y_shuffled = X[indices], y[indices]
        for i in range(0, m, batch_size):
            X_batch = X_shuffled[i:i+batch_size]
            y_batch = y_shuffled[i:i+batch_size]
            gradient = (1/batch_size) * X_batch.T @ (X_batch @ theta - y_batch)
            theta -= lr * gradient
    return theta

class MomentumOptimizer:
    """带动量的 SGD 优化器"""
    def __init__(self, lr=0.01, momentum=0.9):
        self.lr = lr
        self.momentum = momentum
        self.velocity = None
    
    def step(self, params, gradients):
        if self.velocity is None:
            self.velocity = np.zeros_like(params)
        
        # 经典动量: v = βv - η∇L
        self.velocity = self.momentum * self.velocity - self.lr * gradients
        params = params + self.velocity
        return params

class NesterovOptimizer:
    """Nesterov 加速梯度下降"""
    def __init__(self, lr=0.01, momentum=0.9):
        self.lr = lr
        self.momentum = momentum
        self.velocity = None
    
    def step(self, params, gradients):
        if self.velocity is None:
            self.velocity = np.zeros_like(params)
        
        # Nesterov: 先用速度走半步，再计算梯度
        self.velocity = self.momentum * self.velocity - self.lr * gradients
        params = params + self.momentum * self.velocity - self.lr * gradients
        return params

class AdaGradOptimizer:
    def __init__(self, lr=0.01, eps=1e-8):
        self.lr = lr
        self.eps = eps
        self.cache = None
    
    def step(self, params, gradients):
        if self.cache is None:
            self.cache = np.zeros_like(params)
        self.cache += gradients ** 2
        params -= self.lr * gradients / (np.sqrt(self.cache) + self.eps)
        return params

class RMSPropOptimizer:
    def __init__(self, lr=0.01, beta2=0.999, eps=1e-8):
        self.lr = lr
        self.beta2 = beta2
        self.eps = eps
        self.cache = None
    
    def step(self, params, gradients):
        if self.cache is None:
            self.cache = np.zeros_like(params)
        # 指数移动平均，非单调递增
        self.cache = self.beta2 * self.cache + (1 - self.beta2) * gradients ** 2
        params -= self.lr * gradients / (np.sqrt(self.cache) + self.eps)
        return params

class AdamOptimizer:
    """Adam 优化器完整实现"""
    def __init__(self, lr=0.001, beta1=0.9, beta2=0.999, eps=1e-8):
        self.lr = lr
        self.beta1 = beta1
        self.beta2 = beta2
        self.eps = eps
        self.m = None  # 一阶矩
        self.v = None  # 二阶矩
        self.t = 0     # 时间步
    
    def step(self, params, gradients):
        if self.m is None:
            self.m = np.zeros_like(params)
            self.v = np.zeros_like(params)
        
        self.t += 1
        
        # 一阶矩（指数移动平均）
        self.m = self.beta1 * self.m + (1 - self.beta1) * gradients
        # 二阶矩（指数移动平均）
        self.v = self.beta2 * self.v + (1 - self.beta2) * gradients ** 2
        
        # 偏差修正
        m_hat = self.m / (1 - self.beta1 ** self.t)
        v_hat = self.v / (1 - self.beta2 ** self.t)
        
        # 参数更新
        params -= self.lr * m_hat / (np.sqrt(v_hat) + self.eps)
        return params

class AdamWOptimizer:
    """AdamW: 解耦权重衰减的 Adam"""
    def __init__(self, lr=0.001, beta1=0.9, beta2=0.999, eps=1e-8, weight_decay=0.01):
        self.lr = lr
        self.beta1 = beta1
        self.beta2 = beta2
        self.eps = eps
        self.weight_decay = weight_decay
        self.m = None
        self.v = None
        self.t = 0
    
    def step(self, params, gradients):
        if self.m is None:
            self.m = np.zeros_like(params)
            self.v = np.zeros_like(params)
        
        self.t += 1
        self.m = self.beta1 * self.m + (1 - self.beta1) * gradients
        self.v = self.beta2 * self.v + (1 - self.beta2) * gradients ** 2
        
        m_hat = self.m / (1 - self.beta1 ** self.t)
        v_hat = self.v / (1 - self.beta2 ** self.t)
        
        # 解耦权重衰减
        params -= self.lr * (m_hat / (np.sqrt(v_hat) + self.eps) + self.weight_decay * params)
        return params

class LionOptimizer:
    """Lion: 符号动量优化器"""
    def __init__(self, lr=0.0003, beta1=0.95, beta2=0.98, weight_decay=0.1):
        self.lr = lr
        self.beta1 = beta1
        self.beta2 = beta2
        self.weight_decay = weight_decay
        self.m = None
    
    def step(self, params, gradients):
        if self.m is None:
            self.m = np.zeros_like(params)
        
        update = self.beta2 * self.m + (1 - self.beta2) * gradients
        self.m = self.beta1 * self.m + (1 - self.beta1) * gradients
        
        # 只取符号 + 权重衰减
        params -= self.lr * (np.sign(update) + self.weight_decay * params)
        return params

# 学习率调度示例（PyTorch 风格）
import math

def cosine_decay(initial_lr, current_step, total_steps):
    """余弦衰减调度"""
    return initial_lr * 0.5 * (1 + math.cos(math.pi * current_step / total_steps))

def warmup_cosine(initial_lr, current_step, total_steps, warmup_steps):
    """Warmup + 余弦衰减"""
    if current_step < warmup_steps:
        # Warmup 阶段：线性增长
        return initial_lr * (current_step / warmup_steps)
    else:
        # 余弦衰减阶段
        progress = (current_step - warmup_steps) / (total_steps - warmup_steps)
        return initial_lr * 0.5 * (1 + math.cos(math.pi * progress))

# 可视化调度曲线
total_steps = 1000
warmup_steps = 100
initial_lr = 0.001

print("Warmup + Cosine 调度曲线:")
for step in [0, 50, 100, 200, 500, 800, 1000]:
    lr = warmup_cosine(initial_lr, step, total_steps, warmup_steps)
    print(f"  step {step:4d} → lr = {lr:.6f}")

# 条件数对优化收敛的影响演示
import numpy as np

def optimize_quadratic(a, b, lr=0.01, steps=100, use_momentum=False, beta=0.9):
    """在二次函数 f(x,y)=0.5*(a*x^2+b*y^2) 上优化"""
    x, y = 3.0, 3.0
    vx, vy = 0.0, 0.0
    for _ in range(steps):
        gx, gy = a * x, b * y  # 梯度
        if use_momentum:
            vx = beta * vx - lr * gx
            vy = beta * vy - lr * gy
            x, y = x + vx, y + vy
        else:
            x, y = x - lr * gx, y - lr * gy
    return x, y

# 高条件数场景 (a=100, b=1)
a, b = 100, 1
sgd_final = optimize_quadratic(a, b, lr=0.005, steps=200)
mom_final = optimize_quadratic(a, b, lr=0.005, steps=200, use_momentum=True)
print(f"条件数: {a/b}")
print(f"SGD 最终位置: ({sgd_final[0]:.6f}, {sgd_final[1]:.6f})")
print(f"Momentum 最终位置: ({mom_final[0]:.6f}, {mom_final[1]:.6f})")

# LLM 训练典型 AdamW 配置
llm_optimizer_config = {
    "lr": 1e-4,           # 学习率（通常 1e-4 ~ 3e-4）
    "betas": (0.9, 0.95), # β₁=0.9, β₂=0.95（注意 β₂<0.999）
    "eps": 1e-8,          # 数值稳定性
    "weight_decay": 0.01, # 权重衰减
    "max_grad_norm": 1.0, # 梯度裁剪
}

def get_llm_lr(step, total_steps, warmup_steps=2000, target_lr=1e-4):
    """LLM 学习率调度：Warmup + Cosine Annealing"""
    import math
    if step < warmup_steps:
        return target_lr * (step / warmup_steps)
    else:
        progress = (step - warmup_steps) / (total_steps - warmup_steps)
        return target_lr * 0.5 * (1 + math.cos(math.pi * progress))

for step in [0, 500, 1000, 2000, 5000, 10000, 50000]:
    lr = get_llm_lr(step, total_steps=50000)
    print(f"  step {step:>5d} → lr = {lr:.2e}")