Spaces:
Running on Zero
Running on Zero
Premchandyadav369
feat: deploy interactive Q-TensorFormer demo with trained weights and Gradio 4 UI
67ded87 | """ | |
| Adaptive TT-Rank Scheduler. | |
| Core novelty of Q-TensorFormer: adjusts tensor rank dynamically | |
| based on per-input complexity, estimated via attention entropy. | |
| r(input) = r_min + α × normalized_entropy × (r_max - r_min) | |
| Supports: | |
| - EMA smoothing to prevent oscillation | |
| - Budget-capped ranks | |
| - Deterministic rounding with hysteresis | |
| """ | |
| import torch | |
| import torch.nn as nn | |
| import math | |
| class RankScheduler(nn.Module): | |
| """ | |
| Attention entropy → TT-rank scheduler. | |
| Parameters | |
| ---------- | |
| r_min : int | |
| Minimum tensor rank (maximum compression). | |
| r_max : int | |
| Maximum tensor rank (minimum compression). | |
| alpha : float | |
| Sensitivity: how much entropy changes the rank. | |
| alpha=0 → fixed rank r_min. | |
| alpha=1 → rank fully spans r_min to r_max. | |
| alpha=2.0 → aggressive scaling (default). | |
| smoothing : float | |
| EMA decay factor (0.9 = smooth, 0 = no history). | |
| """ | |
| def __init__(self, r_min: int = 2, r_max: int = 8, | |
| alpha: float = 2.0, smoothing: float = 0.9): | |
| super().__init__() | |
| self.r_min = r_min | |
| self.r_max = r_max | |
| self.alpha = alpha | |
| self.smoothing = smoothing | |
| self.register_buffer("_ema_entropy", torch.tensor(0.5)) | |
| self.register_buffer("_ema_rank", torch.tensor((r_min + r_max) // 2, dtype=torch.float)) | |
| self.register_buffer("_counter", torch.tensor(0, dtype=torch.long)) | |
| # Optionally learn alpha | |
| self.learned_alpha = nn.Parameter(torch.tensor(float(alpha)), requires_grad=False) | |
| def forward(self, entropy: torch.Tensor, seq_len: int = None) -> int: | |
| """ | |
| Compute rank from attention entropy. | |
| Args: | |
| entropy: Scalar or 0-dim tensor (mean attention entropy). | |
| seq_len: Sequence length for normalization (optional). | |
| Returns: | |
| Integer tensor rank. | |
| """ | |
| if entropy.dim() > 0: | |
| entropy = entropy.mean() | |
| # Normalize entropy to [0, 1] | |
| if seq_len is not None and seq_len > 1: | |
| norm_factor = math.log(seq_len) | |
| normalized = torch.clamp(entropy / max(norm_factor, 1e-8), 0.0, 1.0) | |
| else: | |
| normalized = torch.clamp(torch.tanh(entropy / 2.0), 0.0, 1.0) | |
| # EMA smoothing | |
| self._ema_entropy.mul_(self.smoothing).add_(normalized, alpha=1.0 - self.smoothing) | |
| smoothed = self._ema_entropy | |
| # Map to rank: r = r_min + alpha * norm * (r_max - r_min) | |
| alpha_val = self.learned_alpha.item() | |
| span = self.r_max - self.r_min | |
| raw = self.r_min + alpha_val * smoothed.item() * span | |
| # Round with hysteresis | |
| self._ema_rank.mul_(0.7).add_(raw, alpha=0.3) | |
| rank = int(torch.round(self._ema_rank).item()) | |
| # Clamp + counter | |
| rank = max(self.r_min, min(self.r_max, rank)) | |
| self._counter.add_(1) | |
| return rank | |
| def reset(self): | |
| """Reset EMA state.""" | |
| self._ema_entropy.fill_(0.5) | |
| self._ema_rank.fill_((self.r_min + self.r_max) / 2.0) | |
| self._counter.fill_(0) | |
| def current_rank(self) -> float: | |
| return self._ema_rank.item() | |
| def current_entropy(self) -> float: | |
| return self._ema_entropy.item() | |
| class BudgetAwareScheduler(nn.Module): | |
| """ | |
| Extends RankScheduler with deployment budget constraints. | |
| Automatically caps tensor rank to meet: | |
| - Max parameter budget | |
| - Max latency target | |
| - Max energy per query | |
| """ | |
| def __init__(self, scheduler: RankScheduler, | |
| max_params: int = None, | |
| max_latency_ms: float = None, | |
| max_energy_uj: float = None): | |
| super().__init__() | |
| self.scheduler = scheduler | |
| self.max_params = max_params | |
| self.max_latency_ms = max_latency_ms | |
| self.max_energy_uj = max_energy_uj | |
| def forward(self, entropy: torch.Tensor, seq_len: int = None, | |
| param_factors: dict = None) -> int: | |
| """ | |
| Compute rank with budget constraints. | |
| Args: | |
| entropy: Attention entropy. | |
| seq_len: Sequence length. | |
| param_factors: Dict mapping rank → estimated total parameters. | |
| Returns: | |
| Budget-constrained rank. | |
| """ | |
| rank = self.scheduler(entropy, seq_len) | |
| if param_factors and self.max_params: | |
| # Find highest rank that meets budget | |
| while rank > self.scheduler.r_min: | |
| est = param_factors.get(rank, float("inf")) | |
| if est <= self.max_params: | |
| break | |
| rank -= 1 | |
| return rank | |
| def reset(self): | |
| self.scheduler.reset() | |