Reverse-mode vs Forward-mode Autodiff: JVP, VJP, Dual Numbers, and Which to Use When in LLMs

import numpy as np
import sympy as sp
 
# 1. Sembolik
x_sym = sp.Symbol('x')
f_sym = sp.sin(x_sym) * sp.exp(x_sym)
df_sym = sp.diff(f_sym, x_sym)
print("Sembolik f'(x):", df_sym)
# cos(x)*exp(x) + sin(x)*exp(x)
f_at_2 = float(df_sym.subs(x_sym, 2.0))
print(f"f'(2) sembolik = {f_at_2:.6f}")
 
# 2. Numerik (finite difference)
def f(x):
    return np.sin(x) * np.exp(x)
 
h = 1e-5
df_num = (f(2.0 + h) - f(2.0 - h)) / (2 * h)
print(f"f'(2) numerik   = {df_num:.6f}")
 
# 3. Otomatik türev (JAX ile)
import jax
import jax.numpy as jnp
 
def f_jax(x):
    return jnp.sin(x) * jnp.exp(x)
 
df_auto = jax.grad(f_jax)(2.0)
print(f"f'(2) autodiff  = {df_auto:.6f}")
 
# Hepsi aynı değer — ama performans karakteristikleri çok farklı

import math
 
class Dual:
    """Dual number for forward-mode autodiff."""
    def __init__(self, primal, tangent=0.0):
        self.p = primal     # değer
        self.t = tangent    # türev
 
    def __repr__(self):
        return f"Dual(primal={self.p}, tangent={self.t})"
 
    def __add__(self, other):
        if isinstance(other, Dual):
            return Dual(self.p + other.p, self.t + other.t)
        return Dual(self.p + other, self.t)
 
    def __mul__(self, other):
        if isinstance(other, Dual):
            # (x + ε ẋ)(y + ε ẏ) = xy + ε(ẋy + xẏ)
            return Dual(self.p * other.p, self.t * other.p + self.p * other.t)
        return Dual(self.p * other, self.t * other)
 
    def __pow__(self, n):
        # d/dx x^n = n x^(n-1)
        return Dual(self.p ** n, n * self.p ** (n - 1) * self.t)
 
    def __radd__(self, other): return self + other
    def __rmul__(self, other): return self * other
 
def sin(x):
    if isinstance(x, Dual):
        return Dual(math.sin(x.p), math.cos(x.p) * x.t)
    return math.sin(x)
 
def cos(x):
    if isinstance(x, Dual):
        return Dual(math.cos(x.p), -math.sin(x.p) * x.t)
    return math.cos(x)
 
# Kullanım: x̂ = 2 + ε·1, türevi hesapla
x = Dual(2.0, 1.0)  # tangent=1 → "x'e göre türev al"
y = x ** 2 + sin(x)
print(y)            # Dual(primal=4.909..., tangent=3.584...)
print(f"f(2) = {y.p:.6f}")
print(f"f'(2) = {y.t:.6f}")
 
# Doğrulama: f'(x) = 2x + cos(x) → f'(2) = 4 + cos(2) ≈ 3.584
print(f"Beklenen: {2*2 + math.cos(2):.6f}")

# Karşılaştırma: 100-input, 1-output (typical loss)
import jax
import jax.numpy as jnp
import time
 
def loss_fn(x):
    return jnp.sum(jnp.sin(x) ** 2 + x ** 2)
 
n = 100
x = jnp.array([0.5] * n)
 
# Reverse-mode (grad): 1 VJP, O(N)
grad_rev = jax.grad(loss_fn)
t0 = time.perf_counter()
g = grad_rev(x)
t1 = time.perf_counter()
print(f"Reverse-mode (grad): {(t1-t0)*1e6:.1f}μs, grad shape: {g.shape}")
 
# Forward-mode (n kez JVP): O(n·N)
def forward_grad(f, x):
    """Forward-mode ile gradient — n kez JVP."""
    grad = jnp.zeros_like(x)
    for i in range(len(x)):
        v = jnp.zeros_like(x).at[i].set(1.0)
        _, jvp_i = jax.jvp(f, (x,), (v,))
        grad = grad.at[i].set(jvp_i)
    return grad
 
t0 = time.perf_counter()
g_fwd = forward_grad(loss_fn, x)
t1 = time.perf_counter()
print(f"Forward-mode (n JVPs): {(t1-t0)*1e6:.1f}μs")
 
# n=100 → forward 100x yavaş
# n=10^9 (LLM) → forward imkansız
 
# Tersi senaryo: 1 input, 100 output
def multi_output(x):
    return jnp.array([jnp.sin(x * i) for i in range(100)])
 
# Forward: 1 JVP (1 input)
y, jvp_val = jax.jvp(multi_output, (0.5,), (1.0,))
print(f"Forward 1 JVP for m=100: instant")
 
# Reverse: 100 kez VJP (her output için)
# Çok daha yavaş

import jax
import jax.numpy as jnp
 
# Test function
def f(x):
    return jnp.sin(x[0]) * jnp.cos(x[1]) + x[2] ** 2
 
x = jnp.array([1.0, 2.0, 3.0])
 
# 1. Gradient (reverse-mode)
grad = jax.grad(f)
print("grad:", grad(x))  # ∇f at x
# array([0.4161..., -0.4546..., 6.0])
 
# 2. Value + Gradient (tek pass)
val, gr = jax.value_and_grad(f)(x)
print(f"f(x) = {val:.4f}, ∇f(x) = {gr}")
 
# 3. JVP: f(x + ε v) ≈ f(x) + ε · J·v
v = jnp.array([1.0, 0.0, 0.0])  # x[0] yönünde tangent
y, jvp_val = jax.jvp(f, (x,), (v,))
print(f"f(x) = {y:.4f}, JVP·v = {jvp_val:.4f}")
# JVP·v = ∂f/∂x[0] = cos(1)·cos(2) ≈ 0.416
 
# 4. VJP: u^T · J
y, vjp_fn = jax.vjp(f, x)
u = 1.0  # scalar output için
grad_via_vjp = vjp_fn(u)
print(f"VJP grad: {grad_via_vjp[0]}")
 
# 5. Jacobian (forward vs reverse) — multi-output
def g(x):
    return jnp.array([x[0]**2, x[1] * x[2], jnp.sin(x[0])])
 
J_fwd = jax.jacfwd(g)(x)  # 3 input, 3 output → 3x3 Jacobian
J_rev = jax.jacrev(g)(x)
print("Forward Jacobian:")
print(J_fwd)
print("Reverse Jacobian:")
print(J_rev)
print("Aynı mı:", jnp.allclose(J_fwd, J_rev))  # True
 
# 6. Hessian
H = jax.hessian(f)(x)
print("Hessian:")
print(H)  # 3x3 symmetric

import jax
import jax.numpy as jnp
 
def f(x):
    return jnp.sum(x ** 4) + jnp.sum(jnp.sin(x) ** 2)
 
x = jnp.array([1.0, 2.0, 3.0, 4.0])
v = jnp.array([1.0, 0.0, 1.0, 0.0])  # bazı yönlerde
 
# Hessian-vector product — forward-of-reverse
def hvp(f, x, v):
    """H(f, x) · v — tam Hessian saklamadan."""
    return jax.jvp(jax.grad(f), (x,), (v,))[1]
 
h_v = hvp(f, x, v)
print("H·v:", h_v)
 
# Doğrulama: tam Hessian'la
H = jax.hessian(f)(x)
h_v_direct = H @ v
print("H @ v (direct):", h_v_direct)
print("Aynı mı:", jnp.allclose(h_v, h_v_direct))
 
# Karmaşıklık: HVP O(N), tam Hessian O(N²)
# n=1000 için: HVP ~10μs, tam Hessian ~10ms (1000x fark)

import torch
from torch.func import grad, jvp, vjp, hessian, vmap
 
# Bir mini transformer'ın forward pass'i (toy)
def mini_attention(qkv_weights, x):
    """qkv_weights: (3 * d_model, d_model), x: (seq, d_model)"""
    d = x.shape[-1]
    qkv = x @ qkv_weights.T                # (seq, 3*d)
    q, k, v = qkv.chunk(3, dim=-1)
    scores = q @ k.T / (d ** 0.5)
    attn = torch.softmax(scores, dim=-1)
    return (attn @ v).sum()                # scalar loss
 
# Test data
torch.manual_seed(0)
d_model = 8
seq_len = 4
W = torch.randn(3 * d_model, d_model, requires_grad=True)
x = torch.randn(seq_len, d_model)
 
# Gradient w.r.t. W (reverse-mode — gradient için doğru seçim)
grad_W = grad(lambda W: mini_attention(W, x))(W)
print(f"grad shape: {grad_W.shape}")
 
# Hessian-vector product
v = torch.randn_like(W)
hv = jvp(grad(lambda W: mini_attention(W, x)), (W,), (v,))[1]
print(f"HVP shape: {hv.shape}")
 
# vmap ile batched gradient
batch_x = torch.randn(16, seq_len, d_model)
batched_grad = vmap(lambda x_b: grad(lambda W: mini_attention(W, x_b))(W))(batch_x)
print(f"Batched grad shape: {batched_grad.shape}")  # (16, 24, 8)