File size: 4,594 Bytes
e3f7919 | 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 | """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 # mean row norm of Jais embed_tokens (design doc, section 2.2)
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))) # learnable output 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)))
# normalise in fp32 for stability under bf16 autocast, then rescale to the embedding norm
return (F.normalize(y.float(), dim=-1) * self.scale.float()).to(y.dtype)
# ---------------------------------------------------------------- PORTAL-style warm start (stage-1 ablation; default OFF)
@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) # per-axis spread along each principal direction
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 # [1536, 3072]
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())
|