Is RA-SAE able to process embedding features from feed-forward 3D vision backbones?

[Zifeng-Zhang]

Hi, thanks on your great work on vision model interpretations! I also love your work and findings in the latest "Into the Rabbit Hull" paper!

Yet I'm wondring if it's do-able to process embeddings from feed-forward 3D backbones (e.g. mast3r family, vggt etc.) which are also attention-based models, and it's supre interesting to reveal what 3D concepts do they learned. Is there anything to pay attention to when swtich to 3D vision models?

Meanwhile, I'll trying to train a RA-SAE on VGGT's embeddings (the attention layer output) and I'm glad to share it if anyone find it interesting.

[fel-thomas]

Hey, :)

Wow, that’s an awesome idea and thanks a lot for the kind words!

I was actually planning to open-source the SAE part on Hugging Face soon, so that would fit perfectly. We should totally try to put out a small notebook if possible, I’d add it with pleasure :)

Running it on 3D backbones like MAST3R or VGGT should be fine, just watch how you handle the spatial tokens or feature shapes before feeding them to the SAE. Otherwise, same logic applies. And training a RA-SAE on VGGT embeddings sounds super cool — please share it when you can, I’d love to see what kind of 3D concepts show up!

Cheers,
Thomas

[SanBast]

Hello @Zifeng-Zhang and @fel-thomas !

Joining the conversation as 1) I'd like to hugely thank you @fel-thomas for this amazing library and work (I'm no1 fan of learned representations and manifold disentanglement, and your work is on point 😄 ), and 2) I'm also interested in using RA-SAE for my 3D manifold learning and I was wondering if @Zifeng-Zhang did you manage to make the full machine going?

Unfortunately, I can't find the RA-SAE module yet :/ I just found the RA-based dictionary learning module. I don't want to put pressure @fel-thomas, but please let me know when you upload a demo for it! (I suppose the RA-SAE Dictionary is a plug-and-play module for TopKSAE for example, but I'd' like to be sure to not break anything ahah)

Thanks in advance to both!
Best,
Vincenzo

[Zifeng-Zhang]

Hi @SanBast, and thanks @fel-thomas for your kind reply!

Sorry for my late answerings, I was trying to train a SAE on VGGT the last week yet my experiments showed negative findings that might suggest SAE is not a good fit for interpreting VGGT's (not sure for other 3D backbones) latent vector, or maybe I missed something that did not work out the SAE training on these embeddings (I still need a lot of work to do ablations). Here I post what I've done and some primary (no gaurantee) findings.

From what I understand, SAE takes latent representations as input and learned a sparse dictionary that could intreprete that representation. So the first step is to build a dataset of VGGT's latent representation.

According to original VGGT, it has an aggregator module before feeding the latent into decoders (camera and DPT heads). The aggregator consists of a DINOv2 style image encoder for PatchEmbed (originally, it's trained from scracth, not using DINOv2's pretrained weights), and an alternating attention block ( the patched tokens are feed into both frame/global attention, repeating 24 layers). The DPT head only takes the output representation from attention layer_id [4, 11, 17, 23] as input and form into a feature pyrimid way to fuse different hierarchies of feature.

As my naive approach, I sampled 16 views across 1.2k ScanNetv2 indoor scenes (averagely each scenes contains around 100 views), and feed 16 views in one batch to VGGT, and ONLY extract the feature form the last attention (23rd) layer (where feature_dim = 2048, concat from frame attn d_f=1024 and global attn d_g=1024) in one forward run. Each scenes with 16 views would output a latent and I stored them into a pt tensor file. Then I built a DataLoader to feed these latent into a BatchTopKSAE for training where topk = 64. I used k_means to extract candidate points.

I trained three times where the only change is the n_concept (10k, 1k, 100) using adamW and a cosine scheduler for lr. Tensorboard is shown below:

All settings have the same bad results where dead concept ratio is almost close to 100%, implying that no meaningful concepts are learned via SAE training.

Part II below

[Zifeng-Zhang]

Part II

I also did a simple PCA analysis on a subset of the extracted VGGT's feature tensor:

Tensor Norm Statistics:
mean: 82.7441
std: 11.0032
min: 57.3918
max: 132.4062

inner sparsity（|x|<0.01）: 0.68%

Dimension STD Statistics:
mean: 1.1663
min: 0.3835
max: 25.8120
dead dimension（std<0.01）: 0/2048

Calculate Token Similarity (sub sample 1000 items)...
Token Cosine Similarity:
mean: 0.9105
median: 0.9196
std: 0.0536

PCA Analysis...
top 10 principle component explained var ratio:
PC1: 30.09%
PC2: 27.46%
PC3: 10.37%
PC4: 4.15%
PC5: 3.63%
PC6: 2.33%
PC7: 1.80%
PC8: 1.72%
PC9: 1.26%
PC10: 1.20%
To explain 95% std needs: 34 amount of principle component

The above findings might indicate that the VGGT learned only a sparse combination of "concepts", this is also being pointed out from the FastVGGT paper as feature "degradation", where they stated that:

In DINO, while the CLS token becomes increasingly discriminative, patch-level features gradually lose local consistency as they converge toward the CLS token, which undermines performance on dense prediction tasks.

A similar limitation appears in VGGT: its global self-attention layers aggregate information across all tokens via weighted averaging, and without explicit regularization or task-specific constraints, the representational space is repeatedly compressed. This compression drives tokens to collapse along a dominant direction, eroding their distinctiveness.

Below is the visualizations of the Global Attention maps in VGGT, showing that at every stage the attention patterns of different tokens exhibit a strong degree of similarity, referenced from FastVGGT:

Based on the above observation -- if the training has no mistake, I need further ablations to double check -- I personally conclude:

1. Global attention mechanism force tokens to converge into a unified representation (tokens are highly similar)
2. 3D geometric constraints maybe easier than 2D vision-language semantics like DINO, CLIP etc.

Yet I'm still confused on:

1. why the dead concept ratio is always close to 100%, is it normal during SAE training?
2. need further experiements to show if 16 views are too small to capture dense 3D constraints (dataset is not large enough), or extracting different layer other than 23rd layer could make any difference.
3. since each pt file is obtained from 16 views per scene, how could I visualize the heatmaps back to those images?

I'm open to any thoughts and questions from yall both!

cheers!