Awesome Kernels Blog

Implementing the MMD Two-Sample Test in JAX

2024-01-25T00:00:00.000Z

The Maximum Mean Discrepancy (MMD) is a powerful kernel-based statistic for comparing distributions. In this post, we'll implement the MMD two-sample test from scratch using JAX and explore its applications.

The Two-Sample Testing Problem

Given samples:

$X = \{x_1, \ldots, x_n\} \sim P$
$Y = \{y_1, \ldots, y_m\} \sim Q$

We want to test:

$H_0: P = Q$ (distributions are the same)
$H_1: P \neq Q$ (distributions differ)

Maximum Mean Discrepancy

The MMD measures the distance between mean embeddings:

$\text{MMD}^2[P, Q] = \|\mu_P - \mu_Q\|_{\mathcal{H}}^2$

Expanding:

$\text{MMD}^2 = \mathbb{E}[k(x, x')] + \mathbb{E}[k(y, y')] - 2\mathbb{E}[k(x, y)]$

where $x, x' \sim P$ and $y, y' \sim Q$ .

Empirical MMD Estimators

Biased Estimator (V-statistic)

$\widehat{\text{MMD}}_b^2 = \frac{1}{n^2}\sum_{i,j} k(x_i, x_j) + \frac{1}{m^2}\sum_{i,j} k(y_i, y_j) - \frac{2}{nm}\sum_{i,j} k(x_i, y_j)$

Unbiased Estimator (U-statistic)

$\widehat{\text{MMD}}_u^2 = \frac{1}{n(n-1)}\sum_{i \neq j} k(x_i, x_j) + \frac{1}{m(m-1)}\sum_{i \neq j} k(y_i, y_j) - \frac{2}{nm}\sum_{i,j} k(x_i, y_j)$

Linear-Time Estimator

$\widehat{\text{MMD}}_l^2 = \frac{2}{n}\sum_{i=1}^{n/2}[k(x_{2i-1}, x_{2i}) + k(y_{2i-1}, y_{2i}) - k(x_{2i-1}, y_{2i}) - k(x_{2i}, y_{2i-1})]$

JAX Implementation

import jax
import jax.numpy as jnp
from jax import random, Array
from typing import Optional, Dict, Any


def rbf_kernel(X: Array, Y: Array, gamma: float) -> Array:
    """Compute RBF kernel matrix."""
    sq_norm_X = jnp.sum(X**2, axis=1, keepdims=True)
    sq_norm_Y = jnp.sum(Y**2, axis=1, keepdims=True)
    sq_dist = sq_norm_X + sq_norm_Y.T - 2 * jnp.dot(X, Y.T)
    return jnp.exp(-gamma * sq_dist)


def median_heuristic(X: Array, Y: Array) -> float:
    """Compute kernel bandwidth using the median heuristic."""
    XY = jnp.vstack([X, Y])
    n = len(XY)
    sq_norm = jnp.sum(XY**2, axis=1, keepdims=True)
    sq_dist = sq_norm + sq_norm.T - 2 * jnp.dot(XY, XY.T)
    mask = jnp.triu(jnp.ones((n, n), dtype=bool), k=1)
    dists = jnp.sqrt(jnp.maximum(sq_dist[mask], 0.0))
    median_dist = jnp.median(dists)
    return 1.0 / (2 * median_dist**2 + 1e-10)


class MMDTest:
    """
    Maximum Mean Discrepancy Two-Sample Test in JAX.
    
    Parameters
    ----------
    kernel : str, default='rbf'
        Kernel function name
    gamma : float, default=None
        RBF kernel bandwidth. If None, uses median heuristic.
    """
    
    def __init__(self, kernel: str = 'rbf', gamma: Optional[float] = None):
        self.kernel = kernel
        self.gamma = gamma
        self.gamma_: Optional[float] = None
        
    def _compute_kernel(self, X: Array, Y: Array) -> Array:
        """Compute kernel matrix."""
        if self.kernel == 'rbf':
            return rbf_kernel(X, Y, self.gamma_)
        else:
            raise ValueError(f"Unknown kernel: {self.kernel}")
    
    def compute_mmd_squared(
        self, X: Array, Y: Array, unbiased: bool = True
    ) -> float:
        """Compute MMD^2 statistic."""
        if self.gamma is None:
            self.gamma_ = median_heuristic(X, Y)
        else:
            self.gamma_ = self.gamma
            
        n, m = len(X), len(Y)
        
        K_XX = self._compute_kernel(X, X)
        K_YY = self._compute_kernel(Y, Y)
        K_XY = self._compute_kernel(X, Y)
        
        if unbiased:
            term_XX = (jnp.sum(K_XX) - jnp.trace(K_XX)) / (n * (n - 1))
            term_YY = (jnp.sum(K_YY) - jnp.trace(K_YY)) / (m * (m - 1))
            term_XY = jnp.sum(K_XY) / (n * m)
        else:
            term_XX = jnp.mean(K_XX)
            term_YY = jnp.mean(K_YY)
            term_XY = jnp.mean(K_XY)
            
        return float(term_XX + term_YY - 2 * term_XY)
    
    def permutation_test(
        self,
        X: Array,
        Y: Array,
        n_permutations: int = 1000,
        alpha: float = 0.05,
        random_state: Optional[int] = None,
    ) -> Dict[str, Any]:
        """Perform MMD permutation test."""
        n, m = len(X), len(Y)
        observed_mmd = self.compute_mmd_squared(X, Y)
        combined = jnp.vstack([X, Y])
        
        key = random.PRNGKey(random_state or 0)
        null_distribution = []
        
        for i in range(n_permutations):
            key, subkey = random.split(key)
            perm = random.permutation(subkey, n + m)
            X_perm = combined[perm[:n]]
            Y_perm = combined[perm[n:]]
            null_distribution.append(self.compute_mmd_squared(X_perm, Y_perm))
        
        null_distribution = jnp.array(null_distribution)
        p_value = float(
            (jnp.sum(null_distribution >= observed_mmd) + 1) / (n_permutations + 1)
        )
        threshold = float(jnp.quantile(null_distribution, 1 - alpha))
        
        return {
            'statistic': observed_mmd,
            'p_value': p_value,
            'reject_null': p_value < alpha,
            'threshold': threshold,
            'null_distribution': null_distribution,
        }

Example: Testing Identical Distributions

Under $H_0$ , the MMD should be close to zero:

import matplotlib.pyplot as plt
from jax import random

key = random.PRNGKey(42)
key1, key2 = random.split(key)

# Same distribution (H0 is true)
n, m = 200, 200
X = random.normal(key1, (n, 2))
Y = random.normal(key2, (m, 2))

mmd_test = MMDTest()
result = mmd_test.permutation_test(X, Y, random_state=42)

print(f"MMD² = {result['statistic']:.6f}")
print(f"p-value = {result['p_value']:.4f}")
print(f"Reject H₀: {result['reject_null']}")

# Plot null distribution
plt.figure(figsize=(10, 5))
plt.hist(result['null_distribution'], bins=50, density=True, alpha=0.7, label='Null distribution')
plt.axvline(result['statistic'], color='red', linestyle='--', linewidth=2, label=f'Observed MMD² = {result["statistic"]:.4f}')
plt.axvline(result['threshold'], color='green', linestyle=':', linewidth=2, label=f'Threshold (α=0.05)')
plt.xlabel('MMD²')
plt.ylabel('Density')
plt.title('MMD Test: Same Distribution (H₀ true)')
plt.legend()
plt.show()

Example: Testing Different Distributions

# Different distributions (H1 is true)
key = random.PRNGKey(42)
key1, key2 = random.split(key)

X = random.normal(key1, (200, 2))
Y = random.normal(key2, (200, 2)) + jnp.array([1.0, 0.5])  # Shifted mean

mmd_test = MMDTest()
result = mmd_test.permutation_test(X, Y, random_state=42)

print(f"MMD² = {result['statistic']:.6f}")
print(f"p-value = {result['p_value']:.4f}")
print(f"Reject H₀: {result['reject_null']}")

# Visualize
fig, axes = plt.subplots(1, 2, figsize=(14, 5))

# Scatter plot
axes[0].scatter(X[:, 0], X[:, 1], alpha=0.5, label='P')
axes[0].scatter(Y[:, 0], Y[:, 1], alpha=0.5, label='Q')
axes[0].legend()
axes[0].set_title('Samples from P and Q')
axes[0].set_xlabel('x₁')
axes[0].set_ylabel('x₂')

# Null distribution
axes[1].hist(result['null_distribution'], bins=50, density=True, alpha=0.7, label='Null distribution')
axes[1].axvline(result['statistic'], color='red', linestyle='--', linewidth=2, label=f'Observed MMD²')
axes[1].axvline(result['threshold'], color='green', linestyle=':', linewidth=2, label='Threshold')
axes[1].set_xlabel('MMD²')
axes[1].set_ylabel('Density')
axes[1].set_title('MMD Test: Different Distributions (H₁ true)')
axes[1].legend()

plt.tight_layout()
plt.show()

Power Analysis

Test power increases with:

Sample size
Distribution separation
Appropriate kernel bandwidth

def compute_power(dist_shift, n_samples, n_trials=100):
    """Estimate test power for given shift and sample size."""
    rejections = 0
    
    for trial in range(n_trials):
        key = random.PRNGKey(trial)
        key1, key2 = random.split(key)
        
        X = random.normal(key1, (n_samples, 2))
        Y = random.normal(key2, (n_samples, 2)) + dist_shift
        
        mmd_test = MMDTest()
        result = mmd_test.permutation_test(X, Y, n_permutations=200, random_state=trial)
        
        if result['reject_null']:
            rejections += 1
    
    return rejections / n_trials

# Power vs shift
shifts = jnp.linspace(0, 1.5, 10)
powers = [compute_power(jnp.array([s, 0]), 100, n_trials=50) for s in shifts]

plt.figure(figsize=(8, 5))
plt.plot(shifts, powers, 'bo-', linewidth=2, markersize=8)
plt.axhline(0.05, color='gray', linestyle='--', label='α = 0.05')
plt.xlabel('Distribution Shift')
plt.ylabel('Power (rejection rate)')
plt.title('MMD Test Power vs Distribution Shift')
plt.legend()
plt.grid(True, alpha=0.3)
plt.show()

JAX Benefits for MMD

Using JAX for MMD computation provides:

JIT Compilation: Speed up permutation tests with jax.jit
Parallelization: Use jax.vmap to vectorize over permutations
GPU Acceleration: Compute kernel matrices on GPU for large samples
Automatic Differentiation: Optimize kernel parameters if needed

Applications

1. GAN Evaluation

Compare generated samples to real data:

def evaluate_gan(real_samples, generated_samples):
    """Evaluate GAN using MMD."""
    mmd_test = MMDTest()
    result = mmd_test.permutation_test(real_samples, generated_samples)
    print(f"MMD² = {result['statistic']:.6f}")
    print(f"p-value = {result['p_value']:.4f}")
    return result

2. Covariate Shift Detection

def detect_covariate_shift(train_features, test_features):
    """Detect if train and test distributions differ."""
    mmd_test = MMDTest()
    result = mmd_test.permutation_test(train_features, test_features)
    
    if result['reject_null']:
        print("⚠️ Covariate shift detected!")
    else:
        print("✓ No significant distribution shift")
    return result

3. Feature Selection

def feature_importance_mmd(X, Y):
    """Rank features by their contribution to MMD."""
    importances = []
    
    for i in range(X.shape[1]):
        mmd_test = MMDTest()
        mmd_i = mmd_test.compute_mmd_squared(X[:, i:i+1], Y[:, i:i+1])
        importances.append((i, mmd_i))
    
    return sorted(importances, key=lambda x: x[1], reverse=True)

Conclusion

The MMD two-sample test is:

Powerful: Can detect any difference between distributions (with characteristic kernel)
Nonparametric: No distributional assumptions
Interpretable: MMD has clear geometric meaning

Key takeaways:

Use unbiased estimator for hypothesis testing
Median heuristic is a good default for bandwidth
Permutation test controls Type I error correctly

The complete JAX implementation is available in our examples folder.

Next: Implementing kernel mean embeddings for conditional distributions!

Random Fourier Features - Scalable Kernel Approximation

2024-01-20T00:00:00.000Z

Random Fourier Features (RFF) is a beautiful technique that approximates kernel functions with explicit feature maps, enabling linear-time training while maintaining kernel-like performance. Let's implement it from scratch using JAX!

The Bottleneck of Kernel Methods

Standard kernel methods require:

Computing and storing an $n \times n$ kernel matrix: $O(n^2)$ space
Solving linear systems: $O(n^3)$ time

For millions of data points, this becomes impractical.

Bochner's Theorem to the Rescue

For shift-invariant kernels $K(x, y) = k(x - y)$ , Bochner's theorem tells us:

$k(\delta) = \int p(\omega) e^{i\omega^T\delta} d\omega$

where $p(\omega)$ is the spectral density (Fourier transform of $k$ ).

The Key Insight

We can approximate the integral via Monte Carlo:

$k(x-y) \approx \frac{1}{D}\sum_{j=1}^D e^{i\omega_j^T(x-y)}$

where $\omega_1, \ldots, \omega_D \sim p(\omega)$ .

Using Euler's formula and taking real parts:

$k(x-y) \approx z(x)^T z(y)$

where:

$z(x) = \sqrt{\frac{2}{D}}[\cos(\omega_1^T x + b_1), \ldots, \cos(\omega_D^T x + b_D)]$

with $b_j \sim \text{Uniform}(0, 2\pi)$ .

JAX Implementation

import jax
import jax.numpy as jnp
from jax import random, Array
from typing import Optional, Literal

KernelType = Literal["rbf", "laplacian"]


class RandomFourierFeatures:
    """
    Random Fourier Features for kernel approximation in JAX.
    
    Parameters
    ----------
    n_features : int, default=100
        Number of random features
    gamma : float, default=1.0
        RBF kernel parameter
    kernel : str, default='rbf'
        Kernel type: 'rbf', 'laplacian'
    random_state : int, default=None
        Random seed for reproducibility
    """
    
    def __init__(
        self,
        n_features: int = 100,
        gamma: float = 1.0,
        kernel: KernelType = "rbf",
        random_state: Optional[int] = None,
    ):
        self.n_features = n_features
        self.gamma = gamma
        self.kernel = kernel
        self.random_state = random_state
        self.omega_: Optional[Array] = None
        self.b_: Optional[Array] = None
        
    def fit(self, X: Array) -> "RandomFourierFeatures":
        """Sample random frequencies from the spectral density."""
        key = random.PRNGKey(self.random_state or 0)
        n_features_in = X.shape[1]
        key1, key2 = random.split(key)
        
        # Sample frequencies from spectral density
        if self.kernel == "rbf":
            # Gaussian RBF: spectral density is Gaussian
            self.omega_ = (
                random.normal(key1, (n_features_in, self.n_features))
                * jnp.sqrt(2 * self.gamma)
            )
        elif self.kernel == "laplacian":
            # Laplacian kernel: spectral density is Cauchy
            self.omega_ = (
                random.cauchy(key1, (n_features_in, self.n_features)) * self.gamma
            )
        else:
            raise ValueError(f"Unknown kernel: {self.kernel}")
        
        # Sample phase offsets uniformly
        self.b_ = random.uniform(key2, (self.n_features,)) * 2 * jnp.pi
        
        return self
    
    def transform(self, X: Array) -> Array:
        """Compute random Fourier features."""
        if self.omega_ is None or self.b_ is None:
            raise ValueError("Transformer not fitted. Call fit() first.")
        projection = jnp.dot(X, self.omega_) + self.b_
        Z = jnp.sqrt(2.0 / self.n_features) * jnp.cos(projection)
        return Z
    
    def fit_transform(self, X: Array) -> Array:
        """Fit and transform in one step."""
        return self.fit(X).transform(X)


class RFFRidgeRegression:
    """
    Ridge regression with Random Fourier Features in JAX.
    
    Parameters
    ----------
    n_features : int, default=100
        Number of random features
    gamma : float, default=1.0
        RBF kernel parameter
    alpha : float, default=1.0
        Regularization strength
    """
    
    def __init__(
        self,
        n_features: int = 100,
        gamma: float = 1.0,
        alpha: float = 1.0,
        random_state: Optional[int] = None,
    ):
        self.n_features = n_features
        self.gamma = gamma
        self.alpha = alpha
        self.random_state = random_state
        self.rff_: Optional[RandomFourierFeatures] = None
        self.weights_: Optional[Array] = None
        
    def fit(self, X: Array, y: Array) -> "RFFRidgeRegression":
        """Fit the model."""
        self.rff_ = RandomFourierFeatures(
            n_features=self.n_features, 
            gamma=self.gamma,
            random_state=self.random_state,
        )
        Z = self.rff_.fit_transform(X)
        
        # Solve regularized least squares in feature space
        A = jnp.dot(Z.T, Z) + self.alpha * jnp.eye(self.n_features)
        b = jnp.dot(Z.T, y)
        self.weights_ = jax.scipy.linalg.solve(A, b, assume_a="pos")
        
        return self
    
    def predict(self, X: Array) -> Array:
        """Make predictions."""
        if self.rff_ is None or self.weights_ is None:
            raise ValueError("Model not fitted. Call fit() first.")
        Z = self.rff_.transform(X)
        return jnp.dot(Z, self.weights_)

Comparing Exact vs Approximate Kernels

Let's verify the approximation quality with JAX:

import matplotlib.pyplot as plt
from jax import random

# Generate data
key = random.PRNGKey(42)
n = 500
d = 10
X = random.normal(key, (n, d))

# Exact RBF kernel
gamma = 0.5
sq_norm = jnp.sum(X**2, axis=1, keepdims=True)
sq_dist = sq_norm + sq_norm.T - 2 * jnp.dot(X, X.T)
K_exact = jnp.exp(-gamma * sq_dist)

# Approximate kernel with different D
D_values = [10, 50, 100, 500, 1000]
errors = []

for D in D_values:
    rff = RandomFourierFeatures(n_features=D, gamma=gamma, random_state=42)
    Z = rff.fit_transform(X)
    K_approx = jnp.dot(Z, Z.T)
    
    error = jnp.linalg.norm(K_exact - K_approx, 'fro') / jnp.linalg.norm(K_exact, 'fro')
    errors.append(float(error))
    print(f"D = {D:4d}: Relative error = {error:.4f}")

plt.figure(figsize=(8, 5))
plt.loglog(D_values, errors, 'bo-', linewidth=2, markersize=8)
plt.xlabel('Number of Random Features (D)')
plt.ylabel('Relative Frobenius Error')
plt.title('RFF Approximation Quality')
plt.grid(True, alpha=0.3)
plt.show()

Example: Regression with RFF

# Generate nonlinear regression data
key = random.PRNGKey(42)
X_train = jnp.sort(random.uniform(key, (200,), minval=0, maxval=2*jnp.pi)).reshape(-1, 1)
key, subkey = random.split(key)
y_train = jnp.sin(X_train).ravel() + 0.1 * random.normal(subkey, (200,))

X_test = jnp.linspace(0, 2*jnp.pi, 300).reshape(-1, 1)

# Compare different numbers of features
fig, axes = plt.subplots(2, 2, figsize=(12, 10))

for ax, D in zip(axes.flat, [10, 50, 200, 1000]):
    model = RFFRidgeRegression(n_features=D, gamma=1.0, alpha=0.001, random_state=42)
    model.fit(X_train, y_train)
    y_pred = model.predict(X_test)
    
    ax.scatter(X_train, y_train, c='blue', alpha=0.3, s=20)
    ax.plot(X_test, y_pred, 'r-', linewidth=2, label='RFF prediction')
    ax.plot(X_test, jnp.sin(X_test), 'g--', linewidth=2, label='True function')
    ax.set_title(f'D = {D} features')
    ax.legend()

plt.tight_layout()
plt.show()

Computational Comparison

Method	Training	Prediction	Storage
Exact Kernel	$O(n^3)$	$O(n)$ per sample	$O(n^2)$
RFF	$O(nD^2)$	$O(D)$ per sample	$O(nD)$

For $D \ll n$ , RFF is dramatically faster!

import time

n_samples = [100, 500, 1000, 2000, 5000]
D = 500

times_exact = []
times_rff = []

for n in n_samples:
    key = random.PRNGKey(42)
    X = random.normal(key, (n, 10))
    y = random.normal(key, (n,))
    
    # Exact kernel
    start = time.time()
    sq_norm = jnp.sum(X**2, axis=1, keepdims=True)
    sq_dist = sq_norm + sq_norm.T - 2 * jnp.dot(X, X.T)
    K = jnp.exp(-0.5 * sq_dist)
    alpha = jax.scipy.linalg.solve(K + 0.01 * jnp.eye(n), y)
    jax.block_until_ready(alpha)  # Wait for JAX async computation
    times_exact.append(time.time() - start)
    
    # RFF
    start = time.time()
    model = RFFRidgeRegression(n_features=D, gamma=0.5, alpha=0.01)
    model.fit(X, y)
    jax.block_until_ready(model.weights_)
    times_rff.append(time.time() - start)
    
    print(f"n = {n:5d}: Exact = {times_exact[-1]:.3f}s, RFF = {times_rff[-1]:.3f}s")

plt.figure(figsize=(8, 5))
plt.plot(n_samples, times_exact, 'bo-', label='Exact Kernel', linewidth=2)
plt.plot(n_samples, times_rff, 'ro-', label=f'RFF (D={D})', linewidth=2)
plt.xlabel('Number of Samples')
plt.ylabel('Training Time (seconds)')
plt.legend()
plt.title('Scalability Comparison')
plt.show()

JAX Benefits for RFF

Using JAX for Random Fourier Features provides:

JIT Compilation: Compile the transform function for faster repeated calls
Automatic Batching: Use vmap for vectorized operations
GPU Acceleration: Seamlessly run on GPUs for large datasets
Composability: Easy to chain with other JAX transformations

When to Use RFF

✅ Use RFF when:

Dataset is large ( $n > 10,000$ )
Shift-invariant kernels (RBF, Laplacian, Matérn)
Online/streaming learning needed
Memory is limited

❌ Avoid RFF when:

Small datasets (exact kernel is fine)
Non-shift-invariant kernels (polynomial)
Very high accuracy required

Conclusion

Random Fourier Features transform kernel methods from $O(n^3)$ to $O(nD^2)$ , making them practical for large-scale learning. The key insights are:

Bochner's theorem connects kernels to Fourier transforms
Monte Carlo approximation gives explicit features
Standard linear methods can then be applied

The complete JAX implementation is available in our examples folder.

Next post: Implementing the Nyström approximation for data-dependent features!

Implementing Kernel Ridge Regression from Scratch

2024-01-15T00:00:00.000Z

Kernel Ridge Regression (KRR) is one of the most elegant kernel methods, combining the simplicity of ridge regression with the power of the kernel trick. In this post, we'll implement KRR from scratch using JAX and understand each component.

The Problem

Given training data $(x_1, y_1), \ldots, (x_n, y_n)$ , we want to find a function $f$ that minimizes:

$\min_{f \in \mathcal{H}} \frac{1}{n}\sum_{i=1}^n (y_i - f(x_i))^2 + \lambda \|f\|_{\mathcal{H}}^2$

where $\mathcal{H}$ is a Reproducing Kernel Hilbert Space.

The Representer Theorem

Thanks to the representer theorem, we know the solution has the form:

$f^*(x) = \sum_{i=1}^n \alpha_i K(x_i, x)$

This transforms our infinite-dimensional optimization into a finite-dimensional one!

Deriving the Solution

Substituting into the objective:

$L(\alpha) = \frac{1}{n}\|y - K\alpha\|^2 + \lambda \alpha^T K \alpha$

Taking the derivative and setting to zero:

$\frac{\partial L}{\partial \alpha} = -\frac{2}{n}K(y - K\alpha) + 2\lambda K\alpha = 0$

Solving (assuming $K$ is invertible or using regularization):

$\alpha = (K + n\lambda I)^{-1}y$

JAX Implementation

import jax
import jax.numpy as jnp
from jax import Array
from typing import Optional, Literal

KernelType = Literal["rbf", "linear", "polynomial"]


def rbf_kernel(X1: Array, X2: Array, gamma: float = 1.0) -> Array:
    """Compute RBF (Gaussian) kernel matrix: K(x, y) = exp(-gamma * ||x - y||^2)"""
    sq_norm1 = jnp.sum(X1**2, axis=1, keepdims=True)
    sq_norm2 = jnp.sum(X2**2, axis=1, keepdims=True)
    sq_dist = sq_norm1 + sq_norm2.T - 2 * jnp.dot(X1, X2.T)
    return jnp.exp(-gamma * sq_dist)


def linear_kernel(X1: Array, X2: Array) -> Array:
    """Compute linear kernel matrix: K(x, y) = x^T y"""
    return jnp.dot(X1, X2.T)


def polynomial_kernel(
    X1: Array, X2: Array, gamma: float = 1.0, coef0: float = 1.0, degree: int = 3
) -> Array:
    """Compute polynomial kernel: K(x, y) = (gamma * x^T y + coef0)^degree"""
    return (gamma * jnp.dot(X1, X2.T) + coef0) ** degree


class KernelRidgeRegression:
    """
    Kernel Ridge Regression implementation in JAX.
    
    Parameters
    ----------
    kernel : str, default='rbf'
        Kernel type: 'rbf', 'linear', 'polynomial'
    gamma : float, default=1.0
        RBF kernel parameter
    degree : int, default=3
        Polynomial kernel degree
    coef0 : float, default=1.0
        Polynomial kernel coefficient
    alpha : float, default=1.0
        Regularization strength
    """
    
    def __init__(
        self,
        kernel: KernelType = "rbf",
        gamma: float = 1.0,
        degree: int = 3,
        coef0: float = 1.0,
        alpha: float = 1.0,
    ):
        self.kernel = kernel
        self.gamma = gamma
        self.degree = degree
        self.coef0 = coef0
        self.alpha = alpha
        self.X_train_: Optional[Array] = None
        self.dual_coef_: Optional[Array] = None
        
    def _compute_kernel(self, X1: Array, X2: Array) -> Array:
        """Compute kernel matrix between X1 and X2."""
        if self.kernel == "rbf":
            return rbf_kernel(X1, X2, self.gamma)
        elif self.kernel == "linear":
            return linear_kernel(X1, X2)
        elif self.kernel == "polynomial":
            return polynomial_kernel(X1, X2, self.gamma, self.coef0, self.degree)
        else:
            raise ValueError(f"Unknown kernel: {self.kernel}")
    
    def fit(self, X: Array, y: Array) -> "KernelRidgeRegression":
        """Fit the model using JAX's linear solver."""
        self.X_train_ = X
        n = len(X)
        K = self._compute_kernel(X, X)
        # Use JAX's optimized linear solver
        self.dual_coef_ = jax.scipy.linalg.solve(
            K + n * self.alpha * jnp.eye(n), y, assume_a="pos"
        )
        return self
    
    def predict(self, X: Array) -> Array:
        """Make predictions."""
        if self.X_train_ is None or self.dual_coef_ is None:
            raise ValueError("Model not fitted. Call fit() first.")
        K = self._compute_kernel(X, self.X_train_)
        return jnp.dot(K, self.dual_coef_)

Example: Fitting a Nonlinear Function

import matplotlib.pyplot as plt
from jax import random

# Generate synthetic data
key = random.PRNGKey(42)
X = jnp.linspace(0, 2*jnp.pi, 100).reshape(-1, 1)
y = jnp.sin(X).ravel() + 0.1 * random.normal(key, (100,))

# Fit model
krr = KernelRidgeRegression(kernel='rbf', gamma=1.0, alpha=0.01)
krr.fit(X, y)

# Predict
X_test = jnp.linspace(0, 2*jnp.pi, 200).reshape(-1, 1)
y_pred = krr.predict(X_test)

# Plot
plt.figure(figsize=(10, 6))
plt.scatter(X, y, c='blue', alpha=0.5, label='Training data')
plt.plot(X_test, y_pred, 'r-', linewidth=2, label='KRR prediction')
plt.plot(X_test, jnp.sin(X_test), 'g--', linewidth=2, label='True function')
plt.legend()
plt.xlabel('x')
plt.ylabel('y')
plt.title('Kernel Ridge Regression (JAX)')
plt.show()

Understanding the Kernel Parameter

The kernel width $\gamma$ controls the flexibility:

Large $\gamma$ : Local influence, can fit complex patterns (may overfit)
Small $\gamma$ : Global influence, smoother predictions (may underfit)

fig, axes = plt.subplots(1, 3, figsize=(15, 4))

for ax, gamma in zip(axes, [0.1, 1.0, 10.0]):
    krr = KernelRidgeRegression(kernel='rbf', gamma=gamma, alpha=0.01)
    krr.fit(X, y)
    y_pred = krr.predict(X_test)
    
    ax.scatter(X, y, c='blue', alpha=0.3)
    ax.plot(X_test, y_pred, 'r-', linewidth=2)
    ax.set_title(f'γ = {gamma}')
    
plt.tight_layout()
plt.show()

Connection to Gaussian Processes

KRR is equivalent to the posterior mean of a Gaussian Process with:

Prior: $f \sim \mathcal{GP}(0, K/\lambda)$
Likelihood: $y | f \sim \mathcal{N}(f(X), \sigma^2 I)$

where $\sigma^2 = n\lambda$ .

JAX Benefits

Using JAX for kernel methods provides several advantages:

JIT Compilation: Functions can be compiled for faster execution
Automatic Differentiation: Easy to compute gradients for optimization
GPU/TPU Support: Seamlessly scale to accelerators
Functional Design: Clean, composable code

Computational Considerations

The main bottleneck is solving the $n \times n$ linear system:

Time complexity: $O(n^3)$
Space complexity: $O(n^2)$

For large datasets, consider:

Nyström approximation: Use subset of data
Random Fourier Features: Explicit feature map
Iterative solvers: Conjugate gradient

Next Steps

In upcoming posts, we'll implement:

Random Fourier Features for scalable KRR
Kernel PCA for dimensionality reduction
MMD for two-sample testing

Full Code

The complete JAX implementation is available in our examples folder.

Have questions? Open an issue on GitHub or join the discussion!

Awesome Kernels Blog

Implementing the MMD Two-Sample Test in JAX

The Two-Sample Testing Problem​

Maximum Mean Discrepancy​

Empirical MMD Estimators​

Biased Estimator (V-statistic)​

Unbiased Estimator (U-statistic)​

Linear-Time Estimator​

JAX Implementation​

Example: Testing Identical Distributions​

Example: Testing Different Distributions​

Power Analysis​

JAX Benefits for MMD​

Applications​

1. GAN Evaluation​

2. Covariate Shift Detection​

3. Feature Selection​

Conclusion​

Random Fourier Features - Scalable Kernel Approximation

The Bottleneck of Kernel Methods​

Bochner's Theorem to the Rescue​

The Key Insight​

JAX Implementation​

Comparing Exact vs Approximate Kernels​

Example: Regression with RFF​

Computational Comparison​

JAX Benefits for RFF​

When to Use RFF​

Conclusion​

Implementing Kernel Ridge Regression from Scratch

The Problem​

The Representer Theorem​

Deriving the Solution​

JAX Implementation​

Example: Fitting a Nonlinear Function​

Understanding the Kernel Parameter​

Connection to Gaussian Processes​

JAX Benefits​

Computational Considerations​

Next Steps​

Full Code​

The Two-Sample Testing Problem

Maximum Mean Discrepancy

Empirical MMD Estimators

Biased Estimator (V-statistic)

Unbiased Estimator (U-statistic)

Linear-Time Estimator

JAX Implementation

Example: Testing Identical Distributions

Example: Testing Different Distributions

Power Analysis

JAX Benefits for MMD

Applications

1. GAN Evaluation

2. Covariate Shift Detection

3. Feature Selection

Conclusion

The Bottleneck of Kernel Methods

Bochner's Theorem to the Rescue

The Key Insight

JAX Implementation

Comparing Exact vs Approximate Kernels

Example: Regression with RFF

Computational Comparison

JAX Benefits for RFF

When to Use RFF

Conclusion

The Problem

The Representer Theorem

Deriving the Solution

JAX Implementation

Example: Fitting a Nonlinear Function

Understanding the Kernel Parameter

Connection to Gaussian Processes

JAX Benefits

Computational Considerations

Next Steps

Full Code