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

import numpy as np
import matplotlib.pyplot as plt
 
def analyze_vanishing_conditions(hidden_size, sequence_lengths, num_trials=100):
    """
    Comprehensive analysis of vanishing gradient conditions.
    
    Simulates gradient flow with realistic activation statistics
    to demonstrate vanishing at different sequence lengths.
    """
    results = {
        'tanh': {'lengths': sequence_lengths, 'final_norms': []},
        'sigmoid': {'lengths': sequence_lengths, 'final_norms': []},
        'relu': {'lengths': sequence_lengths, 'final_norms': []}
    }
    
    for activation in ['tanh', 'sigmoid', 'relu']:
        print(f"\nAnalyzing {activation} activation:")
        print("-" * 40)
        
        final_norms_per_length = []
        
        for T in sequence_lengths:
            trial_norms = []
            
            for trial in range(num_trials):
                # Initialize recurrent weights (orthogonal for stability)
                Whh = np.linalg.qr(np.random.randn(hidden_size, hidden_size))[0]
                
                # Simulate forward pass to get realistic activation statistics
                h = np.random.randn(hidden_size, 1) * 0.5
                pre_activations = []
                
                for t in range(T):
                    x = np.random.randn(hidden_size, 1) * 0.1
                    z = Whh @ h + x
                    
                    if activation == 'tanh':
                        h = np.tanh(z)
                        deriv = 1 - h**2
                    elif activation == 'sigmoid':
                        h = 1 / (1 + np.exp(-z))
                        deriv = h * (1 - h)
                    else:  # relu
                        h = np.maximum(0, z)
                        deriv = (z > 0).astype(float)
                    
                    pre_activations.append(deriv)
                
                # Backward pass: compute gradient norm
                grad = np.ones((hidden_size, 1))
                grad = grad / np.linalg.norm(grad)
                
                for t in range(T-1, -1, -1):
                    D = np.diag(pre_activations[t].flatten())
                    grad = Whh.T @ D @ grad
                
                final_norm = np.linalg.norm(grad)
                trial_norms.append(final_norm)
            
            mean_norm = np.mean(trial_norms)
            std_norm = np.std(trial_norms)
            final_norms_per_length.append(mean_norm)
            
            print(f"  T={T:3d}: gradient norm = {mean_norm:.2e} ± {std_norm:.2e}")
        
        results[activation]['final_norms'] = final_norms_per_length
    
    return results
 
 
def theoretical_vanishing_curves():
    """
    Plot theoretical vanishing curves for different gamma * sigma values.
    """
    T = np.arange(0, 101)
    
    products = [0.99, 0.95, 0.9, 0.8, 0.7, 0.5]
    
    fig, ax = plt.subplots(figsize=(10, 6))
    
    for prod in products:
        decay = prod ** T
        ax.semilogy(T, decay, label=f'γσ = {prod}')
    
    ax.axhline(y=1e-6, color='r', linestyle='--', label='Practical threshold')
    ax.set_xlabel('Temporal Distance (t - k)')
    ax.set_ylabel('Gradient Magnitude (relative)')
    ax.set_title('Theoretical Gradient Decay Curves')
    ax.legend()
    ax.grid(True, alpha=0.3)
    ax.set_xlim(0, 100)
    ax.set_ylim(1e-20, 10)
    
    return fig
 
 
def effective_memory_calculation():
    """
    Calculate effective memory for different configurations.
    """
    print("\n" + "=" * 60)
    print("EFFECTIVE MEMORY ANALYSIS")
    print("=" * 60)
    print("\n(Memory = timesteps until gradient falls below 1e-6)")
    print("-" * 60)
    
    # Different configurations
    configs = [
        ('Sigmoid, small W', 0.2, 0.9),
        ('Sigmoid, large W', 0.2, 1.2),
        ('Tanh, small W', 0.5, 0.9),
        ('Tanh, typical W', 0.7, 1.0),
        ('Tanh, orthogonal W', 0.8, 1.0),
        ('ReLU (no dying)', 1.0, 0.9),
        ('ReLU (no dying)', 1.0, 1.0),
        ('LSTM-like (gated)', 0.95, 1.0),
    ]
    
    threshold = 1e-6
    
    for name, gamma, sigma in configs:
        product = gamma * sigma
        
        if product >= 1:
            memory = float('inf')
            memory_str = "∞ (exploding)"
        else:
            # Solve: product^T = threshold
            memory = np.log(threshold) / np.log(product)
            memory_str = f"{int(memory)}"
        
        print(f"{name:25s}  γ={gamma:.2f}, σ={sigma:.2f}, product={product:.2f}, memory={memory_str}")
 
 
# Run analysis
np.random.seed(42)
hidden_size = 100
sequence_lengths = [10, 25, 50, 75, 100]
 
print("=" * 60)
print("VANISHING GRADIENT ANALYSIS")
print("=" * 60)
 
results = analyze_vanishing_conditions(hidden_size, sequence_lengths)
 
# Theoretical curves
fig = theoretical_vanishing_curves()
plt.savefig('vanishing_gradient_curves.png', dpi=150, bbox_inches='tight')
 
# Effective memory
effective_memory_calculation()

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

import numpy as np
import matplotlib.pyplot as plt
 
def activation_derivative_analysis():
    """
    Comprehensive analysis of activation derivatives and their
    impact on gradient flow.
    """
    z = np.linspace(-6, 6, 1000)
    
    # Activations
    sigmoid = 1 / (1 + np.exp(-z))
    tanh = np.tanh(z)
    relu = np.maximum(0, z)
    
    # Derivatives
    sigmoid_deriv = sigmoid * (1 - sigmoid)
    tanh_deriv = 1 - tanh**2
    relu_deriv = (z > 0).astype(float)
    
    # Create visualization
    fig, axes = plt.subplots(2, 2, figsize=(14, 10))
    
    # Plot activations
    axes[0, 0].plot(z, sigmoid, label='Sigmoid', linewidth=2)
    axes[0, 0].plot(z, tanh, label='Tanh', linewidth=2)
    axes[0, 0].plot(z, relu, label='ReLU', linewidth=2)
    axes[0, 0].axhline(y=0, color='gray', linestyle='--', alpha=0.5)
    axes[0, 0].axvline(x=0, color='gray', linestyle='--', alpha=0.5)
    axes[0, 0].set_xlabel('z')
    axes[0, 0].set_ylabel('φ(z)')
    axes[0, 0].set_title('Activation Functions')
    axes[0, 0].legend()
    axes[0, 0].grid(True, alpha=0.3)
    axes[0, 0].set_ylim(-1.5, 3)
    
    # Plot derivatives
    axes[0, 1].plot(z, sigmoid_deriv, label='Sigmoid′', linewidth=2)
    axes[0, 1].plot(z, tanh_deriv, label='Tanh′', linewidth=2)
    axes[0, 1].plot(z, relu_deriv, label='ReLU′', linewidth=2)
    axes[0, 1].axhline(y=1, color='gray', linestyle='--', alpha=0.5, label='Ideal (no shrinkage)')
    axes[0, 1].axhline(y=0.25, color='red', linestyle=':', alpha=0.5, label='Sigmoid max')
    axes[0, 1].set_xlabel('z')
    axes[0, 1].set_ylabel("φ'(z)")
    axes[0, 1].set_title('Activation Derivatives')
    axes[0, 1].legend()
    axes[0, 1].grid(True, alpha=0.3)
    axes[0, 1].set_ylim(-0.1, 1.2)
    
    # Gradient decay over timesteps
    timesteps = np.arange(0, 51)
    
    # Simulate with different average derivatives
    avg_derivs = {
        'Sigmoid (avg=0.15)': 0.15,
        'Tanh (avg=0.5)': 0.5,
        'Tanh (avg=0.7)': 0.7,
        'Tanh (avg=0.85)': 0.85,
        'ReLU (50% active)': 0.5,
        'Ideal (γ=1)': 1.0
    }
    
    # Assume spectral norm = 1 (orthogonal weights)
    sigma = 1.0
    
    axes[1, 0].axhline(y=1e-6, color='red', linestyle='--', label='Practical threshold')
    
    for name, gamma in avg_derivs.items():
        decay = (gamma * sigma) ** timesteps
        axes[1, 0].semilogy(timesteps, decay, label=name, linewidth=2)
    
    axes[1, 0].set_xlabel('Timesteps Backward')
    axes[1, 0].set_ylabel('Relative Gradient Magnitude')
    axes[1, 0].set_title('Gradient Decay by Activation Type')
    axes[1, 0].legend(fontsize=9)
    axes[1, 0].grid(True, alpha=0.3)
    axes[1, 0].set_ylim(1e-15, 10)
    
    # Distribution of derivatives in typical operation
    # Simulate hidden states and compute derivative distribution
    np.random.seed(42)
    n_samples = 10000
    
    # Hidden states tend to be distributed around 0 with some variance
    h_samples = np.random.randn(n_samples) * 0.8
    
    sigmoid_derivs = sigmoid_deriv_at = (1 / (1 + np.exp(-h_samples))) * (1 - 1 / (1 + np.exp(-h_samples)))
    tanh_derivs = 1 - np.tanh(h_samples)**2
    
    axes[1, 1].hist(sigmoid_derivs, bins=50, alpha=0.7, label='Sigmoid derivatives', density=True)
    axes[1, 1].hist(tanh_derivs, bins=50, alpha=0.7, label='Tanh derivatives', density=True)
    axes[1, 1].axvline(x=np.mean(sigmoid_derivs), color='blue', linestyle='--', 
                       label=f'Sigmoid mean: {np.mean(sigmoid_derivs):.3f}')
    axes[1, 1].axvline(x=np.mean(tanh_derivs), color='orange', linestyle='--',
                       label=f'Tanh mean: {np.mean(tanh_derivs):.3f}')
    axes[1, 1].set_xlabel("φ'(z)")
    axes[1, 1].set_ylabel('Density')
    axes[1, 1].set_title('Distribution of Derivatives in Typical Operation')
    axes[1, 1].legend()
    axes[1, 1].grid(True, alpha=0.3)
    
    plt.tight_layout()
    return fig
 
 
def saturation_analysis():
    """
    Analyze how saturation affects gradient flow in practice.
    """
    print("\n" + "=" * 60)
    print("SATURATION ANALYSIS")
    print("=" * 60)
    
    def compute_saturation_stats(activations, threshold=0.1):
        """Compute fraction of saturated units (derivative < threshold)."""
        return np.mean(activations < threshold)
    
    # Simulate RNN dynamics
    np.random.seed(42)
    hidden_size = 256
    seq_length = 50
    n_trials = 100
    
    # Different initialization scales
    scales = [0.5, 1.0, 1.5, 2.0]
    
    print("\nSaturation fraction (% of units with tanh'(z) < 0.1):")
    print("-" * 60)
    
    for scale in scales:
        sat_fractions = []
        
        for trial in range(n_trials):
            Whh = np.random.randn(hidden_size, hidden_size) * scale / np.sqrt(hidden_size)
            Wxh = np.random.randn(hidden_size, hidden_size) * 0.1
            
            h = np.zeros((hidden_size, 1))
            
            trial_saturations = []
            for t in range(seq_length):
                x = np.random.randn(hidden_size, 1) * 0.1
                z = Whh @ h + Wxh @ x
                h = np.tanh(z)
                
                # Compute tanh derivative
                tanh_deriv = 1 - h**2
                sat_frac = np.mean(tanh_deriv < 0.1)
                trial_saturations.append(sat_frac)
            
            sat_fractions.append(np.mean(trial_saturations))
        
        print(f"  Scale={scale:.1f}: {100*np.mean(sat_fractions):.1f}% ± {100*np.std(sat_fractions):.1f}%")
 
 
# Run analyses
fig = activation_derivative_analysis()
plt.savefig('activation_analysis.png', dpi=150, bbox_inches='tight')
 
saturation_analysis()

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
222
223
224
225
226
227
228
229

import numpy as np
import torch
import torch.nn as nn
from typing import Dict, List, Tuple
import matplotlib.pyplot as plt
 
class VanishingGradientDetector:
    """
    Comprehensive toolkit for detecting vanishing gradients in RNNs.
    """
    
    def __init__(self, model: nn.Module, verbose: bool = True):
        self.model = model
        self.verbose = verbose
        self.gradient_history: List[Dict] = []
    
    def compute_hidden_state_gradients(
        self, 
        inputs: torch.Tensor,  # [seq_len, batch, input_size]
        targets: torch.Tensor,
        criterion: nn.Module
    ) -> Dict[str, np.ndarray]:
        """
        Compute gradient norms at each hidden state.
        """
        seq_len = inputs.size(0)
        hidden_grads = []
        
        # Forward pass with gradient tracking
        self.model.zero_grad()
        hidden = None
        hiddens = []
        
        for t in range(seq_len):
            if hidden is None:
                output, hidden = self.model(inputs[t:t+1])
            else:
                output, hidden = self.model(inputs[t:t+1], hidden)
            
            # Store hidden state (detach for storage, keep graph intact)
            if isinstance(hidden, tuple):  # LSTM returns (h, c)
                hiddens.append(hidden[0].clone())
            else:
                hiddens.append(hidden.clone())
        
        # Compute loss
        loss = criterion(output, targets)
        
        # Backward pass
        loss.backward()
        
        # Now compute gradient norms at each timestep
        # Re-run with gradient accumulation tracking
        self.model.zero_grad()
        
        accum_grads = []
        hidden = None
        
        for t in range(seq_len):
            if hidden is None:
                output, hidden = self.model(inputs[t:t+1])
            else:
                # Ensure hidden requires grad
                if isinstance(hidden, tuple):
                    hidden = (hidden[0].detach().requires_grad_(True),
                             hidden[1].detach().requires_grad_(True))
                else:
                    hidden = hidden.detach().requires_grad_(True)
                output, hidden = self.model(inputs[t:t+1], hidden)
        
        loss = criterion(output, targets)
        loss.backward()
        
        # Collect gradient information
        gradient_info = {
            'loss': loss.item(),
            'param_grad_norms': {},
        }
        
        for name, param in self.model.named_parameters():
            if param.grad is not None:
                gradient_info['param_grad_norms'][name] = param.grad.norm().item()
        
        return gradient_info
    
    def analyze_gradient_flow(
        self,
        inputs: torch.Tensor,
        targets: torch.Tensor,
        criterion: nn.Module
    ) -> Dict:
        """
        Comprehensive gradient flow analysis.
        """
        seq_len = inputs.size(0)
        batch_size = inputs.size(1)
        
        # Method: Compute gradients at each hidden state by running BPTT
        # and checking gradient norms for hidden->hidden weights
        
        self.model.zero_grad()
        
        # Get all hidden states
        hiddens = []
        hidden = None
        
        for t in range(seq_len):
            x_t = inputs[t:t+1]
            if hidden is None:
                out, hidden = self.model(x_t)
            else:
                out, hidden = self.model(x_t, hidden)
            
            if isinstance(hidden, tuple):
                hiddens.append(hidden[0])
            else:
                hiddens.append(hidden)
        
        # Compute loss
        loss = criterion(out, targets)
        loss.backward()
        
        # Analyze recurrent weight gradients
        recurrent_grads = {}
        for name, param in self.model.named_parameters():
            if 'weight_hh' in name or 'Whh' in name:
                if param.grad is not None:
                    recurrent_grads[name] = {
                        'norm': param.grad.norm().item(),
                        'max': param.grad.abs().max().item(),
                        'mean': param.grad.abs().mean().item(),
                    }
        
        # Compute vanishing indicators
        analysis = {
            'sequence_length': seq_len,
            'loss': loss.item(),
            'recurrent_gradients': recurrent_grads,
            'vanishing_detected': False,
            'severity': 'none'
        }
        
        # Check for vanishing: if recurrent gradients are extremely small
        if recurrent_grads:
            max_grad = max(g['norm'] for g in recurrent_grads.values())
            if max_grad < 1e-10:
                analysis['vanishing_detected'] = True
                analysis['severity'] = 'severe'
            elif max_grad < 1e-6:
                analysis['vanishing_detected'] = True
                analysis['severity'] = 'moderate'
            elif max_grad < 1e-3:
                analysis['vanishing_detected'] = True
                analysis['severity'] = 'mild'
        
        return analysis
    
    @staticmethod
    def gradient_flow_test(hidden_size: int = 128, seq_lengths: List[int] = None):
        """
        Run a standard test for gradient flow in vanilla RNN.
        """
        if seq_lengths is None:
            seq_lengths = [10, 25, 50, 100, 200]
        
        print("=" * 60)
        print("GRADIENT FLOW TEST")
        print("=" * 60)
        print(f"Hidden size: {hidden_size}")
        print("-" * 60)
        
        results = []
        
        for T in seq_lengths:
            # Create simple RNN
            rnn = nn.RNN(input_size=32, hidden_size=hidden_size, 
                        num_layers=1, batch_first=False)
            linear = nn.Linear(hidden_size, 10)
            
            # Forward pass
            x = torch.randn(T, 1, 32)
            target = torch.randn(1, 10)
            
            h0 = torch.zeros(1, 1, hidden_size)
            output, hn = rnn(x, h0)
            pred = linear(output[-1])
            
            loss = nn.MSELoss()(pred, target)
            loss.backward()
            
            # Get gradient norms
            rnn_grad_norm = rnn.weight_hh_l0.grad.norm().item()
            
            results.append({
                'T': T,
                'grad_norm': rnn_grad_norm,
                'loss': loss.item()
            })
            
            print(f"T={T:4d}: gradient norm = {rnn_grad_norm:.2e}")
        
        # Check decay rate
        if len(results) >= 2:
            log_norms = [np.log(r['grad_norm'] + 1e-20) for r in results]
            log_Ts = [np.log(r['T']) for r in results]
            
            # Fit decay: log(grad) = a * T + b
            # Actually for exponential: log(grad) = log(c) + T * log(r)
            # Linear fit: log_norm vs T
            Ts = [r['T'] for r in results]
            coeffs = np.polyfit(Ts, log_norms, 1)
            decay_rate = np.exp(coeffs[0])
            
            print("-" * 60)
            print(f"Estimated decay rate per timestep: {decay_rate:.4f}")
            if decay_rate < 0.95:
                print("⚠️  VANISHING GRADIENTS DETECTED")
                print(f"   Effective memory: ~{int(np.log(1e-6) / np.log(decay_rate))} timesteps")
            else:
                print("✓  Gradient flow appears stable")
        
        return results
 
 
# Run the gradient flow test
results = VanishingGradientDetector.gradient_flow_test(
    hidden_size=128,
    seq_lengths=[10, 25, 50, 75, 100, 150, 200]
)

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

import numpy as np
import torch
import torch.nn as nn
 
class LeakyIntegrationRNN(nn.Module):
    """
    Leaky integration RNN - provides direct path for gradient flow.
    
    h_t = (1 - alpha) * h_{t-1} + alpha * tanh(W_hh @ h_{t-1} + W_xh @ x_t)
    
    The (1 - alpha) term allows gradients to flow with less shrinkage.
    """
    
    def __init__(self, input_size, hidden_size, alpha=0.1):
        super().__init__()
        self.hidden_size = hidden_size
        self.alpha = alpha
        
        self.Wxh = nn.Linear(input_size, hidden_size, bias=False)
        self.Whh = nn.Linear(hidden_size, hidden_size, bias=False)
        self.bias = nn.Parameter(torch.zeros(hidden_size))
    
    def forward(self, x, h=None):
        """
        x: [seq_len, batch, input_size]
        h: [batch, hidden_size] or None
        """
        seq_len, batch_size, _ = x.size()
        
        if h is None:
            h = torch.zeros(batch_size, self.hidden_size, device=x.device)
        
        outputs = []
        for t in range(seq_len):
            # Leaky integration
            h_new = torch.tanh(self.Whh(h) + self.Wxh(x[t]) + self.bias)
            h = (1 - self.alpha) * h + self.alpha * h_new
            outputs.append(h)
        
        return torch.stack(outputs), h
 
 
class IRNN(nn.Module):
    """
    Identity-initialized RNN (Le et al., 2015).
    
    - W_hh initialized to identity matrix
    - ReLU activation instead of tanh
    - Biases initialized to zero
    """
    
    def __init__(self, input_size, hidden_size):
        super().__init__()
        self.hidden_size = hidden_size
        
        self.Wxh = nn.Linear(input_size, hidden_size, bias=False)
        self.Whh = nn.Linear(hidden_size, hidden_size, bias=False)
        self.bias = nn.Parameter(torch.zeros(hidden_size))
        
        # Identity initialization for recurrent weights
        nn.init.eye_(self.Whh.weight)
        
        # Small initialization for input weights
        nn.init.normal_(self.Wxh.weight, std=0.001)
    
    def forward(self, x, h=None):
        seq_len, batch_size, _ = x.size()
        
        if h is None:
            h = torch.zeros(batch_size, self.hidden_size, device=x.device)
        
        outputs = []
        for t in range(seq_len):
            # ReLU activation (not tanh)
            h = torch.relu(self.Whh(h) + self.Wxh(x[t]) + self.bias)
            outputs.append(h)
        
        return torch.stack(outputs), h
 
 
def compare_gradient_flow():
    """
    Compare gradient flow across different RNN architectures.
    """
    input_size = 32
    hidden_size = 64
    seq_lengths = [10, 25, 50, 100, 200]
    
    architectures = {
        'Vanilla RNN': nn.RNN(input_size, hidden_size, batch_first=False),
        'Leaky (α=0.1)': LeakyIntegrationRNN(input_size, hidden_size, alpha=0.1),
        'Leaky (α=0.5)': LeakyIntegrationRNN(input_size, hidden_size, alpha=0.5),
        'IRNN': IRNN(input_size, hidden_size),
        'LSTM': nn.LSTM(input_size, hidden_size, batch_first=False),
    }
    
    results = {name: [] for name in architectures}
    
    print("=" * 70)
    print("GRADIENT FLOW COMPARISON ACROSS ARCHITECTURES")
    print("=" * 70)
    
    for T in seq_lengths:
        print(f"\nSequence length T = {T}:")
        print("-" * 50)
        
        x = torch.randn(T, 1, input_size)
        target = torch.randn(1, hidden_size)
        
        for name, model in architectures.items():
            model.zero_grad()
            
            if isinstance(model, nn.RNN) or isinstance(model, nn.LSTM):
                output, _ = model(x)
                pred = output[-1]
            else:
                output, _ = model(x)
                pred = output[-1]
            
            loss = nn.MSELoss()(pred, target)
            loss.backward()
            
            # Get recurrent weight gradient norm
            grad_norm = 0
            for pname, param in model.named_parameters():
                if 'hh' in pname.lower() or 'whh' in pname.lower():
                    if param.grad is not None:
                        grad_norm += param.grad.norm().item() ** 2
            grad_norm = np.sqrt(grad_norm)
            
            results[name].append(grad_norm)
            print(f"  {name:20s}: grad_norm = {grad_norm:.2e}")
    
    return results, seq_lengths
 
 
# Run comparison
results, seq_lengths = compare_gradient_flow()
 
# Print summary
print("\n" + "=" * 70)
print("SUMMARY: Gradient norm at T=200 relative to T=10")
print("=" * 70)
for name, norms in results.items():
    ratio = norms[-1] / norms[0] if norms[0] > 0 else 0
    print(f"{name:20s}: {ratio:.2e}")