# Poet Orthogonal LLM Training > Improve LLM training stability and generalization by reparameterizing weight matrices as orthogonal transformations, achieving better perplexity than AdamW with fewer trainable parameters. - Skill: `adu2021/poet-orthogonal-llm-training` (Agent Skill) - Install (CLI): `npx skillmds@latest add adu2021/poet-orthogonal-llm-training` - Raw SKILL.md: https://api.skillmd.com/api/skills/adu2021/poet-orthogonal-llm-training/raw - Safety review: pending - Works with: Claude Code, Claude.ai, OpenAI Codex - Category: AI & ML - License: MIT - Author: adu2021 (https://skillmd.com/u/adu2021) - Updated: 2026-09-17 - Page: https://skillmd.com/skills/adu2021/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 ```python 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 ```python 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 ```python 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 ```python 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)