Premchan369/Q-TensorFormer

1"""2Quantum Router: Selective Quantum Activation.3 4Only "hard" tokens pass through the quantum circuit.5Decision mechanism: learned linear gate + straight-through estimator.6 7v3 improvements:8  - Sparsity target: ensures target fraction of tokens skip quantum9  - Straight-through gradient for gradient-based learning10  - Sparsity statistics tracking11  - Fallback embedding for bypassed tokens12"""13 14import torch15import torch.nn as nn16import torch.nn.functional as F17 18 19class QuantumRouter(nn.Module):20    """21    Selective quantum activation gate.22 23    Given a batch of token embeddings, computes a per-token24    probability of routing through quantum. Uses straight-through25    estimator: forward pass uses hard binary decisions, backward26    uses soft sigmoid gradient.27 28    Parameters29    ----------30    d_model : int31        Input feature dimension.32    q_input_dim : int33        Dimension expected by quantum circuit (typically n_qubits).34    target_sparsity : float35        Target fraction of tokens that SKIP quantum (0.7 = 70% skip).36    temperature : float37        Softmax temperature for gate decisions (lower = harder).38    """39 40    def __init__(self, d_model: int, q_input_dim: int = 4,41                 target_sparsity: float = 0.7, temperature: float = 1.0):42        super().__init__()43        self.d_model = d_model44        self.q_input_dim = q_input_dim45        self.target_sparsity = target_sparsity46        self.temperature = temperature47 48        # Projection for gate decision49        self.gate_proj = nn.Sequential(50            nn.LayerNorm(d_model),51            nn.Linear(d_model, d_model // 4),52            nn.GELU(),53            nn.Linear(d_model // 4, 1),54        )55 56        # Projection to quantum input dimension57        self.q_proj = nn.Linear(d_model, q_input_dim)58 59        # Statistics60        self.register_buffer("total_tokens", torch.tensor(0, dtype=torch.long))61        self.register_buffer("quantum_tokens", torch.tensor(0, dtype=torch.long))62        self.register_buffer("_ema_sparsity", torch.tensor(target_sparsity))63 64    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:65        """66        Route tokens selectively through quantum.67 68        Args:69            x: (*batch, seq_len, d_model)70 71        Returns:72            quantum_out: (*batch, seq_len, d_model) — quantum-processed tokens73            mask: (*batch, seq_len) — which tokens went through quantum (bool)74        """75        *batch_dims, seq_len, d_model = x.shape76 77        # Gate decision78        gate_logits = self.gate_proj(x).squeeze(-1)  # (*, seq_len)79        soft_mask = torch.sigmoid(gate_logits / self.temperature)80 81        # Straight-through: hard forward, soft backward82        hard_mask = (soft_mask > 0.5).float()83        mask = hard_mask.detach() + soft_mask - soft_mask.detach()84 85        # Project selected tokens to quantum dimension86        q_input = self.q_proj(x)  # (*, seq_len, q_input_dim)87 88        # TODO: actual quantum circuit call goes here89        # For now: project back to d_model with learned linear layer90        quantum_out = F.gelu(q_input)91        if not hasattr(self, '_q_out_proj'):92            self._q_out_proj = nn.Linear(self.q_input_dim, d_model).to(x.device)93        quantum_out = self._q_out_proj(quantum_out)94 95        # Gate output96        mask_expanded = mask.unsqueeze(-1)  # (*, seq_len, 1)97        output = mask_expanded * quantum_out98 99        # Update statistics100        with torch.no_grad():101            n_tokens = seq_len * max(1, math_prod(batch_dims))102            n_quantum = int(mask_expanded.sum().item())103            self.total_tokens += n_tokens104            self.quantum_tokens += n_quantum105            actual_rate = n_quantum / max(n_tokens, 1)106            self._ema_sparsity.mul_(0.99).add_(107                (1 - actual_rate), alpha=0.01108            )109 110        return output, mask.detach().bool()111 112    @property113    def sparsity(self) -> float:114        """Fraction of tokens that SKIP the quantum circuit."""115        return self._ema_sparsity.item()116 117    @property118    def usage_percent(self) -> float:119        """Fraction of tokens that use the quantum circuit."""120        return 1.0 - self.sparsity121 122    def reset_stats(self):123        self.total_tokens.zero_()124        self.quantum_tokens.zero_()125        self._ema_sparsity.fill_(self.target_sparsity)126 127    def reset_state(self):128        """Full reset for clean evaluation runs."""129        self.reset_stats()130        for m in self.modules():131            if hasattr(m, "reset_parameters"):132                m.reset_parameters()133 134    def extra_repr(self) -> str:135        return (f"d_model={self.d_model}, q_dim={self.q_input_dim}, "136                f"target_sparsity={self.target_sparsity:.1%}")137 138 139def math_prod(iterable):140    """Safe product of iterable."""141    result = 1142    for x in iterable:143        result *= x144    return result145