Spectral Normalization

Problem

Implement spectral normalization for a weight matrix. Spectral normalization constrains the spectral norm (largest singular value) of a weight matrix to 1, which stabilizes GAN training by enforcing Lipschitz continuity.

Solution

import numpy as np

def spectral_normalization(W: np.ndarray, num_iters: int = 1, u: np.ndarray = None) -> tuple[np.ndarray, np.ndarray]:
    height = W.shape[0]
    width = np.prod(W.shape[1:])
    W_mat = W.reshape(height, width)

    if u is None:
        u = np.random.randn(height)
        u = u / np.linalg.norm(u)

    for _ in range(num_iters):
        # Power iteration
        v = W_mat.T @ u
        v = v / np.linalg.norm(v)
        u = W_mat @ v
        u = u / np.linalg.norm(u)

    sigma = u @ W_mat @ v
    W_normalized = W / sigma

    return W_normalized, u

Explanation

1. Reshape the weight tensor into a 2D matrix for SVD-related computation.
2. Use power iteration to approximate the largest singular value: alternate computing v = W^T u and u = W v, normalizing each time.
3. The spectral norm sigma is then u^T W v.
4. Divide the entire weight matrix by sigma so the spectral norm becomes 1.
5. Return the updated u vector so it can be reused across training steps for efficiency.

Complexity

• Time: O(num_iters m n) where m, n are the matrix dimensions
• Space: O(m + n) for the u and v vectors