Foundry: Distilling 3D Foundation Models for the Edge

AI summary

Overview

• Research area: Computer vision / 3D representation learning — specifically knowledge distillation and model compression for 3D point cloud foundation models.
• Technical level: Advanced. The paper assumes familiarity with Transformer architectures, self-supervised learning, attention mechanisms, and knowledge distillation.
• Scope: The paper proposes Foundation Model Distillation (FMD), a new compression paradigm, and introduces Foundry, its first implementation for 3D point cloud Transformers, which distills a large frozen Point-JEPA teacher into a compact student that stays a general-purpose feature extractor.

What This Paper Is About

Large self-supervised 3D foundation models are powerful general-purpose feature extractors, but their size and quadratic attention cost make them impossible to run on robots, AR/VR headsets, and other resource-constrained devices. Standard knowledge distillation can shrink them, but it produces narrow "specialist" models that lose the downstream-agnostic generality of the original. Foundry addresses this by compressing the teacher's entire latent representation space into a small set of learnable "SuperTokens" that can be decompressed back into the teacher's token-level embeddings, yielding one small student that transfers across classification, part segmentation, and few-shot tasks.

Key Contributions

1. Foundation Model Distillation (FMD): A new paradigm that distills the full representation space of a large self-supervised unimodal model rather than a single task capability, producing a compact general-purpose proxy instead of a specialist.
2. Foundry framework: The first FMD implementation for 3D point clouds, built on a compress-and-reconstruct objective in which a lightweight student learns a small set of SuperTokens as an efficient basis for the teacher's latent space.
3. Two novel modules: Dynamic Supertoken Optimization (DSO) compresses input tokens into SuperTokens via hard-assignment cross-attention, and Cross-Attention Upsampling (CAU) reconstructs the teacher's full token set with a residual connection to the original tokens.
4. A budget-aware variant (Foundry-Gate): An optional gating MLP that lets a user trade accuracy for compute at inference time by controlling the fraction of tokens that bypass compression.

Main Findings

• Generalist beats specialists under transfer: A single Foundry student distilled once on ShapeNet55 reaches 89.95% accuracy on ShapeNet55 classification and 81.85/84.77 mIoU (class/instance) on ShapeNetPart, close to the teacher (90.54 and 83.91/85.73), whereas task-specific specialist students achieve only 75.09% (classification) and 61.88/65.72 (segmentation) even on their native tasks.
• Few-shot capability survives extreme compression: On ModelNet40 few-shot benchmarks, Foundry with a single SuperToken (s=1) reaches 91.8% (5-way 10-shot), 95.1% (5-way 20-shot), 88.3% (10-way 10-shot) and 91.5% (10-way 20-shot), against the Point-JEPA teacher's 96.1, 98.2, 93.4 and 95.3.
• Learnable SuperTokens are essential: Replacing the learned DSO/CAU modules with static K-Means clustering drops ShapeNet55 accuracy from 89.68 to 76.08 (a drop of over 13%), and naive Farthest Point Sampling pre-sampling reaches only 87.56.
• Accuracy stays close to the teacher at small budgets: With s=4 SuperTokens, the gateless student reaches 91.25% on ModelNet40 and 84.67/84.81 mIoU on ShapeNetPart, while the full teacher baseline (c=64) reports 93.02% and 83.91/85.73.
• Stable cross-task behavior: The 'Avg' summary column reports 87.72 for the baseline, 84.33–84.87 for Foundry variants (s=16), versus 73.80 (ToMe), 75.20 (random sampling), and 21.46 (random sampling with c=1). The authors note that averaging Accuracy and mIoU is not strictly rigorous.
• Foundry combines with other compression methods: Applying ToMe, PiToMe, or PatchMerger within the FMD training and fine-tuning pipeline raises their Avg scores to 86.52 (ToMe, s=16), 86.51 (PiToMe, s=16) and 87.59 (PatchMerger, s=16), versus 73.80 for untrained ToMe at s=16.
• Distillation loss predicts transferability: A clear inverse correlation is observed between distillation loss and fine-tuning accuracy, with diminishing gains beyond s=4, suggesting reconstruction fidelity saturates once the student spans the teacher's latent space.
• Real speed and memory savings: On an NVIDIA RTX A3000 6GB Laptop GPU, Foundry cuts FLOPs from the baseline's 478 G to 178 G (s=16), 156 G (s=8) and 137 G (s=1), with average latency falling from 0.09 s to 0.05–0.06 s and no memory overhead (1135.5 MiB vs 1143.0 MiB). The ViT-T variant needs 143.698 GFLOPs at similar accuracy.
• Large scenes become feasible: On 2^18-point indoor scenes, both the baseline and ToMe exceed the 6 GB VRAM limit and fail to execute (OOM), while Foundry completes inference with a 4.0 GB footprint, 4.6 s forward time and throughput of roughly 0.22 obj/s.
• Gating enables dynamic compute control: Foundry-Gate lets users fix a token budget or vary the fusion ratio r between 0.1 and 0.9 to modulate latency predictably; Spearman correlations between gate regularization and fine-tuning performance on ModelNet40 are -0.45 (frozen) and -0.13 (unfrozen).

Methodology in Plain English

The authors keep a large pre-trained teacher frozen and train a small student to reproduce it. First, the teacher processes a point cloud's tokens and produces a set of output embeddings — this is the target. The student then does two things. In the compression step, a module called DSO uses cross-attention to assign each of the many input tokens to one of a small, learnable set of SuperTokens, and averages the assigned tokens into each SuperToken, forming a compressed basis of the teacher's latent space. In the reconstruction step, a lightweight encoder refines those SuperTokens, and a CAU module routes each original token location back to its SuperToken, adds the original token as a residual to retain fine detail, and passes the result through an MLP to produce an approximation of the teacher's output. The whole student is trained with a single Smooth L1 reconstruction loss against the teacher's embeddings — no task labels and no logits. A separate gate that bypasses compression for some tokens allows the model to operate at a chosen compute budget.

Why This Matters

This work opens a third route between task-specific distillation and direct feature mimicry — distilling a model's whole representation space — which could change how foundation models are deployed rather than just how they are trained.

Real-world applications:

• Robotics, where perception models must run onboard under tight power and memory limits.
• Autonomous driving, where large-scale scene understanding must happen in real time.
• AR/VR headsets, which need 3D understanding within mobile-class compute and memory budgets.
• Edge and mobile GPUs generally, where the paper shows a full foundation model and ToMe both fail on 2^18-point scenes that Foundry completes with a 4.0 GB footprint.

Industry relevance: the method is architecture-agnostic in its wrapper design, requires only a frozen teacher and one distillation pass, and produces a standalone student that never needs the teacher again — a practical path to shipping 3D perception on constrained hardware.

Future Directions

• Extending FMD beyond point clouds to other unimodal foundation models and modalities, since the framework is described as a general paradigm.
• Scaling the approach to larger teachers than the ViT-S (12 layers, 384 dimensions) and ViT-T backbones used here, and to tasks beyond classification and part segmentation such as detection or full scene understanding.
• Better understanding and improving the gate's behavior, given that the Spearman anti-correlation between gate regularization and performance was weaker (-0.13) when the student was unfrozen than when frozen (-0.45).
• Tighter integration with token-merging methods like ToMe, PiToMe and PatchMerger, which already came close to baseline performance when trained under the FMD objective.

Target Audience

Researchers and engineers working on 3D computer vision, point cloud Transformers, knowledge distillation, and model compression, as well as practitioners deploying perception models on edge devices, robots, or AR/VR hardware. Readers primarily interested in an accessible introduction to distillation will find the formalism (complexity tables, attention equations) demanding, so the paper is best suited to those already comfortable with Transformer architecture and self-supervised learning.

Authors’ abstract

Foundation models pre-trained with self-supervised learning (SSL) on large-scale datasets have become powerful general-purpose feature extractors. However, their immense size and computational cost make them prohibitive for deployment on edge devices such as robots and AR/VR headsets. Existing compression techniques like standard knowledge distillation create efficient 'specialist' models but sacrifice the crucial, downstream-agnostic generality that makes foundation models so valuable. In this paper, we introduce Foundation Model Distillation (FMD), a new paradigm for compressing large SSL models into compact, efficient, and faithful proxies that retain their general-purpose representational power. We present Foundry, the first implementation of FMD for 3D point clouds. Our approach, Foundry, trains a student to learn a compressed set of SuperTokens that reconstruct the teacher's token-level representations, capturing a compact basis of its latent space. A single distilled model maintains strong transferability across diverse downstream tasks-classification, part segmentation, and few-shot scenarios-approaching full foundation-model performance while using significantly fewer tokens and FLOPs, making such models more practical for deployment on resourceconstrained hardware.