Poet Orthogonal LLM Training

Reparameterized LLM Training via Orthogonal Equivalence Transformation

Core Concept

POET reparameterizes weight matrices as products of fixed random initialization and learnable orthogonal transformations, decoupling spectral control from optimization dynamics. Rather than optimizing weights W directly, POET learns two orthogonal matrices R and P such that W = RW₀P, where W₀ is fixed. This approach preserves singular values (spectrum) while optimizing singular vectors, improving training stability and generalization without increasing total parameters.

Architecture Overview

• Spectrum-Preserving Parameterization: W = RW₀P where W₀ is frozen random, R/P are learned orthogonal transformations
• Three Training Phases: Vector probing analysis reveals distinct phases—conical-shell searching, stable learning on shell, final adjustment
• Efficient Approximations: Stochastic Primitive Optimization (SPO) factorizes large orthogonal matrices into products of smaller primitives; Cayley-Neumann parameterization avoids expensive matrix inversions
• Memory Optimization: Merge-then-reinitialize trick consolidates learned transformations periodically, reducing GPU memory by 30%
• Generalization Guarantee: Maintains small hyperspherical energy under initialization, connecting to established generalization theory

Implementation

Step 1: Implement Orthogonal Matrix Parameterization

import torch
import torch.nn as nn
from torch.linalg import qr

class OrthogonalLinear(nn.Module):
    """
    Reparameterized linear layer: W = RW₀P
    Learns R and P as orthogonal transformations of fixed random W₀.
    """

    def __init__(self, in_features, out_features):
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features

        # Fixed random initialization (normalized Gaussian)
        W0 = torch.randn(out_features, in_features) / (in_features ** 0.5)
        self.register_buffer('W0', W0)

        # Learnable orthogonal factors
        self.R_param = nn.Parameter(torch.eye(out_features))
        self.P_param = nn.Parameter(torch.eye(in_features))

    def forward(self, x):
        # Reconstruct weight matrix from orthogonal factors
        W = self.R_param @ self.W0 @ self.P_param
        return torch.nn.functional.linear(x, W, None)

class StochasticPrimitiveOrthogonal(nn.Module):
    """
    Efficient approximation: factorize large orthogonal matrices into
    products of smaller Householder reflections or Givens rotations.
    Reduces parameters from O(mn) to O(m+n).
    """

    def __init__(self, size, num_primitives=3):
        super().__init__()
        self.size = size
        self.num_primitives = num_primitives
        self.primitives = nn.ParameterList()

        # Initialize as product of Householder reflections
        for _ in range(num_primitives):
            # Householder vector (normalized random vector)
            v = torch.randn(size)
            v = v / (torch.norm(v) + 1e-8)
            self.primitives.append(nn.Parameter(v))

    def forward(self):
        """Construct orthogonal matrix from primitive reflections."""
        Q = torch.eye(self.size, device=self.primitives[0].device)

        for v in self.primitives:
            v_norm = v / (torch.norm(v) + 1e-8)
            # Householder reflection: H = I - 2vv^T
            H = torch.eye(self.size, device=v.device) - 2 * torch.outer(v_norm, v_norm)
            Q = Q @ H

        return Q

Step 2: Cayley-Neumann Parameterization

class CayleyNeumannOrthogonal(nn.Module):
    """
    Approximate orthogonal matrix using Cayley transform with Neumann series.
    Avoids expensive matrix inversions: Q ≈ (I - A)(I + A)⁻¹
    approximated via truncated Neumann series.
    """

    def __init__(self, size, truncation_order=3):
        super().__init__()
        self.size = size
        self.truncation_order = truncation_order

        # Skew-symmetric matrix parameterization
        self.skew_params = nn.Parameter(torch.randn(size, size) * 0.01)

    def forward(self):
        # Ensure skew-symmetry: A = (params - params^T) / 2
        A = (self.skew_params - self.skew_params.T) / 2.0

        # Cayley transform approximation via Neumann series
        # Q ≈ Σ (−1)^k (I+A)^(−k) using series expansion
        I = torch.eye(self.size, device=A.device)

        # Iteratively apply: Q ≈ I + Σ (−A)^k for truncation_order terms
        Q = I.clone()
        term = I.clone()

        for k in range(1, self.truncation_order):
            term = term @ (-A) / k
            Q = Q + term

        return Q

Step 3: Training Loop with Spectrum Preservation

def train_poet_model(model, train_loader, num_epochs=3, learning_rate=1e-3):
    """
    Train LLM with orthogonal reparameterization.
    Monitors spectrum stability throughout training.
    """
    optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
    loss_fn = nn.CrossEntropyLoss()

    spectrum_history = []

    for epoch in range(num_epochs):
        epoch_loss = 0.0

        for batch_idx, (input_ids, labels) in enumerate(train_loader):
            optimizer.zero_grad()

            # Forward pass
            logits = model(input_ids)
            loss = loss_fn(logits.view(-1, logits.size(-1)), labels.view(-1))

            # Backward pass
            loss.backward()
            optimizer.step()

            epoch_loss += loss.item()

            # Periodically log spectral properties
            if batch_idx % 100 == 0:
                spectrum_stats = measure_spectrum(model)
                spectrum_history.append(spectrum_stats)

                print(f"Epoch {epoch}, Batch {batch_idx}: "
                      f"Loss={loss.item():.4f}, "
                      f"Singular Value Mean={spectrum_stats['mean']:.4f}")

    return spectrum_history

def measure_spectrum(model):
    """
    Compute spectral statistics for reparameterized weights.
    Verifies that singular values remain stable.
    """
    singular_values = []

    for name, module in model.named_modules():
        if isinstance(module, OrthogonalLinear):
            W = module.R_param @ module.W0 @ module.P_param
            _, S, _ = torch.linalg.svd(W)
            singular_values.extend(S.detach().cpu().tolist())

    singular_values = torch.tensor(singular_values)

    return {
        'mean': singular_values.mean().item(),
        'std': singular_values.std().item(),
        'max': singular_values.max().item(),
        'min': singular_values.min().item()
    }

Step 4: Merge-Then-Reinitialize Optimization

def merge_and_reinitialize(model, frequency=1000):
    """
    Periodically consolidate learned orthogonal transformations
    and reinitialize parameters to reduce GPU memory footprint.
    """
    for name, module in model.named_modules():
        if isinstance(module, OrthogonalLinear):
            # Consolidate: update W0 to current reconstruction
            with torch.no_grad():
                W_current = module.R_param @ module.W0 @ module.P_param
                module.W0.copy_(W_current)

                # Reset orthogonal factors to identity
                module.R_param.data.copy_(torch.eye(module.out_features))
                module.P_param.data.copy_(torch.eye(module.in_features))

Practical Guidance

Parameter Efficiency Trade-offs:

• Standard training: O(mn) trainable parameters per layer
• POET block-stochastic: O(m+n) parameters with comparable performance
• SPO with 3 primitives: 6 parameters per dimension vs m+n for full orthogonal

Generalization Improvements:

• POET achieves better validation perplexity than AdamW at same model size
• Spectrum preservation stabilizes training, reducing loss spikes
• Works synergistically with other techniques (layer norm, mixed precision)

Scaling Considerations:

• Efficient for models 1B-70B parameters
• Cayley-Neumann approximation reduces computational overhead by 40%
• Merge-then-reinitialize critical for training stability beyond 13B parameters

When to Apply POET:

• Long-context models (where spectrum stability prevents divergence)
• Multi-task training (shared orthogonal transformation improves transfer)
• Distillation targets (spectrum-preserving approach improves student learning)

Reference

• Orthogonal transformations preserve spectral properties: singular values invariant under orthogonal multiplication
• Hyperspherical energy: measure of weight magnitude variance; bounded under zero-mean isotropic Gaussian initialization
• Skew-symmetric matrices: parameterize Lie algebra so(n), directly generate orthogonal group SO(n)