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

import numpy as np
 
class VanillaRNNCell:
    """
    A single RNN cell implementing the canonical hidden state dynamics.
    
    This implementation exposes the mathematical structure explicitly,
    prioritizing clarity over computational efficiency.
    """
    
    def __init__(self, input_dim: int, hidden_dim: int):
        """
        Initialize RNN cell with Xavier initialization for weights.
        
        Args:
            input_dim: Dimensionality of input vectors (n)
            hidden_dim: Dimensionality of hidden state vectors (d)
        """
        self.input_dim = input_dim
        self.hidden_dim = hidden_dim
        
        # Xavier initialization helps maintain gradient flow
        # Scale for input-to-hidden: sqrt(2 / (input_dim + hidden_dim))
        self.W_xh = np.random.randn(hidden_dim, input_dim) * np.sqrt(2.0 / (input_dim + hidden_dim))
        
        # Scale for hidden-to-hidden: sqrt(2 / (2 * hidden_dim))
        self.W_hh = np.random.randn(hidden_dim, hidden_dim) * np.sqrt(1.0 / hidden_dim)
        
        # Bias initialized to zero
        self.b_h = np.zeros((hidden_dim, 1))
    
    def forward(self, x_t: np.ndarray, h_prev: np.ndarray) -> np.ndarray:
        """
        Compute the hidden state update for a single timestep.
        
        Args:
            x_t: Input vector at time t, shape (input_dim, 1)
            h_prev: Hidden state from time t-1, shape (hidden_dim, 1)
        
        Returns:
            h_t: New hidden state at time t, shape (hidden_dim, 1)
        """
        # Compute the two contributions separately for clarity
        input_contribution = self.W_xh @ x_t      # What the current input contributes
        memory_contribution = self.W_hh @ h_prev  # What the previous state contributes
        
        # Combine and apply nonlinearity
        pre_activation = input_contribution + memory_contribution + self.b_h
        h_t = np.tanh(pre_activation)
        
        return h_t
    
    def forward_sequence(self, x_sequence: np.ndarray, h_0: np.ndarray = None) -> list:
        """
        Process an entire sequence, returning all hidden states.
        
        Args:
            x_sequence: Input sequence, shape (seq_len, input_dim, 1)
            h_0: Initial hidden state (optional, defaults to zeros)
        
        Returns:
            List of hidden states for each timestep
        """
        seq_len = x_sequence.shape[0]
        
        # Initialize hidden state to zeros if not provided
        if h_0 is None:
            h_0 = np.zeros((self.hidden_dim, 1))
        
        hidden_states = [h_0]
        h_t = h_0
        
        for t in range(seq_len):
            h_t = self.forward(x_sequence[t], h_t)
            hidden_states.append(h_t)
        
        return hidden_states  # Length: seq_len + 1 (includes h_0)
 
 
# Demonstration: Processing a simple sequence
if __name__ == "__main__":
    # Create a small RNN cell
    rnn = VanillaRNNCell(input_dim=4, hidden_dim=8)
    
    # Create a random sequence of 5 timesteps
    sequence = np.random.randn(5, 4, 1)
    
    # Process the sequence
    hidden_states = rnn.forward_sequence(sequence)
    
    print(f"Processed sequence of length {len(sequence)}")
    print(f"Got {len(hidden_states)} hidden states (including h_0)")
    print(f"Each hidden state has shape: {hidden_states[0].shape}")
    
    # Observe how hidden states evolve
    for t, h in enumerate(hidden_states):
        print(f"h_{t}: norm = {np.linalg.norm(h):.4f}, range = [{h.min():.3f}, {h.max():.3f}]")

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

import numpy as np
import matplotlib.pyplot as plt
from sklearn.decomposition import PCA
 
def visualize_hidden_trajectories(rnn_cell, sequences, labels, figsize=(12, 8)):
    """
    Visualize how different sequences trace different trajectories
    through hidden state space.
    
    This demonstrates how the hidden state encodes sequence identity
    and how similar sequences have similar trajectories.
    """
    fig, axes = plt.subplots(2, 2, figsize=figsize)
    
    # Collect all hidden states for PCA fitting
    all_hidden_states = []
    trajectories = {}
    
    for seq_name, sequence in sequences.items():
        hidden_states = rnn_cell.forward_sequence(sequence)
        # Convert to 2D array: (seq_len+1, hidden_dim)
        hs_array = np.array([h.squeeze() for h in hidden_states])
        trajectories[seq_name] = hs_array
        all_hidden_states.append(hs_array)
    
    # Stack all states for PCA
    all_states = np.vstack(all_hidden_states)
    
    # Fit PCA to reduce to 2D for visualization
    pca = PCA(n_components=2)
    pca.fit(all_states)
    
    # Plot 1: Trajectories in 2D PCA space
    ax1 = axes[0, 0]
    for seq_name, hs_array in trajectories.items():
        projected = pca.transform(hs_array)
        ax1.plot(projected[:, 0], projected[:, 1], 
                 marker='o', markersize=4, label=seq_name)
        # Mark start and end
        ax1.scatter(projected[0, 0], projected[0, 1], 
                   marker='s', s=100, edgecolors='black', zorder=5)
        ax1.scatter(projected[-1, 0], projected[-1, 1], 
                   marker='*', s=150, edgecolors='black', zorder=5)
    ax1.set_xlabel('PC1')
    ax1.set_ylabel('PC2')
    ax1.set_title('Hidden State Trajectories (PCA projection)')
    ax1.legend()
    ax1.grid(True, alpha=0.3)
    
    # Plot 2: Hidden state norms over time
    ax2 = axes[0, 1]
    for seq_name, hs_array in trajectories.items():
        norms = np.linalg.norm(hs_array, axis=1)
        ax2.plot(norms, marker='o', markersize=4, label=seq_name)
    ax2.set_xlabel('Timestep')
    ax2.set_ylabel('||h_t||_2')
    ax2.set_title('Hidden State Magnitude Over Time')
    ax2.legend()
    ax2.grid(True, alpha=0.3)
    
    # Plot 3: Individual dimension activations
    ax3 = axes[1, 0]
    seq_name = list(trajectories.keys())[0]
    hs_array = trajectories[seq_name]
    im = ax3.imshow(hs_array.T, aspect='auto', cmap='RdBu_r', vmin=-1, vmax=1)
    ax3.set_xlabel('Timestep')
    ax3.set_ylabel('Hidden Dimension')
    ax3.set_title(f'Hidden State Activations: {seq_name}')
    plt.colorbar(im, ax=ax3)
    
    # Plot 4: Distance between consecutive hidden states
    ax4 = axes[1, 1]
    for seq_name, hs_array in trajectories.items():
        distances = np.linalg.norm(np.diff(hs_array, axis=0), axis=1)
        ax4.plot(range(1, len(distances)+1), distances, 
                 marker='o', markersize=4, label=seq_name)
    ax4.set_xlabel('Timestep')
    ax4.set_ylabel('||h_t - h_{t-1}||_2')
    ax4.set_title('Hidden State Change Rate')
    ax4.legend()
    ax4.grid(True, alpha=0.3)
    
    plt.tight_layout()
    return fig
 
 
# Example usage
if __name__ == "__main__":
    np.random.seed(42)
    
    # Create RNN cell
    rnn = VanillaRNNCell(input_dim=3, hidden_dim=16)
    
    # Create different types of sequences
    sequences = {
        'constant': np.ones((10, 3, 1)) * 0.5,
        'increasing': np.linspace(0, 1, 10).reshape(10, 1, 1) * np.ones((1, 3, 1)),
        'random': np.random.randn(10, 3, 1),
        'oscillating': np.sin(np.linspace(0, 4*np.pi, 10)).reshape(10, 1, 1) * np.ones((1, 3, 1)),
    }
    
    fig = visualize_hidden_trajectories(rnn, sequences, labels=None)
    plt.savefig('hidden_state_trajectories.png', dpi=150, bbox_inches='tight')
    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
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

import torch
import torch.nn as nn
 
class RNNWithInitialization(nn.Module):
    """
    RNN that supports multiple hidden state initialization strategies.
    """
    
    def __init__(
        self, 
        input_dim: int, 
        hidden_dim: int, 
        init_strategy: str = 'zero',
        context_dim: int = None
    ):
        super().__init__()
        self.hidden_dim = hidden_dim
        self.init_strategy = init_strategy
        
        # Core RNN cell
        self.rnn_cell = nn.RNNCell(input_dim, hidden_dim, nonlinearity='tanh')
        
        # Initialization depends on strategy
        if init_strategy == 'learned':
            # Learned initial state (trainable parameter)
            self.h0 = nn.Parameter(torch.zeros(1, hidden_dim))
        elif init_strategy == 'context':
            # Context-conditional initialization
            if context_dim is None:
                raise ValueError("context_dim required for 'context' strategy")
            self.context_projection = nn.Linear(context_dim, hidden_dim)
            self.h0 = None
        else:
            # Zero initialization (default)
            self.h0 = None
    
    def get_initial_state(self, batch_size: int, context=None):
        """Get initial hidden state based on strategy."""
        if self.init_strategy == 'zero':
            return torch.zeros(batch_size, self.hidden_dim)
        
        elif self.init_strategy == 'learned':
            # Expand learned h0 to batch size
            return self.h0.expand(batch_size, -1)
        
        elif self.init_strategy == 'context':
            if context is None:
                raise ValueError("Context required for context initialization")
            return torch.tanh(self.context_projection(context))
        
        elif self.init_strategy == 'random':
            return torch.randn(batch_size, self.hidden_dim) * 0.01
        
        else:
            raise ValueError(f"Unknown init strategy: {self.init_strategy}")
    
    def forward(self, x, context=None):
        """
        Process a sequence.
        
        Args:
            x: Input tensor of shape (batch, seq_len, input_dim)
            context: Optional context for initialization (batch, context_dim)
        
        Returns:
            hidden_states: All hidden states (batch, seq_len, hidden_dim)
            final_state: Last hidden state (batch, hidden_dim)
        """
        batch_size, seq_len, _ = x.shape
        
        # Get initial state
        h_t = self.get_initial_state(batch_size, context)
        if x.is_cuda:
            h_t = h_t.cuda()
        
        # Collect hidden states
        hidden_states = []
        
        # Process sequence
        for t in range(seq_len):
            h_t = self.rnn_cell(x[:, t, :], h_t)
            hidden_states.append(h_t)
        
        # Stack: (seq_len, batch, hidden) -> (batch, seq_len, hidden)
        hidden_states = torch.stack(hidden_states, dim=1)
        
        return hidden_states, h_t
 
 
# Demonstration: Compare initialization strategies
def compare_initializations():
    """Compare how different initializations affect hidden state trajectories."""
    torch.manual_seed(42)
    
    input_dim, hidden_dim = 4, 8
    batch_size, seq_len = 2, 20
    
    # Create models with different strategies
    models = {
        'zero': RNNWithInitialization(input_dim, hidden_dim, 'zero'),
        'learned': RNNWithInitialization(input_dim, hidden_dim, 'learned'),
        'context': RNNWithInitialization(input_dim, hidden_dim, 'context', context_dim=3),
    }
    
    # Same input sequence for all models
    x = torch.randn(batch_size, seq_len, input_dim)
    context = torch.randn(batch_size, 3)  # For context-based init
    
    results = {}
    for name, model in models.items():
        if name == 'context':
            hidden_states, final = model(x, context)
        else:
            hidden_states, final = model(x)
        
        # Compute statistics
        mean_norm = hidden_states.norm(dim=-1).mean().item()
        final_norm = final.norm(dim=-1).mean().item()
        
        results[name] = {
            'mean_hidden_norm': mean_norm,
            'final_state_norm': final_norm,
            'init_norm': hidden_states[:, 0, :].norm(dim=-1).mean().item()
        }
        
        print(f"\n{name.upper()} initialization:")
        print(f"  Initial state norm: {results[name]['init_norm']:.4f}")
        print(f"  Mean hidden norm: {mean_norm:.4f}")
        print(f"  Final state norm: {final_norm:.4f}")
 
if __name__ == "__main__":
    compare_initializations()

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

import torch
import torch.nn as nn
from collections import defaultdict
 
class RNNWithDiagnostics(nn.Module):
    """
    RNN that tracks hidden state statistics for debugging and monitoring.
    """
    
    def __init__(self, input_dim, hidden_dim):
        super().__init__()
        self.rnn = nn.RNN(input_dim, hidden_dim, batch_first=True)
        self.hidden_dim = hidden_dim
        
        # Diagnostics storage
        self.diagnostics = defaultdict(list)
        self.track_diagnostics = False
    
    def compute_diagnostics(self, hidden_states, name_prefix=""):
        """Compute and store diagnostic statistics."""
        # hidden_states: (batch, seq_len, hidden_dim)
        
        # Saturation: fraction of values near ±1
        abs_values = hidden_states.abs()
        saturation = (abs_values > 0.9).float().mean().item()
        
        # Per-timestep norms
        norms = hidden_states.norm(dim=-1)  # (batch, seq_len)
        mean_norm = norms.mean().item()
        max_norm = norms.max().item()
        
        # Value distribution
        mean_val = hidden_states.mean().item()
        std_val = hidden_states.std().item()
        
        # Temporal gradient: how much states change between timesteps
        if hidden_states.shape[1] > 1:
            diffs = (hidden_states[:, 1:, :] - hidden_states[:, :-1, :]).norm(dim=-1)
            mean_change = diffs.mean().item()
        else:
            mean_change = 0.0
        
        stats = {
            f'{name_prefix}saturation': saturation,
            f'{name_prefix}mean_norm': mean_norm,
            f'{name_prefix}max_norm': max_norm,
            f'{name_prefix}mean_value': mean_val,
            f'{name_prefix}std_value': std_val,
            f'{name_prefix}mean_change': mean_change,
        }
        
        for key, value in stats.items():
            self.diagnostics[key].append(value)
        
        return stats
    
    def forward(self, x, h_0=None):
        hidden_states, h_n = self.rnn(x, h_0)
        
        if self.track_diagnostics:
            self.compute_diagnostics(hidden_states, name_prefix="hidden_")
        
        return hidden_states, h_n
    
    def get_diagnostic_summary(self):
        """Return summary statistics of tracked diagnostics."""
        summary = {}
        for key, values in self.diagnostics.items():
            if len(values) > 0:
                values_tensor = torch.tensor(values)
                summary[key] = {
                    'mean': values_tensor.mean().item(),
                    'std': values_tensor.std().item(),
                    'min': values_tensor.min().item(),
                    'max': values_tensor.max().item(),
                }
        return summary
    
    def check_health(self):
        """Print diagnostic health check."""
        summary = self.get_diagnostic_summary()
        
        print("\n=== Hidden State Health Check ===")
        
        if 'hidden_saturation' in summary:
            sat = summary['hidden_saturation']['mean']
            if sat > 0.5:
                print(f"⚠️  High saturation: {sat:.2%} of values near ±1")
                print("   → Consider reducing learning rate or checking initialization")
            else:
                print(f"✓  Saturation OK: {sat:.2%}")
        
        if 'hidden_mean_norm' in summary:
            norm = summary['hidden_mean_norm']['mean']
            if norm < 0.1:
                print(f"⚠️  Low hidden norms: {norm:.4f}")
                print("   → Possible vanishing activations")
            elif norm > 0.9 * (self.hidden_dim ** 0.5):
                print(f"⚠️  High hidden norms: {norm:.4f}")
                print("   → Possible saturation or explosion")
            else:
                print(f"✓  Hidden norms OK: {norm:.4f}")
        
        if 'hidden_std_value' in summary:
            std = summary['hidden_std_value']['mean']
            if std < 0.1:
                print(f"⚠️  Low variance: {std:.4f}")
                print("   → Hidden states may be collapsing")
            else:
                print(f"✓  Variance OK: {std:.4f}")
        
        if 'hidden_mean_change' in summary:
            change = summary['hidden_mean_change']['mean']
            print(f"   Temporal dynamics: avg change = {change:.4f}")
 
 
# Example: Training loop with diagnostics
def training_with_diagnostics():
    """Example training loop with hidden state monitoring."""
    model = RNNWithDiagnostics(input_dim=10, hidden_dim=32)
    optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
    
    # Enable diagnostics
    model.track_diagnostics = True
    model.diagnostics.clear()
    
    # Simulated training
    for epoch in range(5):
        # Fake batch
        x = torch.randn(16, 20, 10)  # (batch, seq, features)
        targets = torch.randn(16, 20, 32)  # Fake targets
        
        optimizer.zero_grad()
        hidden_states, _ = model(x)
        loss = ((hidden_states - targets) ** 2).mean()
        loss.backward()
        optimizer.step()
        
        print(f"Epoch {epoch+1}, Loss: {loss.item():.4f}")
    
    # Check health after training
    model.check_health()
 
if __name__ == "__main__":
    training_with_diagnostics()