1import torch
 2from torch import nn
 3import torch.nn.functional as F
 4
 5
 6class GeometricallyAttractingRandomRNN(nn.Module):
 7    """RNN with K candidate tanh maps and a time-dependent categorical gate.
 8
 9    Forward uses either a sampled candidate (stochastic=True) or the exact
10    probability-weighted output mixture. The latter is useful for low-variance
11    training. Candidate spectral norms provide a cheap Jacobian upper bound.
12    """
13    def __init__(self, input_size, hidden_size, candidates=2, target_rho=0.95,
14                 gate_input_size=None):
15        super().__init__()
16        self.input_size = input_size
17        self.hidden_size = hidden_size
18        self.candidates = candidates
19        self.target_rho = target_rho
20        self.W = nn.Parameter(torch.empty(candidates, hidden_size, hidden_size))
21        self.U = nn.Parameter(torch.empty(candidates, hidden_size, input_size))
22        self.bias = nn.Parameter(torch.zeros(candidates, hidden_size))
23        nn.init.orthogonal_(self.W[0])
24        for k in range(1, candidates):
25            nn.init.orthogonal_(self.W[k])
26        nn.init.xavier_uniform_(self.U)
27        gate_input_size = input_size if gate_input_size is None else gate_input_size
28        self.gate = nn.Linear(gate_input_size, candidates)
29
30    def gains(self):
31        # Exact matrix-norm upper bound for tanh Jacobian (which is <= 1).
32        return torch.linalg.matrix_norm(self.W, ord=2, dim=(-2, -1))
33
34    def contraction_penalty(self, probabilities, eps=1e-8):
35        """Squared positive excess of log expected gain over log(target_rho)."""
36        expected_gain = (probabilities * self.gains()).sum(dim=-1)
37        excess = torch.log(expected_gain + eps) - torch.log(
38            torch.as_tensor(self.target_rho, device=expected_gain.device))
39        return F.relu(excess).square().mean()
40
41    def step(self, x, h, gate_features=None, stochastic=True):
42        # x: [B,input], h: [B,hidden], gate_features: [B,gate_input].
43        z = x if gate_features is None else gate_features
44        probabilities = torch.softmax(self.gate(z), dim=-1)
45        candidates = torch.tanh(
46            torch.einsum('kij,bj->bki', self.W, h)
47            + torch.einsum('kij,bj->bki', self.U, x)
48            + self.bias[None])
49        if stochastic:
50            index = torch.multinomial(probabilities, 1).squeeze(-1)
51            next_h = candidates[torch.arange(x.shape[0], device=x.device), index]
52        else:
53            next_h = (probabilities[..., None] * candidates).sum(dim=1)
54        return next_h, probabilities
55
56    def forward(self, x_sequence, h0=None, stochastic=True):
57        # x_sequence: [T,B,input]
58        T, B, _ = x_sequence.shape
59        h = x_sequence.new_zeros(B, self.hidden_size) if h0 is None else h0
60        states, routes = [], []
61        for t in range(T):
62            h, p = self.step(x_sequence[t], h, stochastic=stochastic)
63            states.append(h)
64            routes.append(p)
65        return torch.stack(states), torch.stack(routes)