Think-When-Needed

KO | EN

gitlite — search

#hacktoberfest #go #android #rust #deep-learning #docker #nodejs #macos #windows #awesome #cli #vue #pytorch #security

Think-When-Needed

(★ 11)

No description available.

Think-When-Needed 최신버젼 다운로드

최종 버전 다운로드 (.zip)

// repository documentation

Was this content helpful?

★ 0 (0 ratings)

# Think When Needed: Adaptive Reasoning-Driven Multimodal Embeddings with a Dual-LoRA Architecture

[![Paper](https://img.shields.io/badge/arXiv-2605.14448-b31b1b?style=flat-square)](https://arxiv.org/abs/2605.14448) [![Hugging Face Models](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Models-blue?style=flat-square)](https://huggingface.co/collections/zhanglx/think-when-needed-twn) [![Hugging Face Data](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Data-yellow?style=flat-square)](https://huggingface.co/datasets/zhanglx/TWN-training-data) [![License](https://img.shields.io/badge/License-Apache_2.0-green.svg?style=flat-square)](https://opensource.org/licenses/Apache-2.0)

</div>

**TWN (Think When Needed)** is a framework for adaptive reasoning-driven multimodal embeddings. It introduces a dual-LoRA architecture with a learned routing gate that adaptively activates chain-of-thought reasoning.

## News

* **[2026.05.22]** We have released the model weights ([TWN-4B](https://huggingface.co/zhanglx/TWN-4B), [TWN-8B](https://huggingface.co/zhanglx/TWN-8B)), [training data](https://huggingface.co/datasets/zhanglx/TWN-training-data), and training code.

---

## Highlights

* **Dual-LoRA Architecture:** Decouples reasoning and embedding into two specialized LoRA adapters — a *reasoning* adapter that generates chain-of-thought and an *embedding* adapter that extracts the final representation — enabling each to be independently optimized.
* **Adaptive Routing Gate:** An MLP-based gate dynamically decides per query whether CoT reasoning is needed, reducing redundant computation while improving retrieval performance.
* **Two-Stage Training:** SFT with contrastive loss and next-token prediction loss, followed by GRPO-based RL, progressively refining reasoning quality and retrieval metrics.

---

## Quick Start

### 1. Installation

```bash
git clone https://github.com/winterfell00/Think-When-Needed.git
cd Think-When-Needed

conda create -n twn python=3.11 -y
conda activate twn

pip install -r requirements.txt
```

### 2. Project Structure

```
.
├── config/                  # Training configurations
│   ├── sft_4b.yaml          #   SFT config for Qwen3-VL-4B
│   ├── sft_8b.yaml          #   SFT config for Qwen3-VL-8B
│   ├── rl_4b.yaml           #   RL config for Qwen3-VL-4B
│   └── rl_8b.yaml           #   RL config for Qwen3-VL-8B
├── src/                     # Source code
│   ├── sft/                 #   SFT training
│   │   ├── train.py         #     Training entry point
│   │   ├── twn.py           #     TWN model (dual-LoRA + routing gate)
│   │   ├── data_module.py   #     Dataset & data loading
│   │   └── model_module.py  #     Lightning model module
│   ├── rl/                  #   RL training
│   │   ├── train.py         #     Training entry point
│   │   ├── twn_rl.py        #     TWN-RL model
│   │   ├── qwenvl_grpo_trainer.py  # GRPO trainer
│   │   ├── rewards.py       #     Reward functions
│   │   ├── data_module.py   #     RL dataset
│   │   └── precompute_embeddings.py  # Pre-compute target embeddings
│   ├── grad_cache/          #   GradCache for memory-efficient training
│   └── utils.py             #   Shared utilities
├── eval/                    # Evaluation
│   ├── eval.py              #   Evaluation entry point
│   ├── eval.sh              #   Evaluation launch script
│   ├── config/              #   Eval configs (image, video, visdoc)
│   └── src/                 #   Eval utilities & dataset loaders
├── script/                  # Launch scripts
│   ├── run_sft.sh           #   Launch SFT training
│   └── run_rl.sh            #   Launch RL training
└── requirements.txt
```

**Training entry points:**
- SFT: `src/sft/train.py` (launched via `script/run_sft.sh`)
- RL: `src/rl/train.py` (launched via `script/run_rl.sh`)

**Evaluation entry point:**
- `eval/eval.py` (launched via `eval/eval.sh`)

---

## Model Zoo

---

## Training

We provide training scripts for both SFT and RL stages. Configs are available for Qwen3-VL 4B and 8B backbones.

### Prepare Data

1. Download the training data from [Hugging Face](https://huggingface.co/datasets/zhanglx/TWN-training-data) and set the data paths in your config.

2. Download the source media from the original datasets and set `media_roots` in your config to point to local copies:

| URI Scheme | Source Dataset |
| :--- | :--- |
| `mmeb_train://` | [TIGER-Lab/MMEB-train](https://huggingface.co/datasets/TIGER-Lab/MMEB-train) |
| `sharegpt_video://` | [ShareGPT4Video/train_video_and_instruction](https://huggingface.co/datasets/ShareGPT4Video/ShareGPTVideo/tree/main/train_video_and_instruction) |
| `visrag_in_domain://` | [openbmb/VisRAG-Ret-Train-In-domain-data](https://huggingface.co/datasets/openbmb/VisRAG-Ret-Train-In-domain-data) |
| `visrag_synthetic://` | [openbmb/VisRAG-Ret-Train-Synthetic-data](https://huggingface.co/datasets/openbmb/VisRAG-Ret-Train-Synthetic-data) |
| `vidore_colpali://` | [vidore/colpali_train_set](https://huggingface.co/datasets/vidore/colpali_train_set) |

3. Update your config file (e.g., `config/sft_4b.yaml`):

```yaml
dataset:
  params:
    dataset_args:
      media_roots:
        mmeb_train: /path/to/MMEB-train
        sharegpt_video: /path/to/ShareGPTVideo
        visrag_in_domain: /path/to/VisRAG-Ret-Train-In-domain-data
        visrag_synthetic: /path/to/VisRAG-Ret-Train-Synthetic-data
        vidore_colpali: /path/to/colpali_train_set
```

### SFT

Before running, set `pretrained_path` and `processor_config_path` in the config to your local Qwen3-VL checkpoint, and configure `media_roots` as described in [Prepare Data](#prepare-data).

```bash
GPUS=8 ./script/run_sft.sh config/sft_4b.yaml
```

### RL

Set the SFT checkpoint path (`trained_path`) and RL data path in the RL config before launching:

```bash
GPUS=8 ./script/run_rl.sh config/rl_4b.yaml
```

### Configuration

---

## Evaluation

For the complete evaluation pipeline setup, please refer to [VLM2Vec](https://github.com/TIGER-AI-Lab/VLM2Vec). The evaluation code in this repository is adapted to integrate with our TWN model.

```bash
bash eval/eval.sh
```

We evaluate TWN on 78 datasets across three domains: Image (36 datasets), Video (18 datasets), and VisDoc (24 datasets).

---

## Citation

If you find our work useful, please consider citing:

```bibtex
@article{zhang2026thinkneeded,
      title={Think When Needed: Adaptive Reasoning-Driven Multimodal Embeddings with a Dual-LoRA Architecture},
      author={Longxiang Zhang and Weilong Dai and Guanghao Zhang and Hao Jiang and Pipei Huang},
      year={2026},
      eprint={2605.14448},
      archivePrefix={arXiv},
      primaryClass={cs.CV},
      url={https://arxiv.org/abs/2605.14448},
}
```

## License

This project is open-sourced under the [Apache-2.0 License](https://opensource.org/licenses/Apache-2.0).

Think-When-Needed

Think-When-Needed 최신버젼 다운로드

오늘의 추천 프로젝트

Was this content helpful?

Recent Feedback

Think-When-Needed

Think-When-Needed 최신버젼 다운로드

오늘의 추천 프로젝트

Was this content helpful?

Recent Feedback

Download README