huggingface / huggingface/nanoVLM
Add optional NdLinear support for standard Linear layers in ModalityProjector and ViTMLP
- Dominant language
- Python
- Stars
- 5k
- Forks
- 510
- PR merge metrics
- No merged PRs in 30d
Description
I'd like to propose an enhancement to `nanoVLM` by adding optional support for [NdLinear](https://github.com/ensemble-core/NdLinear), a shape-preserving and parameter-efficient drop-in replacement for the standard `nn.Linear` layer.
=> Related paper: [NdLinear Is All You Need for Representation Learning (Ensemble AI, 2025)](https://arxiv.org/pdf/2503.17353)
---
### Motivation
Most neural architectures flatten inputs, discarding critical cross-dimensional information.
**`NdLinear`** avoids this by operating directly on multi-dimensional tensors, enabling more efficient modeling while preserving structural relationships.
It serves as a foundational building block for large-scale foundation models by processing unimodal or multimodal data in its native shape — eliminating the need for flattening or modality-specific preprocessing.
**Potential benefits:**
1. Reduced parameter count
2. Improved inference latency
3. Preservation of structural awareness
4. Comparable performance to `nn.Linear`
These advantages are especially valuable in VLM applications involving spatial layouts, biomedical imagery, or grid-structured tokens — where retaining dimensional structure is crucial for performance.
---
### Proposed Additions
- Add two optional flags to `VLMConfig`:
- `use_ndlinear` → enables NdLinear in `ModalityProjector` (`modality_projector.py`)
- `use_ndlinear_in_vit_mlp` → enables NdLinear in `ViTMLP` (`vision_transformer.py`)
- Modify both modules to conditionally use `NdLinear` as a drop-in replacement
- Include a benchmark script (`benchmark_ndlinear.py`) to evaluate:
- Inference latency (forward pass)
- Model parameter count
All changes will be **fully backward-compatible**, with `NdLinear` only used when explicitly enabled via config.
Would love to know if you're open to a PR , happy to prepare a clean and meaningful contribution!
Thanks!
Contributor guide
No contributing guide indexed for this repository
Research direction
Start by reading modality_projector.py, vision_transformer.py, and the VLMConfig definition to understand the existing Linear paths and configuration flow. Review the NdLinear dependency and the proposed benchmark_ndlinear.py entry point. Done means both flags preserve current behavior by default, enable the requested replacements when set, and the benchmark reports forward-pass latency and parameter count.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- machine-learning
- Issue type
- Feature
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100