| """Vision connector: SigLIP2 NaFlex patch features [B,576,768] -> soft tokens [B,144 (or 64),1536] for the frozen Jais-3-0.9B. |
| pixel-shuffle 2x2 over the 24x24 grid (576x768 -> 144x3072) -> Linear(3072,1536) -> GELU -> Linear(1536,1536) -> per-token L2 normalise * learnable scalar |
| (initialised to the Jais embedding mean row norm 0.896, so soft tokens start at the scale of real text embeddings). 7.08M parameters. |
| low-token mode (64 tokens): the 12x12 shuffled grid is adaptively average-pooled to 8x8 (same MLP, so one connector serves both modes).""" |
| import torch, torch.nn as nn, torch.nn.functional as F |
|
|
| EMBED_NORM = 0.896 |
|
|
|
|
| def pixel_shuffle(x, grid=24, r=2): |
| """[B, grid*grid, D] -> [B, (grid/r)^2, D*r*r]; each output token concatenates a rxr block of neighbouring patches""" |
| B, N, D = x.shape |
| assert N == grid * grid, f"expected {grid*grid} patches, got {N}" |
| g = grid // r |
| x = x.view(B, g, r, g, r, D).permute(0, 1, 3, 2, 4, 5) |
| return x.reshape(B, g * g, r * r * D) |
|
|
|
|
| class Connector(nn.Module): |
| def __init__(self, in_dim=768, hidden=1536, out_dim=1536, grid=24, r=2, target_norm=EMBED_NORM, low_grid=8): |
| super().__init__() |
| self.grid, self.r, self.low_grid = grid, r, low_grid |
| self.fc1 = nn.Linear(in_dim * r * r, hidden) |
| self.fc2 = nn.Linear(hidden, out_dim) |
| self.scale = nn.Parameter(torch.tensor(float(target_norm))) |
|
|
| def tokens(self, feats, low=False): |
| """feats [B,576,768] -> shuffled tokens [B,144,3072] (or [B,64,3072] in low mode)""" |
| x = pixel_shuffle(feats, self.grid, self.r) |
| if low: |
| g = self.grid // self.r |
| x = x.transpose(1, 2).reshape(x.shape[0], -1, g, g) |
| x = F.adaptive_avg_pool2d(x, self.low_grid).flatten(2).transpose(1, 2) |
| return x |
|
|
| def forward(self, feats, low=False): |
| x = self.tokens(feats.to(self.fc1.weight.dtype), low) |
| y = self.fc2(F.gelu(self.fc1(x))) |
| |
| return (F.normalize(y.float(), dim=-1) * self.scale.float()).to(y.dtype) |
|
|
| |
| @torch.no_grad() |
| def portal_init(self, feats_sample, embed_matrix, k=512, min_norm=0.3): |
| """Bures-Wasserstein (Gaussian optimal transport) initialisation after PORTAL (arXiv 2610.18598), as summarised in docs/research/vision_design.md s.3: |
| W1 = V_t diag(sigma_t/sigma_s) V_s^T, b1 = mu_t - W1 mu_s, W2 = identity. Source = shuffled patch tokens (3072-d) of an unlabeled image pool |
| (feats_sample [N,576,768], 1-2k images are enough), target = LLM input-embedding rows (rows with norm < min_norm = reserved/unused are dropped). |
| CAVEATS: written from the design-doc description, NOT from the paper's Appendix F code, and not validated against it; the GELU after layer 1 is kept |
| (as in the summary), so layer-1 output is only approximately in the embedding distribution. The 2-layer MLP has hidden=out=1536 so W2 = I fits. |
| Use: conn.portal_init(feats[:1024].float(), model.get_input_embeddings().weight)""" |
| dev = self.fc1.weight.device |
| src = self.tokens(feats_sample.float().to(dev)).reshape(-1, self.fc1.in_features) |
| tgt = embed_matrix.float().to(dev); tgt = tgt[tgt.norm(dim=1) > min_norm] |
|
|
| def stats(X, k): |
| mu = X.mean(0); Xc = X - mu |
| k = min(k, X.shape[1], X.shape[0] - 1) |
| U, S, V = torch.pca_lowrank(Xc, q=k, center=False, niter=3) |
| sigma = (Xc @ V).std(0) |
| return mu, V, sigma |
| mu_s, Vs, sg_s = stats(src, k); mu_t, Vt, sg_t = stats(tgt, k) |
| kk = min(Vs.shape[1], Vt.shape[1]) |
| W1 = Vt[:, :kk] @ torch.diag(sg_t[:kk] / sg_s[:kk].clamp_min(1e-6)) @ Vs[:, :kk].T |
| self.fc1.weight.copy_(W1); self.fc1.bias.copy_(mu_t - W1 @ mu_s) |
| self.fc2.weight.copy_(torch.eye(self.fc2.out_features, self.fc2.in_features)); self.fc2.bias.zero_() |
| return {"k": kk, "src_tokens": src.shape[0], "tgt_rows": tgt.shape[0]} |
|
|
|
|
| def n_params(m): return sum(p.numel() for p in m.parameters()) |
|
|
|
|
| if __name__ == "__main__": |
| c = Connector(); print("params", n_params(c)) |
| x = torch.randn(2, 576, 768); print(c(x).shape, c(x, low=True).shape, c(x).norm(dim=-1).mean().item()) |
|
|
