docs/ALEXNET_BRIDGE.md

AlexNet-to-BURT Bridge: Vision Backbone Integration

Project: BURT-IMMA
Contact: jessica@collectivekitty.com
License: BSL-1.1

---

1. Overview

This document specifies the integration bridge between AlexNet's convolutional feature extraction and BURT's Boolean reasoning layers. The bridge replaces AlexNet's fully-connected classification head with MMEP-trained Boolean reasoning, creating a hybrid architecture that leverages proven vision features while applying principled constraint-based reasoning for classification.

---

2. Architecture Mapping

2.1 Original AlexNet (Krizhevsky 2012)

Input (224x224x3)
  -> Conv1 (96 filters, 11x11, stride 4) + ReLU + MaxPool + LRN
  -> Conv2 (256 filters, 5x5, pad 2) + ReLU + MaxPool + LRN
  -> Conv3 (384 filters, 3x3, pad 1) + ReLU
  -> Conv4 (384 filters, 3x3, pad 1) + ReLU
  -> Conv5 (256 filters, 3x3, pad 1) + ReLU + MaxPool
  -> Flatten (9216)
  -> FC6 (4096) + ReLU + Dropout
  -> FC7 (4096) + ReLU + Dropout
  -> FC8 (1000) + Softmax

2.2 BURT-AlexNet Bridge

Input (224x224x3)
  -> [AlexNet Conv Backbone: Conv1-Conv5] (feature extraction, retained)
  -> Flatten (9216)
  -> Linear Projection (9216 -> d_model=512)
  -> [BURT Reasoning Block x 2] (replaces FC6, FC7, FC8)
      -> GatesNorm + SmoothLeaky
      -> CIFGMatrixMemory (retrieval)
      -> GatesRouter -> Expert MoE
      -> BooleanConstraintLayer
      -> QuantumInterferenceResolver
  -> Constrained Classification Head (d_model -> num_classes)
  -> Constrained Softmax (entropy <= 0.20)

2.3 Layer Correspondence

AlexNet Layer	BURT-AlexNet	Change
Conv1	Conv1	Retained (frozen or fine-tuned)
Conv2	Conv2	Retained (frozen or fine-tuned)
Conv3	Conv3	Retained (frozen or fine-tuned)
Conv4	Conv4	Retained (frozen or fine-tuned)
Conv5	Conv5	Retained (frozen or fine-tuned)
FC6 (4096)	BURT Block 1	Replaced with MMEP-trained block
FC7 (4096)	BURT Block 2	Replaced with MMEP-trained block
FC8 (1000)	Constrained Head	Replaced with entropy-bounded head
ReLU	SmoothLeaky	Smooth activation for EP compatibility
Dropout	GatesNorm + Sparsity	Structural regularization
LRN	Spectral Normalization	Modern normalization
Softmax	Constrained Softmax	Entropy bounded at 0.20

---

3. Feature Extraction (Conv Backbone)

3.1 Retained Architecture

The convolutional layers (Conv1-Conv5) are retained from AlexNet as the feature extraction backbone. These layers have been proven effective for learning hierarchical visual features:

• Conv1: Edge detectors, Gabor-like filters
• Conv2: Texture and color detectors
• Conv3: Part detectors
• Conv4: Object part combinations
• Conv5: Whole-object representations

3.2 Modifications to Conv Backbone

Modification	Original	Bridge	Rationale
Activation	ReLU	SmoothLeaky	Required for EP energy landscape smoothness
Normalization	LRN	Spectral Norm	Bounded singular values for convergence
Pooling	MaxPool	MaxPool	Retained (compatible with EP)
Initialization	Random	ImageNet pretrained	Transfer learning

3.3 Fine-tuning Strategy

Phase 1 (memorization): Conv backbone frozen, only projection trained
Phase 2 (query):        Conv backbone frozen, BURT blocks trained
Phase 3 (convergence):  Conv backbone unfrozen with 10x lower LR
Phase 4 (ablation):     All layers trainable

Learning rates:

• Conv layers: lr_base * 0.1 (slow adaptation)
• Projection: lr_base (full learning rate)
• BURT blocks: lr_base (full learning rate)
• Classification head: lr_base * 2.0 (faster convergence)

---

4. Replacing FC Layers with MMEP-Trained Boolean Reasoning

4.1 Projection Layer

The flattened convolutional output (9216 dims) is projected to the BURT model dimension:

class ConvToBURT(nn.Module):
    def __init__(self, conv_dim=9216, d_model=512):
        self.projection = nn.Linear(conv_dim, d_model)
        self.norm = GatesNormalization(d_model)
        self.activation = SmoothLeaky(k=1.0)
    
    def forward(self, conv_features):
        x = self.projection(conv_features.flatten(1))
        x = self.norm(x)
        x = self.activation(x)
        return x  # [batch, d_model]

4.2 BURT Reasoning Blocks

Two BURT blocks replace FC6 and FC7:

Block 1 (replaces FC6): Feature integration + memory retrieval

• Retrieves relevant stored patterns from C_global
• Routes to domain-specific experts (animal vs. vehicle vs. scene, etc.)
• Applies Boolean constraints to ensure logical consistency

Block 2 (replaces FC7): Abstract reasoning + classification prep

• Combines expert outputs via interference
• Enforces Huntington postulates on classification features
• Prepares logits for constrained output

4.3 Constrained Classification Head

class ConstrainedClassificationHead(nn.Module):
    def __init__(self, d_model=512, num_classes=1000, entropy_bound=0.20):
        self.linear = SpectralNorm(nn.Linear(d_model, num_classes))
        self.entropy_bound = entropy_bound
    
    def forward(self, x):
        logits = self.linear(x)
        probs = constrained_softmax(logits, entropy_bound=self.entropy_bound)
        return logits, probs

The entropy bound of 0.20 on the classification distribution means the model must be confident in its predictions. A uniform distribution over 1000 classes would have entropy ~6.9 nats; bounding at 0.20 forces the model to concentrate probability mass on very few classes.

---

5. Data Augmentation: BiEncoder Retrieval Replaces Color Jitter

5.1 Original AlexNet Augmentation

Krizhevsky 2012 used:

• Random 224x224 crops from 256x256 images
• Horizontal flips
• Color jitter via PCA (fancy PCA augmentation on RGB channels)

5.2 BURT-AlexNet Augmentation

The color jitter augmentation is replaced by BiEncoder retrieval augmentation:

class BiEncoderAugmentation:
    def __init__(self, memory_bank, k=5):
        self.memory_bank = memory_bank  # Pre-computed feature bank
        self.k = k  # Number of retrieved neighbors
    
    def augment(self, x, features):
        # Retrieve k nearest neighbors from memory bank
        distances, indices = self.memory_bank.search(features, k=self.k)
        neighbors = self.memory_bank.get_images(indices)
        
        # Create augmented batch: original + retrieved context
        augmented = {
            'image': x,
            'context_features': self.memory_bank.get_features(indices),
            'context_labels': self.memory_bank.get_labels(indices),
            'distances': distances
        }
        return augmented

Rationale: Rather than perturbing color channels (which adds noise), BiEncoder retrieval provides semantically relevant context from the training set. This context is stored in the matrix memory and used during inference, enabling the model to reason about an image by comparing it to similar stored examples.

5.3 Retained Augmentations

• Random 224x224 crops (retained)
• Horizontal flips (retained)
• Random erasing (added for robustness)

---

6. Training Schedule (Adapted from Krizhevsky 2012)

6.1 Original Schedule

Krizhevsky trained for 90 epochs with:

• SGD with momentum 0.9
• Weight decay 0.0005
• LR starts at 0.01, divided by 10 when validation error plateaus
• Batch size 128 across 2 GPUs

6.2 BURT-AlexNet Schedule

Adapted for single-GPU MMEP training:

Phase	Epochs	LR (conv)	LR (BURT)	Method	Batch
1: Memorize	5	0 (frozen)	1e-3	Forward only	32
2: Route	20	0 (frozen)	3e-4	Backprop	16
3: Converge	50	1e-5	1e-4	MMEP (EP)	8
4: Fine-tune	15	1e-5	5e-5	MMEP (EP)	8

Total: 90 epochs (matching Krizhevsky's total training duration)

6.3 LR Decay

# Cosine annealing with warm restarts
scheduler = CosineAnnealingWarmRestarts(
    optimizer,
    T_0=10,       # First restart period
    T_mult=2,     # Double period after each restart
    eta_min=1e-6  # Minimum LR
)

6.4 Weight Decay

• Conv layers: 0.0005 (matching Krizhevsky)
• BURT layers: 0.0 (spectral norm + sparsity replaces weight decay)
• Classification head: 0.0001

---

7. Entropy Bound on Classification Head

7.1 Constraint Definition

The output distribution after softmax must satisfy:

H(p) = -sum(p_i * log(p_i)) <= 0.20

For ImageNet (1000 classes), this means:

• Maximum entropy of uniform distribution: log(1000) = 6.91 nats
• Bound of 0.20 nats implies effective support of ~1.22 classes
• The model must be very confident, concentrating > 80% probability on the top class

7.2 Enforcement

During training:

def constrained_classification_loss(logits, targets, entropy_bound=0.20):
    probs = softmax(logits)
    ce_loss = cross_entropy(logits, targets)
    
    # Entropy penalty
    entropy = -sum(probs * log(probs + 1e-10), dim=-1)
    entropy_penalty = relu(entropy - entropy_bound).mean()
    
    return ce_loss + 10.0 * entropy_penalty

During inference:

def constrained_predict(logits, entropy_bound=0.20):
    # Temperature scaling to meet entropy bound
    temperature = find_temperature(logits, entropy_bound)
    probs = softmax(logits / temperature)
    return probs

7.3 Implications

• Forces the model to learn discriminative features
• Prevents hedging across similar classes
• Ensures interpretable predictions (few classes with high probability)
• Acts as implicit confidence calibration

---

8. Spectral Normalization on All Weight Matrices

8.1 Application

Every weight matrix in the BURT-AlexNet bridge is spectrally normalized:

def apply_spectral_normalization(model, lambda_max=0.95):
    for name, module in model.named_modules():
        if isinstance(module, (nn.Linear, nn.Conv2d)):
            # Constrain largest singular value
            spectral_norm(module, name='weight', n_power_iterations=1)
            
            # Additional: clip to lambda_max after normalization
            with torch.no_grad():
                U, S, V = torch.svd(module.weight_orig)
                S_clipped = torch.clamp(S, max=lambda_max)
                module.weight_orig.copy_(U @ torch.diag(S_clipped) @ V.T)

8.2 Where Applied

Layer Type	Spectral Bound	Update Frequency
Conv1-Conv5	0.95	Every step (if unfrozen)
Projection	0.95	Every step
BURT attention (Q, K, V)	0.95	Every step
Expert weights	0.95	Every step
Router weights	0.95	Every step
Classification head	0.95	Every step
Memory gate weights	0.95	Every step

8.3 Interaction with EP Convergence

Spectral normalization is essential for EP convergence in the BURT blocks:

• Guarantees contraction mapping in free phase
• Ensures unique equilibrium for each input
• Prevents activation explosion during nudged phase
• See MMEP_THEORY.md Section 7 for the Lyapunov proof

---

9. Complete Forward Pass

class BURTAlexNet(nn.Module):
    def forward(self, x, memory_state=None):
        # === Conv Backbone (from AlexNet) ===
        x = self.conv1(x)           # [B, 96, 55, 55]
        x = self.smooth_leaky(x)
        x = self.pool1(x)           # [B, 96, 27, 27]
        
        x = self.conv2(x)           # [B, 256, 27, 27]
        x = self.smooth_leaky(x)
        x = self.pool2(x)           # [B, 256, 13, 13]
        
        x = self.conv3(x)           # [B, 384, 13, 13]
        x = self.smooth_leaky(x)
        
        x = self.conv4(x)           # [B, 384, 13, 13]
        x = self.smooth_leaky(x)
        
        x = self.conv5(x)           # [B, 256, 13, 13]
        x = self.smooth_leaky(x)
        x = self.pool5(x)           # [B, 256, 6, 6]
        
        # === Bridge: Conv -> BURT ===
        x = x.flatten(1)            # [B, 9216]
        x = self.projection(x)     # [B, d_model=512]
        x = self.gates_norm_proj(x)
        x = self.smooth_leaky(x)
        
        # === BURT Block 1 (replaces FC6) ===
        # Memory retrieval
        if memory_state is None:
            memory_state = self.C_global
        mem_context = memory_state @ x.unsqueeze(-1)  # [B, d_model, 1]
        x = x + mem_context.squeeze(-1)
        
        # Expert routing
        expert_ids, weights = self.router1(x)
        expert_outs = [self.experts1[k](x) for k in expert_ids]
        x = x + self.resolver1(expert_outs, weights)
        
        # Boolean constraint
        x = x + self.boolean1(self.gates_norm1(x))
        
        # Memory update
        self.update_memory(x)
        
        # === BURT Block 2 (replaces FC7) ===
        mem_context = memory_state @ x.unsqueeze(-1)
        x = x + mem_context.squeeze(-1)
        
        expert_ids, weights = self.router2(x)
        expert_outs = [self.experts2[k](x) for k in expert_ids]
        x = x + self.resolver2(expert_outs, weights)
        
        x = x + self.boolean2(self.gates_norm2(x))
        
        # === Classification Head (replaces FC8) ===
        logits = self.classification_head(x)  # [B, num_classes]
        probs = constrained_softmax(logits, entropy_bound=0.20)
        
        return logits, probs, memory_state

---

10. Expected Performance

10.1 ImageNet Comparison

Model	Top-1 Acc	Top-5 Acc	Parameters
AlexNet (2012)	63.3%	84.6%	61M
AlexNet + BURT (projected)	68-72%	87-90%	45M

The BURT bridge reduces parameters (no 4096-dim FC layers) while improving accuracy through structured reasoning and memory retrieval.

10.2 Reasoning Tasks

On visual reasoning benchmarks (CLEVR, visual arithmetic):

• Expected improvement over baseline AlexNet: 15-25%
• Boolean constraint enforcement ensures logically consistent answers
• Memory retrieval enables few-shot pattern matching

---

11. References

• Krizhevsky, A., Sutskever, I., & Hinton, G. E. (2012). "ImageNet Classification with Deep Convolutional Neural Networks." NeurIPS.
• MMEP_THEORY.md - Training algorithm
• BURT_SPEC.md - Full BURT architecture specification
• kernel_api.md - CUDA implementations for constrained_softmax