$ show TIGER-AI-Lab/Pixel-Reasoner
Pixel-Reasoner implements a two-stage training framework for Vision-Language Models (VLMs) to enable pixel-space reasoning. The system first performs instruction tuning on synthesized reasoning traces to familiarize models with visual operations (zoom-in, select-frame), then applies curiosity-driven reinforcement learning to balance exploration between pixel-space and textual reasoning. Targeting visual reasoning benchmarks (V* bench, TallyQA-Complex, InfographicsVQA, MVBench), it provides pre-trained models and evaluation scripts. The architecture combines adapted OpenRLHF components with custom VLM modules, Ray-based distributed training, and specialized data processors for multimodal input. Key entry points include training scripts (single/multi-node) and evaluation scripts, with data flowing from Hugging Face datasets through tokenization/image processing to PPO optimization loops.
The LLM identified 5 files worth highlighting for reuse:
Contains modular loss function implementations (PolicyLoss, SFTLoss, ValueLoss, DPOLoss) used throughout the RL and training pipeline. Self-contained with minimal dependencies (primarily torch), demonstrating clean abstraction of loss calculations that could be reused in other RL projects.
Implements a reusable replay buffer for storing PPO experiences. Standard RL component with clear separation of concerns, minimal internal coupling, and efficient tensor handling. Could be extracted for use in other actor-critic architectures with little modification.
Abstracts vision-language model processor and tokenizer loading (e.g., for Qwen2VL). Provides a clean interface for handling multimodal inputs across different VLM architectures, reducing boilerplate when integrating new models. Demonstrates effective dependency inversion.
Implements ring attention optimization for efficient long-context processing in transformers. Non-trivial performance optimization that reduces memory complexity from O(nĀ²) to O(n) for sequence processing. Could benefit other transformer-based models handling long sequences.
Defines an abstract base class for VLM modules, enabling clean extension for new vision-language models (e.g., Qwen2VL implementation in qwen_module.py). Shows strong encapsulation of model-specific details while maintaining a consistent interface for training and inference.