huggingface / huggingface/nanoVLM

Add optional NdLinear support for standard Linear layers in ModalityProjector and ViTMLP

Open
#68 2 comments 1 reaction 0 assignees View on GitHub
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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.