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

import numpy as np
import matplotlib.pyplot as plt
 
def gradient_survival(decay_factor, max_T=500):
    """Calculate gradient magnitude through T steps."""
    return decay_factor ** np.arange(max_T)
 
# Typical scenarios
T = np.arange(500)
 
# Vanilla RNN: gradient factor around 0.9 (often worse)
vanilla_rnn_09 = gradient_survival(0.90)
vanilla_rnn_08 = gradient_survival(0.80)
 
# LSTM with various forget gate means
lstm_f099 = gradient_survival(0.99)
lstm_f095 = gradient_survival(0.95)
lstm_f090 = gradient_survival(0.90)
 
plt.figure(figsize=(12, 5))
 
plt.subplot(1, 2, 1)
plt.semilogy(T, vanilla_rnn_09, label='Vanilla RNN (factor=0.9)', color='red')
plt.semilogy(T, vanilla_rnn_08, label='Vanilla RNN (factor=0.8)', color='darkred')
plt.semilogy(T, lstm_f099, label='LSTM (f=0.99)', color='green')
plt.semilogy(T, lstm_f095, label='LSTM (f=0.95)', color='blue')
plt.xlabel('Time Steps Back')
plt.ylabel('Relative Gradient Magnitude (log scale)')
plt.title('Gradient Survival: Vanilla RNN vs LSTM')
plt.legend()
plt.grid(True, alpha=0.3)
 
# Zoomed view on practical learning range
plt.subplot(1, 2, 2)
practical_T = 100
plt.bar(['Vanilla\nRNN (0.9)', 'LSTM\n(f=0.95)', 'LSTM\n(f=0.99)', 'LSTM\n(f=1.0)'],
        [0.9**practical_T, 0.95**practical_T, 0.99**practical_T, 1.0**practical_T],
        color=['red', 'blue', 'green', 'darkgreen'])
plt.ylabel(f'Gradient Surviving {practical_T} Steps')
plt.title(f'Gradient at T={practical_T}')
plt.yscale('log')
 
plt.tight_layout()
plt.savefig('gradient_comparison.png', dpi=150)
plt.show()
 
print(f"Gradient survival at T=100:")
print(f"  Vanilla RNN (0.9): {0.9**100:.2e}")
print(f"  LSTM (f=0.95):     {0.95**100:.2e}")
print(f"  LSTM (f=0.99):     {0.99**100:.4f}")
print(f"  LSTM (f=1.00):     {1.00**100:.4f}")

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23

# Forward equations (for reference):
c_t = f_t ⊙ c_{t-1} + i_t ⊙ c̃_t
h_t = o_t ⊙ tanh(c_t)
 
# Gradient of loss w.r.t. cell state (backward recursion):
# Two paths: through h_t and through c_{t+1}
 
∂L/∂c_t = [∂L/∂h_t ⊙ ∂h_t/∂c_t] + [∂L/∂c_{t+1} ⊙ ∂c_{t+1}/∂c_t]
 
Where:
  ∂h_t/∂c_t = o_t ⊙ (1 - tanh²(c_t))    # output gate × tanh derivative
  ∂c_{t+1}/∂c_t = f_{t+1}                 # just the forget gate!
 
So:
∂L/∂c_t = ∂L/∂h_t ⊙ o_t ⊙ (1 - tanh²(c_t)) + ∂L/∂c_{t+1} ⊙ f_{t+1}
 
# The CEC gradient pathway:
# Following only the c → c path back T steps:
 
∂L/∂c_0 = ∂L/∂c_T ⊙ ∏(t=1 to T) f_t  +  [other paths through h]
 
# If all f_t ≈ 1: ∂L/∂c_0 ≈ ∂L/∂c_T — gradient preserved!
# This is the Constant Error Carousel in action.

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

import torch
import numpy as np
 
def estimate_memory_capacity(lstm_hidden_size, num_trials=1000):
    """
    Estimate how many independent facts an LSTM can reliably store.
    
    Method: Write random patterns, read after delay, measure correlation.
    """
    from torch import nn
    
    class MemoryProber(nn.Module):
        def __init__(self, hidden_size, pattern_size):
            super().__init__()
            self.encoder = nn.Linear(pattern_size, hidden_size * 4)
            self.decoder = nn.Linear(hidden_size, pattern_size)
            self.hidden_size = hidden_size
        
        def encode(self, pattern, cell_state):
            """Encode pattern into cell state."""
            gates = torch.sigmoid(self.encoder(pattern))
            i, f, o, g = gates.chunk(4, dim=-1)
            new_c = f * cell_state + i * torch.tanh(g)
            return new_c
        
        def decode(self, cell_state):
            """Decode pattern from cell state."""
            return self.decoder(cell_state)
    
    results = []
    for pattern_size in [8, 16, 32, 64, 128]:
        model = MemoryProber(lstm_hidden_size, pattern_size)
        
        correct_recalls = 0
        for _ in range(num_trials):
            # Initialize cell state
            c = torch.zeros(1, lstm_hidden_size)
            
            # Generate random pattern
            pattern = torch.randn(1, pattern_size)
            
            # Encode
            c = model.encode(pattern, c)
            
            # Simulate delay (multiple identity passes)
            for _ in range(50):
                c = c * 0.99  # Slight forget
            
            # Decode
            retrieved = model.decode(c)
            
            # Check correlation
            corr = torch.corrcoef(torch.stack([
                pattern.flatten(), retrieved.flatten()
            ]))[0, 1]
            
            if corr > 0.8:
                correct_recalls += 1
        
        results.append((pattern_size, correct_recalls / num_trials))
        print(f"Pattern size {pattern_size}: {correct_recalls/num_trials:.1%} correct")
    
    return results
 
# Run analysis
print("Memory capacity analysis for hidden_size=256:")
estimate_memory_capacity(256)

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

import torch
from collections import deque
 
class CellStateMonitor:
    """Monitor cell state statistics during training/inference."""
    
    def __init__(self, window_size=100):
        self.window_size = window_size
        self.mean_history = deque(maxlen=window_size)
        self.std_history = deque(maxlen=window_size)
        self.max_history = deque(maxlen=window_size)
        self.min_history = deque(maxlen=window_size)
    
    def update(self, cell_state: torch.Tensor):
        """Record statistics of current cell state."""
        with torch.no_grad():
            self.mean_history.append(cell_state.mean().item())
            self.std_history.append(cell_state.std().item())
            self.max_history.append(cell_state.max().item())
            self.min_history.append(cell_state.min().item())
    
    def get_summary(self) -> dict:
        """Get summary statistics over window."""
        import numpy as np
        return {
            'mean_of_means': np.mean(self.mean_history),
            'mean_of_stds': np.mean(self.std_history),
            'max_seen': max(self.max_history),
            'min_seen': min(self.min_history),
            'range': max(self.max_history) - min(self.min_history),
        }
    
    def check_health(self) -> tuple[bool, str]:
        """Check if cell state dynamics are healthy."""
        summary = self.get_summary()
        
        issues = []
        
        # Check for explosion
        if summary['max_seen'] > 10:
            issues.append(f"Cell state exploding (max={summary['max_seen']:.2f})")
        
        # Check for collapse
        if summary['mean_of_stds'] < 0.01:
            issues.append(f"Cell state collapsed (std={summary['mean_of_stds']:.4f})")
        
        # Check for drift
        if abs(summary['mean_of_means']) > 2:
            issues.append(f"Cell state drifting (mean={summary['mean_of_means']:.2f})")
        
        if issues:
            return False, "; ".join(issues)
        return True, "Cell state dynamics healthy"
    
    def plot_history(self, save_path='cell_state_history.png'):
        """Plot cell state statistics over time."""
        import matplotlib.pyplot as plt
        
        fig, axes = plt.subplots(2, 2, figsize=(12, 8))
        
        axes[0, 0].plot(list(self.mean_history))
        axes[0, 0].set_title('Cell State Mean')
        axes[0, 0].axhline(y=0, color='r', linestyle='--', alpha=0.5)
        
        axes[0, 1].plot(list(self.std_history))
        axes[0, 1].set_title('Cell State Std')
        
        axes[1, 0].plot(list(self.max_history), label='Max')
        axes[1, 0].plot(list(self.min_history), label='Min')
        axes[1, 0].set_title('Cell State Range')
        axes[1, 0].legend()
        
        axes[1, 1].fill_between(
            range(len(self.max_history)),
            list(self.min_history),
            list(self.max_history),
            alpha=0.3
        )
        axes[1, 1].plot(list(self.mean_history), color='red')
        axes[1, 1].set_title('Cell State Distribution Over Time')
        
        plt.tight_layout()
        plt.savefig(save_path, dpi=150)
        plt.show()

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

import torch
import torch.nn as nn
import torch.nn.functional as F
 
class OptimizedLSTMCell(nn.Module):
    """
    LSTM cell with best practices for gradient flow.
    """
    def __init__(
        self,
        input_size: int,
        hidden_size: int,
        forget_bias: float = 1.0,
        cell_clip: float = None,
        use_layer_norm: bool = False,
        zoneout_prob: float = 0.0,
    ):
        super().__init__()
        self.hidden_size = hidden_size
        self.cell_clip = cell_clip
        self.use_layer_norm = use_layer_norm
        self.zoneout_prob = zoneout_prob
        
        # Fused weight matrices
        self.weight_ih = nn.Parameter(
            torch.empty(4 * hidden_size, input_size)
        )
        self.weight_hh = nn.Parameter(
            torch.empty(4 * hidden_size, hidden_size)
        )
        self.bias = nn.Parameter(torch.zeros(4 * hidden_size))
        
        # Layer norm (optional)
        if use_layer_norm:
            self.ln_cell = nn.LayerNorm(hidden_size)
            self.ln_hidden = nn.LayerNorm(hidden_size)
        
        # Initialize
        self._reset_parameters(forget_bias)
    
    def _reset_parameters(self, forget_bias):
        # Xavier initialization for input weights
        nn.init.xavier_uniform_(self.weight_ih)
        
        # Orthogonal initialization for hidden weights
        for i in range(4):
            nn.init.orthogonal_(
                self.weight_hh[i*self.hidden_size:(i+1)*self.hidden_size]
            )
        
        # Forget gate bias initialization (critical!)
        with torch.no_grad():
            self.bias[self.hidden_size:2*self.hidden_size].fill_(forget_bias)
    
    def forward(self, x, state):
        h_prev, c_prev = state
        
        # Fused gate computation
        gates = (F.linear(x, self.weight_ih) + 
                 F.linear(h_prev, self.weight_hh) + 
                 self.bias)
        
        i, f, g, o = gates.chunk(4, dim=-1)
        
        # Activations
        i = torch.sigmoid(i)
        f = torch.sigmoid(f)
        g = torch.tanh(g)
        o = torch.sigmoid(o)
        
        # Cell state update
        c = f * c_prev + i * g
        
        # Optional: cell clipping
        if self.cell_clip is not None:
            c = torch.clamp(c, -self.cell_clip, self.cell_clip)
        
        # Optional: layer normalization
        if self.use_layer_norm:
            c = self.ln_cell(c)
        
        # Optional: zoneout (training only)
        if self.training and self.zoneout_prob > 0:
            mask = torch.bernoulli(
                torch.full_like(c, 1 - self.zoneout_prob)
            )
            c = mask * c + (1 - mask) * c_prev
        
        # Hidden state output
        h = o * torch.tanh(c)
        
        if self.use_layer_norm:
            h = self.ln_hidden(h)
        
        return h, c