File size: 5,434 Bytes
b6564d4 | 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 | """
Weight Retrieval Layer - Core Innovation of WRAG 2.0
This layer dynamically retrieves and composes weights during the forward pass,
rather than using fixed weights like traditional neural networks.
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
class WeightRetrievalLayer(nn.Module):
"""
Dynamically retrieves and composes layer weights based on input.
Traditional layer: x → W → output (W is fixed)
This layer: x → retrieve W based on x → W → output (W is dynamic)
Args:
d_in: Input dimension
d_out: Output dimension
num_shards: Number of weight shards in the bank
k: Number of shards to retrieve and compose
"""
def __init__(self, d_in, d_out, num_shards=10, k=3):
super().__init__()
self.d_in = d_in
self.d_out = d_out
self.num_shards = num_shards
self.k = k
# Weight bank: stores multiple weight matrices (shards)
# Shape: [num_shards, d_in, d_out]
self.weight_bank = nn.Parameter(
torch.randn(num_shards, d_in, d_out) * 0.02
)
# Query network: converts input to retrieval query
# This learns what kind of weights to retrieve for each input
self.query_net = nn.Sequential(
nn.Linear(d_in, d_in // 2),
nn.ReLU(),
nn.Linear(d_in // 2, num_shards)
)
# Bias (optional, like standard linear layer)
self.bias = nn.Parameter(torch.zeros(d_out))
def forward(self, x, return_weights=False):
"""
Forward pass with dynamic weight retrieval.
Args:
x: Input tensor [batch, d_in]
return_weights: If True, return dynamic weight matrix instead of applying it
Returns:
If return_weights=False: output [batch, d_out], scores [batch, num_shards]
If return_weights=True: dynamic_W [batch, d_in, d_out], scores [batch, num_shards]
"""
batch = x.shape[0]
# Step 1: Compute retrieval scores
retrieval_scores = self.query_net(x) # [batch, num_shards]
# Step 2: Select top-k shards
top_k_scores, top_k_indices = torch.topk(
retrieval_scores, k=self.k, dim=1
) # [batch, k]
# Step 3: Softmax to get composition weights
composition_weights = F.softmax(top_k_scores, dim=1) # [batch, k]
# Step 4: Retrieve selected shards efficiently
# Instead of indexing which causes memory explosion, use einsum
# Create one-hot encoding for top-k selection
selection_mask = torch.zeros(batch, self.num_shards, device=x.device)
selection_mask.scatter_(1, top_k_indices, composition_weights)
# Step 5: Compose dynamic weight matrix using einsum
# weight_bank: [num_shards, d_in, d_out]
# selection_mask: [batch, num_shards]
dynamic_W = torch.einsum(
'bs,sio->bio',
selection_mask,
self.weight_bank
) # [batch, d_in, d_out]
if return_weights:
return dynamic_W, retrieval_scores
# Step 6: Apply dynamic weights to input
output = torch.bmm(x.unsqueeze(1), dynamic_W).squeeze(1) # [batch, d_out]
# Add bias
output = output + self.bias
return output, retrieval_scores
def get_retrieval_stats(self, x):
"""
Get statistics about which shards are being retrieved.
Useful for debugging and analysis.
"""
if len(x.shape) == 3:
query_input = x.mean(dim=1)
else:
query_input = x
retrieval_scores = self.query_net(query_input)
top_k_scores, top_k_indices = torch.topk(
retrieval_scores, k=self.k, dim=1
)
return {
'top_k_indices': top_k_indices.cpu().numpy(),
'top_k_scores': top_k_scores.cpu().numpy(),
'all_scores': retrieval_scores.cpu().numpy()
}
class SimpleWeightRetrievalNet(nn.Module):
"""
Simple network using WeightRetrievalLayers for testing.
Architecture:
Input → WRLayer1 → ReLU → WRLayer2 → Output
"""
def __init__(self, input_dim, hidden_dim, output_dim, num_shards=10, k=3):
super().__init__()
self.layer1 = WeightRetrievalLayer(
input_dim, hidden_dim, num_shards, k
)
self.layer2 = WeightRetrievalLayer(
hidden_dim, output_dim, num_shards, k
)
def forward(self, x):
x = self.layer1(x)
x = F.relu(x)
x = self.layer2(x)
return x
# Quick test
if __name__ == "__main__":
print("Testing WeightRetrievalLayer...")
# Create layer
layer = WeightRetrievalLayer(
d_in=784, # MNIST flattened
d_out=128,
num_shards=10,
k=3
)
# Test forward pass
batch_size = 4
x = torch.randn(batch_size, 784)
output = layer(x)
print(f"Input shape: {x.shape}")
print(f"Output shape: {output.shape}")
# Check retrieval stats
stats = layer.get_retrieval_stats(x)
print(f"Retrieved shards: {stats['top_k_indices']}")
print("\n✓ WeightRetrievalLayer works!")
|