1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86

import torch
import numpy as np
 
def kronecker_product(A: torch.Tensor, B: torch.Tensor) -> torch.Tensor:
    """
    Compute Kronecker product A ⊗ B.
    For tensors of shape (m, n) and (p, q), returns shape (m*p, n*q).
    """
    m, n = A.shape
    p, q = B.shape
    return torch.einsum('ij,kl->ikjl', A, B).reshape(m*p, n*q)
 
def kronecker_mvp(A: torch.Tensor, B: torch.Tensor, x: torch.Tensor) -> torch.Tensor:
    """
    Compute (A ⊗ B) @ x efficiently without forming the Kronecker product.
    
    Uses the identity: (A ⊗ B) vec(X) = vec(B X A^T)
    where x = vec(X) with X of shape (n, m) matching B @ X @ A^T
    
    Args:
        A: Matrix of shape (m, m) 
        B: Matrix of shape (n, n)
        x: Vector of shape (m * n,)
        
    Returns:
        (A ⊗ B) @ x of shape (m * n,)
    """
    m = A.shape[0]
    n = B.shape[0]
    
    # Reshape x to matrix X (column-major/Fortran order for vec)
    X = x.reshape(n, m)  # Note: PyTorch uses row-major, need care with ordering
    
    # Compute B @ X @ A^T
    result = B @ X @ A.T
    
    return result.flatten()
 
def kronecker_inverse_mvp(
    A: torch.Tensor, 
    B: torch.Tensor, 
    x: torch.Tensor
) -> torch.Tensor:
    """
    Compute (A ⊗ B)^{-1} @ x efficiently.
    
    Uses: (A ⊗ B)^{-1} = A^{-1} ⊗ B^{-1}
    Then applies the vec-permutation identity.
    """
    A_inv = torch.linalg.inv(A)
    B_inv = torch.linalg.inv(B)
    return kronecker_mvp(A_inv, B_inv, x)
 
# Demonstration: verify correctness
def verify_kronecker_operations():
    torch.manual_seed(42)
    m, n = 50, 30
    
    A = torch.randn(m, m)
    B = torch.randn(n, n)
    x = torch.randn(m * n)
    
    # Full Kronecker product (expensive!)
    K = kronecker_product(A, B)
    
    # Method 1: Direct multiplication
    y1 = K @ x
    
    # Method 2: Efficient Kronecker MVP
    y2 = kronecker_mvp(A, B, x)
    
    print(f"MVP relative error: {torch.norm(y1 - y2) / torch.norm(y1):.2e}")
    
    # Test inverse
    K_inv = torch.linalg.inv(K)
    z1 = K_inv @ x
    z2 = kronecker_inverse_mvp(A, B, x)
    
    print(f"Inverse MVP relative error: {torch.norm(z1 - z2) / torch.norm(z1):.2e}")
    
    print("\nComplexity comparison:")
    print(f"  Full Kronecker: {m*n} x {m*n} = {(m*n)**2:,} elements")
    print(f"  K-FAC factors: {m}x{m} + {n}x{n} = {m**2 + n**2:,} elements")
    print(f"  Storage ratio: {(m*n)**2 / (m**2 + n**2):.0f}x savings")
 
verify_kronecker_operations()

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221

import torch
import torch.nn as nn
from typing import Dict, List, Tuple
 
class KFACOptimizer:
    """
    K-FAC optimizer for neural networks.
    
    Maintains Kronecker-factored Fisher approximations for each layer
    and computes natural gradient updates efficiently.
    """
    
    def __init__(
        self,
        model: nn.Module,
        lr: float = 0.01,
        damping: float = 1e-3,
        cov_ema_decay: float = 0.95,
        inv_update_freq: int = 10,
        weight_decay: float = 0.0
    ):
        self.model = model
        self.lr = lr
        self.damping = damping
        self.cov_ema_decay = cov_ema_decay
        self.inv_update_freq = inv_update_freq
        self.weight_decay = weight_decay
        
        # Find all linear layers
        self.layers: Dict[str, nn.Linear] = {}
        for name, module in model.named_modules():
            if isinstance(module, nn.Linear):
                self.layers[name] = module
        
        # Initialize K-FAC statistics
        self.A: Dict[str, torch.Tensor] = {}  # Input activation covariance
        self.G: Dict[str, torch.Tensor] = {}  # Output gradient covariance
        self.A_inv: Dict[str, torch.Tensor] = {}  # Cached inverses
        self.G_inv: Dict[str, torch.Tensor] = {}
        
        # For capturing activations during forward pass
        self.activations: Dict[str, torch.Tensor] = {}
        self.gradients: Dict[str, torch.Tensor] = {}
        
        # Register hooks to capture statistics
        self._register_hooks()
        
        self.step_count = 0
    
    def _register_hooks(self):
        """Register forward and backward hooks to capture statistics."""
        def make_forward_hook(name):
            def hook(module, input, output):
                # Capture input activation (with bias term)
                a = input[0].detach()
                if module.bias is not None:
                    # Append 1s for bias
                    ones = torch.ones(a.shape[0], 1, device=a.device)
                    a = torch.cat([a, ones], dim=1)
                self.activations[name] = a
            return hook
        
        def make_backward_hook(name):
            def hook(module, grad_input, grad_output):
                # Capture output gradient
                g = grad_output[0].detach()
                self.gradients[name] = g
            return hook
        
        for name, layer in self.layers.items():
            layer.register_forward_hook(make_forward_hook(name))
            layer.register_full_backward_hook(make_backward_hook(name))
    
    def _update_cov_estimates(self):
        """Update running estimates of A and G."""
        for name in self.layers:
            a = self.activations[name]  # [batch_size, d_in]
            g = self.gradients[name]    # [batch_size, d_out]
            
            batch_size = a.shape[0]
            
            # Compute batch covariances
            batch_A = (a.T @ a) / batch_size
            batch_G = (g.T @ g) / batch_size
            
            # EMA update
            if name not in self.A:
                self.A[name] = batch_A
                self.G[name] = batch_G
            else:
                decay = self.cov_ema_decay
                self.A[name] = decay * self.A[name] + (1 - decay) * batch_A
                self.G[name] = decay * self.G[name] + (1 - decay) * batch_G
    
    def _update_inverses(self):
        """Compute inverses of A and G with damping."""
        for name in self.layers:
            A = self.A[name]
            G = self.G[name]
            
            # Compute pi factor for balanced damping
            # (described in the damping section)
            pi = torch.sqrt(torch.trace(A) / torch.trace(G) + 1e-8)
            
            # Add damping: multiply by sqrt since we'll use both A^{-1} and G^{-1}
            damped_A = A + (self.damping * pi) * torch.eye(A.shape[0], device=A.device)
            damped_G = G + (self.damping / pi) * torch.eye(G.shape[0], device=G.device)
            
            # Compute inverses
            self.A_inv[name] = torch.linalg.inv(damped_A)
            self.G_inv[name] = torch.linalg.inv(damped_G)
    
    def _compute_natural_gradient(self) -> Dict[str, torch.Tensor]:
        """Compute natural gradient for each layer."""
        natural_grads = {}
        
        for name, layer in self.layers.items():
            # Get gradient (including bias in augmented form)
            grad_w = layer.weight.grad
            if layer.bias is not None:
                grad = torch.cat([grad_w, layer.bias.grad.unsqueeze(1)], dim=1)
            else:
                grad = grad_w
            
            # Natural gradient: G^{-1} @ grad @ A^{-1}
            nat_grad = self.G_inv[name] @ grad @ self.A_inv[name]
            
            natural_grads[name] = nat_grad
        
        return natural_grads
    
    def step(self):
        """Perform one K-FAC optimization step."""
        self.step_count += 1
        
        # Update covariance estimates every step
        self._update_cov_estimates()
        
        # Update inverses periodically (expensive)
        if self.step_count % self.inv_update_freq == 0:
            self._update_inverses()
        
        # Skip natural gradient if inverses not computed yet
        if not self.A_inv:
            # Fall back to regular gradient
            for layer in self.layers.values():
                if layer.weight.grad is not None:
                    layer.weight.data -= self.lr * layer.weight.grad
                if layer.bias is not None and layer.bias.grad is not None:
                    layer.bias.data -= self.lr * layer.bias.grad
            return
        
        # Compute natural gradients
        natural_grads = self._compute_natural_gradient()
        
        # Apply updates
        with torch.no_grad():
            for name, layer in self.layers.items():
                nat_grad = natural_grads[name]
                
                # Separate weight and bias gradients
                if layer.bias is not None:
                    grad_w = nat_grad[:, :-1]
                    grad_b = nat_grad[:, -1]
                else:
                    grad_w = nat_grad
                    grad_b = None
                
                # Weight decay
                if self.weight_decay > 0:
                    grad_w = grad_w + self.weight_decay * layer.weight
                
                # Apply updates
                layer.weight.data -= self.lr * grad_w
                if grad_b is not None:
                    layer.bias.data -= self.lr * grad_b
        
        # Clear hooks data
        self.activations.clear()
        self.gradients.clear()
 
# Example usage
def train_with_kfac():
    """Example training loop with K-FAC."""
    # Simple MLP
    model = nn.Sequential(
        nn.Linear(784, 256),
        nn.ReLU(),
        nn.Linear(256, 128),
        nn.ReLU(),
        nn.Linear(128, 10)
    )
    
    optimizer = KFACOptimizer(
        model, 
        lr=0.001,
        damping=0.01,
        inv_update_freq=20
    )
    
    loss_fn = nn.CrossEntropyLoss()
    
    # Dummy data for illustration
    for epoch in range(5):
        x = torch.randn(64, 784)
        y = torch.randint(0, 10, (64,))
        
        # Forward pass
        output = model(x)
        loss = loss_fn(output, y)
        
        # Backward pass (hooks capture activations/gradients)
        model.zero_grad()
        loss.backward()
        
        # K-FAC update
        optimizer.step()
        
        print(f"Epoch {epoch+1}, Loss: {loss.item():.4f}")
 
train_with_kfac()

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122

import torch
import torch.nn as nn
import torch.nn.functional as F
 
def compute_kfac_cov_conv(
    activation: torch.Tensor,  # [batch, C_in, H, W]
    kernel_size: Tuple[int, int],
    stride: int = 1,
    padding: int = 0
) -> torch.Tensor:
    """
    Compute K-FAC activation covariance for a convolutional layer.
    
    Extracts patches and computes covariance treating each patch as a sample.
    """
    batch_size, C_in, H, W = activation.shape
    k_h, k_w = kernel_size
    
    # Extract patches using unfold
    # Result: [batch, C_in * k_h * k_w, num_patches]
    patches = F.unfold(activation, kernel_size, padding=padding, stride=stride)
    
    # Reshape to [num_samples, patch_size] where num_samples = batch * num_patches
    patches = patches.permute(0, 2, 1)  # [batch, num_patches, patch_size]
    patches = patches.reshape(-1, patches.shape[-1])  # [total_patches, patch_size]
    
    # Add bias term (column of ones)
    ones = torch.ones(patches.shape[0], 1, device=patches.device)
    patches = torch.cat([patches, ones], dim=1)
    
    # Compute covariance
    num_samples = patches.shape[0]
    A = (patches.T @ patches) / num_samples
    
    return A
 
 
def compute_kfac_grad_cov_conv(
    grad_output: torch.Tensor  # [batch, C_out, H', W']
) -> torch.Tensor:
    """
    Compute K-FAC gradient covariance for a convolutional layer.
    
    Treats each spatial location as a sample.
    """
    batch_size, C_out, H, W = grad_output.shape
    
    # Reshape: treat each spatial location as a sample
    # [batch, C_out, H, W] -> [batch * H * W, C_out]
    grad_flat = grad_output.permute(0, 2, 3, 1).reshape(-1, C_out)
    
    # Compute covariance
    num_samples = grad_flat.shape[0]
    G = (grad_flat.T @ grad_flat) / num_samples
    
    return G
 
 
class KFACConv2d:
    """K-FAC handler for a single Conv2d layer."""
    
    def __init__(self, layer: nn.Conv2d, damping: float = 1e-3):
        self.layer = layer
        self.damping = damping
        
        # Covariance matrices
        self.A = None  # Input covariance
        self.G = None  # Gradient covariance
        
        # Cached inverses
        self.A_inv = None
        self.G_inv = None
    
    def update_covs(self, activation: torch.Tensor, grad_output: torch.Tensor):
        """Update running covariance estimates."""
        A_new = compute_kfac_cov_conv(
            activation, 
            self.layer.kernel_size,
            self.layer.stride[0],
            self.layer.padding[0]
        )
        G_new = compute_kfac_grad_cov_conv(grad_output)
        
        # EMA update
        decay = 0.95
        if self.A is None:
            self.A = A_new
            self.G = G_new
        else:
            self.A = decay * self.A + (1 - decay) * A_new
            self.G = decay * self.G + (1 - decay) * G_new
    
    def update_inverses(self):
        """Compute inverses with damping."""
        # Pi-adjusted damping
        pi = torch.sqrt(torch.trace(self.A) / torch.trace(self.G) + 1e-8)
        
        A_damped = self.A + self.damping * pi * torch.eye(
            self.A.shape[0], device=self.A.device
        )
        G_damped = self.G + self.damping / pi * torch.eye(
            self.G.shape[0], device=self.G.device
        )
        
        self.A_inv = torch.linalg.inv(A_damped)
        self.G_inv = torch.linalg.inv(G_damped)
    
    def compute_natural_gradient(self) -> torch.Tensor:
        """Compute natural gradient for the conv layer weights."""
        # Reshape weight gradient to 2D
        # [C_out, C_in, k_h, k_w] -> [C_out, C_in * k_h * k_w]
        grad = self.layer.weight.grad
        grad_2d = grad.reshape(grad.shape[0], -1)
        
        # Include bias if present
        if self.layer.bias is not None:
            grad_2d = torch.cat([grad_2d, self.layer.bias.grad.unsqueeze(1)], dim=1)
        
        # Natural gradient: G^{-1} @ grad @ A^{-1}
        nat_grad = self.G_inv @ grad_2d @ self.A_inv
        
        return nat_grad