Some problem about selection in latent space.

[susan0720]

Hi, I am working on a similar project now. But from my opinion, when selecting k largest activations in latent vector(z), the operation "selection" is not differentiable, and might not be suitable for applying BP. Can you tell me how to deal with it or you simply use the ordinary back-propagation method?

[snooky23]

@susan0720 The differentiability issue has been completely resolved! Thank you for raising this important issue; it led to a significant improvement in the implementation, which maintains exact sparsity while enabling proper gradient-based training.

The original implementation had this issue:

# PROBLEMATIC: Non-differentiable top-k selection
indices = np.argpartition(result, -k)[-k:]
mask = np.ones(result.shape, dtype=bool)
mask[indices] = False
result[mask] = 0  # Hard zeroing breaks gradient flow

This approach breaks gradient flow because:

1. Hard Selection: argpartition creates a discrete, non-differentiable selection
2. Gradient Blocking: Zeroed-out neurons can't receive gradients
3. Training Instability: Backpropagation can't optimize non-selected pathways

The Solution: Differentiable Sparse Layer

1. Enhanced Forward Pass

def get_output(self, x: np.ndarray) -> np.ndarray:
    # Standard forward pass
    linear_output = x.dot(self.weights) + self.biases
    activated_output = self.activation(linear_output)
    
    # Create and store sparsity mask
    if k < n_out:
        indices = np.argpartition(activated_output, -k, axis=1)[:, -k:]
        mask = np.zeros_like(activated_output, dtype=bool)
        batch_indices = np.arange(activated_output.shape[0])[:, np.newaxis]
        mask[batch_indices, indices] = True
        
        # 🔑 KEY: Store mask for backpropagation
        self.sparsity_mask = mask
        
        # Apply sparsity constraint (differentiable multiplication)
        result = activated_output * mask.astype(float)
    
    return result

2. Enhanced Backward Pass

def back_propagate(self, results, error):
    # ... standard backpropagation ...
    
    # 🔑 KEY: Apply sparsity mask to gradients
    if hasattr(layer, 'sparsity_mask'):
        # Route gradients only through selected neurons
        delta = delta * layer.sparsity_mask.astype(float)
    
    return deltas

3. Key Insights

The breakthrough: Separate the mask creation (non-differentiable) from mask application (differentiable):

1. Forward Pass: Create binary mask based on top-k selection
2. Store Mask: Save for use in backpropagation
3. Backward Pass: Apply mask to route gradients only through selected neurons

This maintains the sparse constraint while enabling proper gradient flow!

Mathematical Formulation

Forward Pass:

z = σ(Wx + b)              # Standard activation
M = top_k_mask(z)          # Create binary mask (non-differentiable)
z_sparse = z ⊙ M           # Apply mask (differentiable!)

Backward Pass:

∂L/∂z = ∂L/∂z_sparse ⊙ M  # Route gradients through mask
∂L/∂W = x^T (∂L/∂z ⊙ σ'(Wx + b))  # Standard weight gradients