Conditional Random Fields

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

import numpy as np
from typing import List, Tuple, Dict
 
def compute_gradient(
    sequences: List[Tuple[List[str], List[int]]],  # (tokens, labels) pairs
    weights: np.ndarray,                            # Current weights (K,)
    num_labels: int,
    feature_extractor,                              # Extract features
    forward_backward_fn,                            # Compute marginals
) -> Tuple[np.ndarray, float]:
    """
    Compute gradient of log-likelihood with respect to weights.
    
    Returns:
        gradient: Gradient vector (K,)
        log_likelihood: Current log-likelihood value
    """
    K = len(weights)
    gradient = np.zeros(K)
    total_log_likelihood = 0.0
    
    for tokens, correct_labels in sequences:
        n = len(tokens)
        
        # =====================
        # Step 1: Observed features (for correct label sequence)
        # =====================
        observed_features = np.zeros(K)
        for i in range(n):
            y = correct_labels[i]
            y_prev = correct_labels[i-1] if i > 0 else -1
            
            # Get active features and accumulate
            active = feature_extractor.get_features(tokens, i, y, y_prev)
            for feat_id, value in active:
                observed_features[feat_id] += value
        
        # =====================
        # Step 2: Expected features (under current model)
        # =====================
        # Run forward-backward to get marginals
        log_alpha, log_beta, log_Z, marginals, pairwise = forward_backward_fn(
            tokens, weights, num_labels, feature_extractor
        )
        
        expected_features = np.zeros(K)
        for i in range(n):
            if i == 0:
                # First position: only unigram marginals
                for y in range(num_labels):
                    prob = marginals[i, y]
                    active = feature_extractor.get_features(tokens, i, y, -1)
                    for feat_id, value in active:
                        expected_features[feat_id] += prob * value
            else:
                # Other positions: pairwise marginals
                for y_prev in range(num_labels):
                    for y in range(num_labels):
                        prob = pairwise[i-1, y_prev, y]
                        active = feature_extractor.get_features(tokens, i, y, y_prev)
                        for feat_id, value in active:
                            expected_features[feat_id] += prob * value
        
        # =====================
        # Step 3: Accumulate gradient and likelihood
        # =====================
        # Gradient = observed - expected
        gradient += observed_features - expected_features
        
        # Log-likelihood = score(correct) - log Z
        score = np.dot(weights, observed_features)
        total_log_likelihood += score - log_Z
    
    return gradient, total_log_likelihood
 
 
def compute_gradient_with_regularization(
    gradient: np.ndarray,
    weights: np.ndarray,
    regularization: str = 'l2',
    reg_strength: float = 1.0
) -> np.ndarray:
    """
    Add regularization gradient component.
    
    For L2: penalty = (c/2) ||λ||² → gradient contribution = -c * λ
    For L1: penalty = c ||λ||₁ → gradient contribution = -c * sign(λ)
    
    Note: We're MAXIMIZING log-likelihood, so we ADD the penalty gradient
    (which is negative, pulling weights toward zero).
    """
    if regularization == 'l2':
        # L2 regularization: gradient -= c * λ
        reg_gradient = -reg_strength * weights
    elif regularization == 'l1':
        # L1 regularization: gradient -= c * sign(λ)
        reg_gradient = -reg_strength * np.sign(weights)
    else:
        reg_gradient = np.zeros_like(weights)
    
    return gradient + reg_gradient
 
 
# Numerical gradient check (for verification)
def numerical_gradient(
    sequences: List[Tuple[List[str], List[int]]],
    weights: np.ndarray,
    compute_ll_fn,
    epsilon: float = 1e-5
) -> np.ndarray:
    """
    Compute gradient numerically for verification.
    
    For each weight dimension:
    ∂L/∂λ_k ≈ (L(λ+ε·e_k) - L(λ-ε·e_k)) / (2ε)
    """
    K = len(weights)
    numerical_grad = np.zeros(K)
    
    for k in range(K):
        weights_plus = weights.copy()
        weights_plus[k] += epsilon
        ll_plus = compute_ll_fn(sequences, weights_plus)
        
        weights_minus = weights.copy()
        weights_minus[k] -= epsilon
        ll_minus = compute_ll_fn(sequences, weights_minus)
        
        numerical_grad[k] = (ll_plus - ll_minus) / (2 * epsilon)
    
    return numerical_grad
 
 
print("Gradient computation for CRF training")
print("=" * 50)
print("Key formula: ∇L = F_observed - E[F_expected]")
print("")
print("This is the same gradient structure as:")
print("  - Logistic regression")
print("  - Maximum entropy models")
print("  - All exponential family distributions")

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

import numpy as np
from scipy.optimize import minimize
from typing import List, Tuple, Callable
 
class CRFTrainer:
    """
    CRF trainer supporting multiple optimization algorithms.
    """
    
    def __init__(
        self,
        num_features: int,
        num_labels: int,
        feature_extractor,
        regularization: str = 'l2',
        reg_strength: float = 1.0
    ):
        self.num_features = num_features
        self.num_labels = num_labels
        self.feature_extractor = feature_extractor
        self.regularization = regularization
        self.reg_strength = reg_strength
        self.weights = None
    
    def _objective_and_gradient(
        self,
        weights: np.ndarray,
        sequences: List[Tuple[List[str], List[int]]]
    ) -> Tuple[float, np.ndarray]:
        """
        Compute negative log-likelihood and its gradient.
        
        Note: We return NEGATIVE because scipy.optimize.minimize minimizes.
        """
        # Compute log-likelihood and gradient
        gradient, log_likelihood = compute_gradient(
            sequences, weights, self.num_labels,
            self.feature_extractor, self.forward_backward
        )
        
        # Add regularization
        if self.regularization == 'l2':
            reg_term = 0.5 * self.reg_strength * np.sum(weights ** 2)
            reg_grad = self.reg_strength * weights
        elif self.regularization == 'l1':
            reg_term = self.reg_strength * np.sum(np.abs(weights))
            reg_grad = self.reg_strength * np.sign(weights)
        else:
            reg_term = 0.0
            reg_grad = np.zeros_like(weights)
        
        # Return negative (for minimization)
        neg_ll = -(log_likelihood - reg_term)
        neg_grad = -(gradient - reg_grad)
        
        return neg_ll, neg_grad
    
    def train_lbfgs(
        self,
        sequences: List[Tuple[List[str], List[int]]],
        max_iter: int = 100,
        verbose: bool = True
    ) -> np.ndarray:
        """
        Train CRF using L-BFGS optimizer.
        
        This is the standard approach for classical CRFs.
        """
        # Initialize weights
        weights = np.zeros(self.num_features)
        
        # Track iterations
        iteration = [0]
        
        def callback(xk):
            if verbose and iteration[0] % 10 == 0:
                neg_ll, _ = self._objective_and_gradient(xk, sequences)
                print(f"Iteration {iteration[0]}: neg_ll = {neg_ll:.4f}")
            iteration[0] += 1
        
        # Run L-BFGS-B (bounded L-BFGS)
        result = minimize(
            fun=lambda w: self._objective_and_gradient(w, sequences),
            x0=weights,
            method='L-BFGS-B',
            jac=True,  # Objective returns gradient
            options={'maxiter': max_iter, 'disp': verbose},
            callback=callback
        )
        
        self.weights = result.x
        return self.weights
    
    def train_sgd(
        self,
        sequences: List[Tuple[List[str], List[int]]],
        epochs: int = 10,
        learning_rate: float = 0.1,
        batch_size: int = 32,
        lr_decay: float = 0.01,
        verbose: bool = True
    ) -> np.ndarray:
        """
        Train CRF using Stochastic Gradient Descent.
        
        Better for large datasets; required for neural CRFs.
        """
        weights = np.zeros(self.num_features)
        n_sequences = len(sequences)
        
        for epoch in range(epochs):
            # Shuffle data
            indices = np.random.permutation(n_sequences)
            epoch_loss = 0.0
            
            # Mini-batch iteration
            for batch_start in range(0, n_sequences, batch_size):
                batch_indices = indices[batch_start:batch_start + batch_size]
                batch = [sequences[i] for i in batch_indices]
                
                # Compute gradient on batch
                gradient, batch_ll = compute_gradient(
                    batch, weights, self.num_labels,
                    self.feature_extractor, self.forward_backward
                )
                
                # Add regularization gradient
                if self.regularization == 'l2':
                    gradient -= self.reg_strength * weights
                    reg_term = 0.5 * self.reg_strength * np.sum(weights ** 2)
                else:
                    reg_term = 0
                
                # Decayed learning rate
                t = epoch * (n_sequences // batch_size) + batch_start // batch_size
                lr = learning_rate / (1 + lr_decay * t)
                
                # Update weights (gradient ascent for maximization)
                weights += lr * gradient
                epoch_loss += -batch_ll + reg_term
            
            if verbose:
                avg_loss = epoch_loss / (n_sequences // batch_size)
                print(f"Epoch {epoch+1}/{epochs}: avg_loss = {avg_loss:.4f}")
        
        self.weights = weights
        return self.weights
    
    def train_adam(
        self,
        sequences: List[Tuple[List[str], List[int]]],
        epochs: int = 10,
        learning_rate: float = 0.001,
        batch_size: int = 32,
        beta1: float = 0.9,
        beta2: float = 0.999,
        epsilon: float = 1e-8,
        verbose: bool = True
    ) -> np.ndarray:
        """
        Train CRF using Adam optimizer.
        
        Modern variant of SGD with adaptive learning rates.
        Standard for neural CRF training.
        """
        weights = np.zeros(self.num_features)
        m = np.zeros_like(weights)  # First moment
        v = np.zeros_like(weights)  # Second moment
        t = 0
        
        n_sequences = len(sequences)
        
        for epoch in range(epochs):
            indices = np.random.permutation(n_sequences)
            
            for batch_start in range(0, n_sequences, batch_size):
                t += 1
                batch_indices = indices[batch_start:batch_start + batch_size]
                batch = [sequences[i] for i in batch_indices]
                
                # Compute gradient
                gradient, _ = compute_gradient(
                    batch, weights, self.num_labels,
                    self.feature_extractor, self.forward_backward
                )
                
                # Add regularization
                if self.regularization == 'l2':
                    gradient -= self.reg_strength * weights
                
                # Negate for gradient descent (we're maximizing)
                gradient = -gradient
                
                # Update moments
                m = beta1 * m + (1 - beta1) * gradient
                v = beta2 * v + (1 - beta2) * (gradient ** 2)
                
                # Bias correction
                m_hat = m / (1 - beta1 ** t)
                v_hat = v / (1 - beta2 ** t)
                
                # Update weights
                weights -= learning_rate * m_hat / (np.sqrt(v_hat) + epsilon)
            
            if verbose:
                print(f"Epoch {epoch+1}/{epochs} complete")
        
        self.weights = weights
        return self.weights
 
 
# Algorithm comparison
print("CRF Optimization Algorithms Comparison")
print("=" * 60)
print(f"{'Algorithm':<15} {'Best For':<25} {'Hyperparameters':<20}")
print("-" * 60)
print(f"{'L-BFGS':<15} {'Small/medium data':<25} {'None (line search)':<20}")
print(f"{'SGD':<15} {'Large data, online':<25} {'learning_rate, decay':<20}")
print(f"{'Adam':<15} {'Neural CRFs':<25} {'lr, beta1, beta2':<20}")

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

import torch
import torch.nn as nn
from typing import List, Tuple
 
class BiLSTMCRF(nn.Module):
    """
    BiLSTM-CRF model for sequence labeling.
    
    Combines neural feature learning with structured prediction.
    """
    
    def __init__(
        self,
        vocab_size: int,
        embedding_dim: int,
        hidden_dim: int,
        num_labels: int,
        padding_idx: int = 0
    ):
        super().__init__()
        
        self.num_labels = num_labels
        self.hidden_dim = hidden_dim
        
        # Embedding layer
        self.embedding = nn.Embedding(
            vocab_size, embedding_dim, padding_idx=padding_idx
        )
        
        # BiLSTM encoder
        self.lstm = nn.LSTM(
            embedding_dim, hidden_dim // 2,
            num_layers=1, bidirectional=True, batch_first=True
        )
        
        # Emission projection: hidden -> label scores
        self.emission = nn.Linear(hidden_dim, num_labels)
        
        # CRF transition matrix: T[i,j] = score of transitioning from i to j
        # +2 for START and END tags
        self.transitions = nn.Parameter(
            torch.randn(num_labels + 2, num_labels + 2)
        )
        self.START_TAG = num_labels
        self.END_TAG = num_labels + 1
        
        # Constrain impossible transitions
        self.transitions.data[self.START_TAG, :] = -10000  # Can't transition TO start
        self.transitions.data[:, self.END_TAG] = -10000    # Can't transition FROM end
    
    def _forward_alg(self, emissions: torch.Tensor) -> torch.Tensor:
        """
        Forward algorithm in log-space.
        
        Args:
            emissions: Emission scores (seq_len, num_labels)
        
        Returns:
            log_Z: Log partition function (scalar)
        """
        seq_len = emissions.size(0)
        
        # Initialize: α[START] = 0, others = -inf
        # In expanded space: α[0, START] = 0
        alpha = torch.full((self.num_labels + 2,), -10000.0, device=emissions.device)
        alpha[self.START_TAG] = 0.0
        
        for i in range(seq_len):
            emit = emissions[i]  # (num_labels,)
            
            # Expand for broadcasting
            # alpha: (num_labels+2,) -> (num_labels+2, 1)
            # transitions: (num_labels+2, num_labels+2)
            # emit: (num_labels,) -> (1, num_labels)
            
            alpha_expanded = alpha.unsqueeze(1)  # (L+2, 1)
            trans = self.transitions[:, :self.num_labels]  # (L+2, L)
            emit_expanded = emit.unsqueeze(0)  # (1, L)
            
            # Next alpha via log-sum-exp
            scores = alpha_expanded + trans + emit_expanded  # (L+2, L)
            alpha_new = torch.logsumexp(scores, dim=0)  # (L,)
            
            # Expand back to include START and END
            alpha = torch.full((self.num_labels + 2,), -10000.0, device=emissions.device)
            alpha[:self.num_labels] = alpha_new
        
        # Terminal: add transition to END
        terminal = alpha + self.transitions[:, self.END_TAG]
        log_Z = torch.logsumexp(terminal, dim=0)
        
        return log_Z
    
    def _score_sentence(
        self, 
        emissions: torch.Tensor, 
        tags: torch.Tensor
    ) -> torch.Tensor:
        """
        Compute score of a specific tag sequence.
        
        Args:
            emissions: Emission scores (seq_len, num_labels)
            tags: Gold tag sequence (seq_len,)
        
        Returns:
            score: Score of this tag sequence (scalar)
        """
        seq_len = emissions.size(0)
        
        score = self.transitions[self.START_TAG, tags[0]]
        score += emissions[0, tags[0]]
        
        for i in range(1, seq_len):
            score += self.transitions[tags[i-1], tags[i]]
            score += emissions[i, tags[i]]
        
        score += self.transitions[tags[-1], self.END_TAG]
        
        return score
    
    def neg_log_likelihood(
        self, 
        sentence: torch.Tensor, 
        tags: torch.Tensor
    ) -> torch.Tensor:
        """
        Compute negative log-likelihood for training.
        
        Args:
            sentence: Word indices (seq_len,)
            tags: Gold tag indices (seq_len,)
        
        Returns:
            loss: Negative log-likelihood (scalar)
        """
        # Get emissions from BiLSTM
        embeds = self.embedding(sentence)
        lstm_out, _ = self.lstm(embeds.unsqueeze(0))
        lstm_out = lstm_out.squeeze(0)  # (seq_len, hidden_dim)
        emissions = self.emission(lstm_out)  # (seq_len, num_labels)
        
        # Compute forward score (log Z)
        forward_score = self._forward_alg(emissions)
        
        # Compute gold score
        gold_score = self._score_sentence(emissions, tags)
        
        # NLL = log Z - gold_score
        return forward_score - gold_score
    
    def viterbi_decode(self, sentence: torch.Tensor) -> List[int]:
        """
        Find best tag sequence using Viterbi.
        
        Args:
            sentence: Word indices (seq_len,)
        
        Returns:
            best_path: List of tag indices
        """
        # Get emissions
        embeds = self.embedding(sentence)
        lstm_out, _ = self.lstm(embeds.unsqueeze(0))
        lstm_out = lstm_out.squeeze(0)
        emissions = self.emission(lstm_out)
        
        seq_len = emissions.size(0)
        
        # Viterbi algorithm
        backpointers = []
        
        # Initialize
        viterbi = torch.full((self.num_labels + 2,), -10000.0)
        viterbi[self.START_TAG] = 0.0
        
        for i in range(seq_len):
            emit = emissions[i]
            viterbi_expanded = viterbi.unsqueeze(1)
            trans = self.transitions[:, :self.num_labels]
            emit_expanded = emit.unsqueeze(0)
            
            scores = viterbi_expanded + trans + emit_expanded
            best_prev = torch.argmax(scores, dim=0)
            viterbi_new = scores[best_prev, torch.arange(self.num_labels)]
            
            backpointers.append(best_prev[:self.num_labels])
            
            viterbi = torch.full((self.num_labels + 2,), -10000.0)
            viterbi[:self.num_labels] = viterbi_new
        
        # Terminal
        terminal_scores = viterbi + self.transitions[:, self.END_TAG]
        best_last = torch.argmax(terminal_scores).item()
        
        # Backtrack
        best_path = [best_last]
        for bp in reversed(backpointers):
            best_path.append(bp[best_path[-1]].item())
        best_path.reverse()
        best_path = best_path[:-1]  # Remove START artifact
        
        return best_path
 
 
# Example usage
print("BiLSTM-CRF Architecture")
print("=" * 50)
print("Key components:")
print("  1. Word embeddings → BiLSTM → hidden states")
print("  2. Hidden states → Linear → emission scores")
print("  3. Emission scores + transitions → CRF score")
print("  4. Training: minimize -log P(y|x)")
print("  5. Inference: Viterbi decoding")