Spaces:
Running
Running
| import torch | |
| import torch.nn as nn | |
| import geoopt | |
| class HyperbolicAdapter(nn.Module): | |
| """ | |
| A lightweight adapter that maps Euclidean embeddings (e.g., from CLIP/ResNet) | |
| into the Poincaré ball. | |
| This solves the "Representation Collapse" by projecting crowded Euclidean | |
| vectors into a space with exponentially growing volume. | |
| """ | |
| def __init__(self, input_dim=512, output_dim=128, c=1.0): | |
| super().__init__() | |
| self.c = c # Curvature of the manifold | |
| self.manifold = geoopt.PoincareBall(c=self.c) | |
| # 1. Linear Layer (Euclidean transformation) | |
| # We first compress/transform the features in Euclidean space | |
| self.encoder = nn.Linear(input_dim, output_dim) | |
| # 2. Mobius Linear (Hyperbolic transformation - Optional but powerful) | |
| # For a simple adapter, we can just project the output of the linear layer. | |
| # But a true hyperbolic layer operates on the manifold. | |
| # self.hyp_layer = geoopt.MobiusLinear(output_dim, output_dim, c=self.c) | |
| def forward(self, x): | |
| """ | |
| Args: | |
| x: Euclidean input tensor (Batch, Input_Dim) | |
| Returns: | |
| x_hyp: Hyperbolic output tensor (Batch, Output_Dim) on the Poincaré ball | |
| """ | |
| # Step 1: Standard Euclidean processing | |
| x_euc = self.encoder(x) | |
| # Step 2: Exponential Map (The "Bridge") | |
| # This projects the Euclidean vector onto the Hyperbolic manifold. | |
| # Ideally, we map to the tangent space at the origin (which is Euclidean-like) | |
| # and then use expmap0 to project to the manifold. | |
| # expmap0(v) maps a vector v in the tangent space at 0 to the manifold. | |
| x_hyp = self.manifold.expmap0(x_euc) | |
| return x_hyp | |
| def demo(): | |
| print("--- Hyperbolic Adapter Demo ---") | |
| # Simulate CLIP embeddings (Batch=32, Dim=512) | |
| # These represent our "crowded" Euclidean data | |
| batch_size = 32 | |
| input_dim = 512 | |
| clip_embeddings = torch.randn(batch_size, input_dim) | |
| print(f"Input (Euclidean): {clip_embeddings.shape}") | |
| # Initialize Adapter | |
| # We project to a lower dimension (128) which is common for hyperbolic spaces | |
| # as they are more efficient (can capture hierarchy in fewer dimensions). | |
| adapter = HyperbolicAdapter(input_dim=input_dim, output_dim=128) | |
| # Forward Pass | |
| hyp_embeddings = adapter(clip_embeddings) | |
| print(f"Output (Hyperbolic): {hyp_embeddings.shape}") | |
| # Verify they are on the manifold | |
| # In Poincaré ball, norm must be < 1/sqrt(c) (usually < 1) | |
| norms = hyp_embeddings.norm(dim=-1) | |
| max_norm = norms.max().item() | |
| print(f"Max Norm of embeddings: {max_norm:.4f} (Should be < 1.0)") | |
| if max_norm < 1.0: | |
| print("SUCCESS: All points are validly projected onto the Poincaré disk.") | |
| else: | |
| print("FAILURE: Points escaped the manifold!") | |
| if __name__ == "__main__": | |
| demo() | |