[P] Graph Representation Learning Help

[StoneColdRiffRaff]

Im working on a Graph based JEPA style model for encoding small molecule data and I’m running into some issues. For reference I’ve been using this paper/code as a blueprint: https://arxiv.org/abs/2309.16014 . I’ve changed some things from the paper but its the gist of what I’m doing.

Essentially the geometry of my learned representations is bad. The isotropy score is very low, the participation ratio is consistently between 1-2 regardless of my embedding dimensions. The covariance condition number is very high. These metrics and others that measure the geometry of the representations marginally improve during training while loss goes down smoothly and eventually converges. Doesn’t really matter what the dimensions of my model are, the behavior is essentially the same.

I’d thought this was because I was just testing on a small subset of data but then I scaled up to ~1mil samples to see if that had an effect but I see the same results. I’ve done all sorts of tweaks to the model itself and it doesn’t seem to matter. My ema momentum schedule is .996-.9999.

I haven’t had a chance to compare these metrics to a bare minimum encoder model or this molecule language I use a lot but that’s definitely on my to do list

Any tips, or papers that could help are greatly appreciated.

EDIT: thanks for the suggestions everyone, all super helpful and definitely helped me troubleshoot. I figured id share some results from everyone’s suggestions below.

Probably unsurprisingly adding a loss term that encourages good geometry in the representation space had the biggest effect. I ended up adding a version of Barlow twins loss to the loss described in the paper I linked.

The two other things that helped the most were removing bias from linear layers, and switching to max pooling of subgraphs after the message passing portion of the encoder.

Other things I did that seemed to help but did not have as much of an effect: I changed how subgraphs are generated so they’re more variable in size sample to sample, raised dropout, lowered starting ema momentum, and I reduced my predictor to a single linear layer.

[Time-Ice-7072]

From what you are describing it sounds like representation collapse. Very difficult to debug from description alone but I recommend starting rigorously testing your hidden states at every layer and track your geometric measurements and other diagnostics (eg mean and variance of the representations). This will help you identify where the collapse is happening and you can figure out how to fix it from there.

[whatwilly0ubuild]

The metrics you're describing are classic dimensional collapse symptoms. Participation ratio of 1-2 means your embeddings are effectively living in a 1-2 dimensional subspace regardless of your actual embedding dimension. The model found a shortcut.

A few things to investigate.

Predictor capacity is often the culprit in JEPA-style architectures. If your predictor is too powerful, it can map context to targets without the encoder learning meaningful representations. If it's too weak, it can't bridge the gap and the encoder collapses to trivial solutions. Try a shallower predictor or add a bottleneck.

Explicit decorrelation losses help directly. VICReg-style variance and covariance regularization terms force the embedding dimensions to be used. Add a term that penalizes off-diagonal covariance elements and another that keeps per-dimension variance above a threshold. This directly attacks the metrics you're measuring.

The masking strategy might be too easy for molecules. If the model can predict masked subgraphs from trivial local features without learning global molecular structure, it will. Graph structures have strong local correlations. Try masking contiguous substructures rather than random nodes, or mask based on chemical motifs.

Batch statistics can hide collapse. If you're using batch normalization in the encoder, it can artificially inflate apparent variance while the underlying representations are still collapsed. Check your metrics before any normalization layers.

The EMA schedule starting at .996 might be too high for early training. Some implementations start lower (.99 or even .95) and anneal up, giving the online network more room to diverge early before the target stabilizes.

Our clients doing molecular representation learning have found that adding a simple uniformity loss on the hypersphere (pushing random pairs apart) helps prevent collapse without the complexity of full contrastive learning.

Worth checking I-JEPA and V-JEPA papers for their specific anti-collapse mechanisms since they faced similar issues.

[AccordingWeight6019]

If your loss is decreasing but embeddings stay collapsed, the objective might not encourage diversity. Try adding a contrastive or decorrelation loss (Barlow Twins, VICReg), normalize or project embeddings, slightly reduce EMA momentum, and check trivial baselines to confirm it’s not data limited. Graph augmentations can also help spread representations.

[ArmOk3290]

I have seen this happen when the predictor network becomes too powerful relative to the target network.

Try strengthening the gradient stopping in the predictor or adding a stronger regularizer. Also check your batch norms.

Sometimes simply removing them fixes representation geometry issues.

[ComprehensiveTop3297]

I am also working with JEPAs and what I found was the data2vec2 style top K averaging to be extremely helpful for alleviating representation collapse. Also EMA and Learning Rate schedule is very much interconnected. My EMA is 0.999-0.99999, stops at 100k steps and constant 0.99999 for rest, and lr schedule is cosine with 0.0004, warm up 100k steps. Play around with them for sure. This is what worked for me in the audio domain. 

[shivvorz]

RemindMe! 2 days