HP-JEPA: Multi-Resolution Graph Learning Without Manual Tuning | Graph Neural Networks

Loading…

HP-JEPA: Multi-Resolution Graph Learning Without Manual Tuning | Graph Neural Networks | graphwiz.ai

Graph self-supervised learning aims to learn transferable representations from large-scale unlabeled graph data. Joint-Embedding Predictive Architectures (JEPAs) avoid explicit negative-pair construction and raw-input reconstruction by predicting masked targets directly in latent space.

However, existing graph JEPAs rely on a single predefined graph partition, biasing learned representations toward one structural granularity and limiting their ability to capture complementary patterns at different scales.

HP-JEPA (Hierarchical Partitioning for Multi-Resolution Graph JEPA) organises each graph into an ordered bank of coarse-to-fine partition resolutions and performs context-target latent prediction separately at each resolution.

The result: outperforms fixed-resolution Graph-JEPA on 6 of 8 tasks, with size-stratified analysis showing consistent improvement across graph-size quartiles.

The Multi-Resolution Problem

Graphs have structure at multiple scales:

Social Network Example:

Fine resolution (individual nodes):
  Alice ──► Bob ──► Charlie
   │   │     │   │     │
   ▼   ▼     ▼   ▼     ▼
  interests: ML, graphs, AI

Medium resolution (communities):
  [ML Researchers] ──► [Graph Databases] ──► [AI Systems]
        │                    │                    │
        └───────► [Overlap: Graph ML] ◄──────────┘

Coarse resolution (global structure):
  [Academic Network] ──► [Industry Network]
         │                    │
         └──────► [Collaboration Bridges] ◄────────┘

Single-resolution Graph-JEPA picks one partition:

Fine: Captures local patterns, misses global structure
Coarse: Captures global patterns, misses local detail
Manual tuning: Required to pick the "right" resolution

HP-JEPA learns at all resolutions simultaneously.

HP-JEPA Architecture

1. Hierarchical Partition Bank

class HierarchicalPartition:
    def __init__(self, graph, n_resolutions=5):
        self.graph = graph
        self.resolutions = []
        
        # Create coarse-to-fine partitions
        for level in range(n_resolutions):
            # Coarsen graph at this level
            coarsened = self._coarsen(graph, level)
            partitions = self._cluster(coarsened)
            self.resolutions.append({
                'level': level,
                'graph': coarsened,
                'partitions': partitions
            })
    
    def _coarsen(self, graph, level):
        # Graph coarsening (e.g., METIS, Louvain)
        # Higher level = coarser graph
        ...
    
    def _cluster(self, graph):
        # Community detection at this resolution
        ...

Key insight: Each resolution captures complementary structural patterns.

2. Multi-Resolution JEPA

┌─────────────────────────────────────────────────────────┐
│                    HP-JEPA                              │
│                                                          │
│  Input Graph                                             │
│      │                                                   │
│      ├─► Resolution 0 (coarse) ──► JEPA-0 ──┐          │
│      ├─► Resolution 1 ───────────► JEPA-1 ──┤          │
│      ├─► Resolution 2 ───────────► JEPA-2 ──┤  ┌─────┐ │
│      ├─► Resolution 3 ───────────► JEPA-3 ──┼─►│Concat││
│      └─► Resolution 4 (fine) ───► JEPA-4 ──┘  └──┬──┘ │
│                                                   │     │
│                                                   ▼     │
│                                            Final Embedding│
└─────────────────────────────────────────────────────────┘

Each JEPA branch:

Online encoder: Processes context at this resolution
Target encoder: EMA of online encoder (predicts masked targets)
Predictor: Maps context to masked target embeddings

3. Resolution-Specific Prediction

class MultiResolutionJEPA:
    def __init__(self, n_resolutions, partitions=None):
        self.n_resolutions = n_resolutions
        self.encoders = [JEPAEncoder() for _ in range(n_resolutions)]
        self.predictors = [JEAPPredictor() for _ in range(n_resolutions)]
        self.target_encoders = [EMAEncoder(enc) for enc in self.encoders]
        self.partitions = partitions  # HierarchicalPartition instance
    
    def forward(self, graph, mask):
        embeddings = []
        
        for level, (encoder, predictor, target) in enumerate(
            zip(self.encoders, self.predictors, self.target_encoders)
        ):
            # Get graph at this resolution
            graph_level = self.partitions[level].graph
            
            # Encode context
            context = encoder(graph_level, mask=mask[level])
            
            # Predict masked targets
            prediction = predictor(context)
            
            # Target (EMA encoder, no gradient)
            target = target(graph_level, mask=mask[level])
            
            # Loss: predict target from context
            loss = mse_loss(prediction, target)
            
            embeddings.append(context)
        
        # Concatenate all resolutions
        return torch.cat(embeddings, dim=-1)

4. Resolution Integration

Final representation can be:

Concatenation: All resolution embeddings concatenated
Task-specific weighting: Learn weights per task
Attention: Attend to relevant resolutions

class ResolutionIntegrator:
    def __init__(self, n_resolutions, embedding_dim):
        self.task_weights = torch.nn.Parameter(torch.ones(n_resolutions))
    
    def integrate(self, embeddings, task=None):
        if task is None:
            # Simple concatenation
            return torch.cat(embeddings, dim=-1)
        else:
            # Task-specific weighting
            weights = self.task_weights  # Learned per task
            weighted = [w * e for w, e in zip(weights, embeddings)]
            return torch.sum(torch.stack(weighted), dim=0)

What the Benchmarks Say

HP-JEPA was evaluated on seven graph classification benchmarks and one graph regression benchmark:

Benchmark	Task	Graph-JEPA (fixed)	HP-JEPA (multi)	Improvement
COLLAB	Scientific collaboration	0.78	0.82	+5.1%
REDDIT-Binary	Forum threads	0.84	0.89	+6.0%
PROTEINS	Protein structure	0.76	0.81	+6.6%
DD	Drug design	0.81	0.87	+7.4%
IMDB-BINARY	Movie collaboration	0.72	0.76	+5.6%
IMDB-MULTI	Movie collaboration	0.51	0.54	+5.9%
MNIST	Image graphs	0.97	0.98	+1.0%
ZINC	Molecular properties	0.48 MAE	0.43 MAE	-10.4% (lower better)

Key findings:

6 of 8 tasks: HP-JEPA outperforms fixed-resolution baseline
Largest gains: Tasks with multi-scale structure (PROTEINS, DD)
Smallest gains: Tasks with single-scale structure (MNIST)
Size-stratified: Consistent improvement across all graph-size quartiles

Production Considerations

Computational Cost

HP-JEPA adds multi-resolution overhead:

Component	Single JEPA	HP-JEPA (5 resolutions)	Overhead
Training time	1x	2.5x	+150%
Memory	1x	1.8x	+80%
Inference	1x	1.5x	+50%
Embedding size	512	2560 (5×512)	+400%

Mitigation:

Reduce resolutions: 3 resolutions instead of 5
Shared weights: Share encoder weights across resolutions
Pruning: Drop irrelevant resolutions at inference

Pre-training Strategy

# Pre-training pipeline
def pretrain_hp_jepa(graphs, n_epochs=100):
    model = HPJEPA(n_resolutions=5)
    optimizer = Adam(model.parameters(), lr=1e-4)
    
    for epoch in range(n_epochs):
        for graph in graphs:
            # Create hierarchical partitions
            partitions = HierarchicalPartition(graph, n_resolutions=5)
            
            # Create masks at each resolution
            masks = [create_mask(partitions.resolutions[i]['graph']) for i in range(5)]
            
            # Forward pass (returns per-resolution embeddings during training)
            per_resolution = model.forward_per_resolution(graph, masks)
            
            # Compute loss (sum across resolutions)
            loss = 0
            for i in range(5):
                target = model.target_encoders[i](partitions.resolutions[i]['graph'], masks[i])
                prediction = model.predictors[i](per_resolution[i])
                loss += mse_loss(prediction, target)
            
            # Backward pass
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
        
        # Update target encoders (EMA)
        for target, encoder in zip(model.target_encoders, model.encoders):
            target.update_ema(encoder)
    
    return model

Transfer Learning

# Fine-tuning on downstream task
def fine_tune(model, downstream_data, task='classification'):
    # Freeze all but last layer
    for param in model.encoders.parameters():
        param.requires_grad = False
    
    # Add task-specific head
    head = ClassificationHead(input_dim=2560, n_classes=10)
    
    # Fine-tune head + resolution weights
    optimizer = Adam([
        {'params': head.parameters(), 'lr': 1e-3},
        {'params': model.resolution_weights.parameters(), 'lr': 1e-4}
    ])
    
    for epoch in range(50):
        for graph, label in downstream_data:
            embedding = model(graph)
            prediction = head(embedding)
            loss = cross_entropy(prediction, label)
            
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
    
    return model, head

Where the Field Is Heading

HP-JEPA reveals three trends:

1. Multi-Scale Is Default

Single-resolution learning is the exception, not the rule. Multi-scale representations will become standard.

2. Self-Supervised Over Supervised

JEPAs learn from unlabelled data at scale. Supervised fine-tuning requires minimal labelled data.

3. Resolution-Aware Downstream Tasks

Downstream models will learn which resolutions matter for their specific task. Resolution weighting becomes a learned parameter.

The Bottom Line

HP-JEPA solves the multi-resolution problem by:

Organising graphs into coarse-to-fine partitions
Learning representations at each resolution separately
Integrating resolution-specific embeddings for downstream tasks

For graph representation learning, the implication is clear: multi-scale is necessary. Single-resolution learning misses complementary patterns.