Factually-Augmented RLHF: How LLaVA Advances Multimodal Alignment

Introduction to Vision-Language Models

We’re still figuring out how to make machines talk sense, let alone see sense. Multimodal models are the next frontier beyond just text wrangling. They attempt to fuse different data streams – in the case of vision-language models (VLMs), sight and speech – letting them supposedly understand and generate stuff that connects pixels to prose.

While text-slinging Large Language Models (LLMs) like GPT-4 grabbed headlines, the real grind in AI alignment has shifted to these multimodal beasts. One notable player in the open-source arena is LLaVA (Large Language and Vision Assistant), which took a stab at aligning vision-language models using refined human feedback techniques.

This piece dissects how LLaVA implements factually-augmented RLHF (Reinforcement Learning from Human Feedback), why this matters for tackling the hallucination plague in multimodal systems, and how you might borrow these ideas for your own silicon creations.

The Evolution of Vision-Language Models

Before we gut LLaVA, let’s trace the lineage of trying to bolt eyes onto language models:

Early Fumbling (2015-2020): Primitive attempts like Show and Tell and CLIP basically duct-taped separate vision and language models together. Limited integration, clumsy results.
Forced Unity (2020-2022): Models like DALL-E and VisualBERT aimed for more unified architectures, processing pixels and words in shared conceptual spaces. Getting warmer.
Alignment Efforts (2022-Present): The current crop, including LLaVA, GPT-4V, and Gemini, are wrestling with sophisticated alignment techniques. The goal: make their outputs palatable to humans and, crucially, tethered to reality.

It’s the same story we saw with text-only models, just compounded by the messiness of bridging vision and language.

The Hallucination Problem in Multimodal Systems

Hallucinations – the polite term for when the machine confidently spews nonsense – get particularly nasty in multimodal setups:

Cross-modal confusion: The model’s visual processing wires get crossed with its language generation. It “sees” a cat but describes a dog.
Visual ambiguity invites invention: Fuzzy pixels or partial views tempt the model to fill in the blanks, often overconfidently.
Confusing pixels with knowledge: Models struggle to differentiate what’s actually in the image versus what requires looking things up or reasoning beyond the frame.

Picture this: show a VLM an image of a parrot on a perch. A hallucinating model might chirp, “a vibrant parrot playing with a tiny red ball,” even if no ball exists. This confabulation shows the model happily weaving fiction into fact. Naive training objectives, often focused on sounding plausible, don’t punish this enough. Truth takes a backseat to coherence.

What is LLaVA?

LLaVA, cooked up initially by researchers from UW-Madison, Microsoft, and Columbia University back in 2023, was a notable step in open-source multimodal systems. The name – Large Language and Vision Assistant – spells out its ambition.

(source: paper)

Architecture and Development

LLaVA essentially bolts together existing pieces:

Vision Encoder: A pre-trained vision transformer (usually CLIP ViT-L/14) to digest the image input.
Projection Layer: A trainable bit of connective tissue to map visual features into the language model’s world.
Language Model: A standard LLM (started with Vicuna, later used LLaMA 2) to generate text, now supposedly informed by both pixels and prompts.

The interesting bit isn’t the plumbing, though. It’s how they tried to tame it – specifically, using factually-augmented RLHF to try and ground the model’s outputs in visual reality.

Factually-Augmented RLHF: The LLaVA Approach

Vanilla RLHF, common in text models, optimizes for human “vibes” – what sounds helpful, engaging, or stylistically pleasing. Factual accuracy often gets trampled in the process. Factually-augmented RLHF attempts to fix this by force-feeding the model a dose of reality via the feedback loop.

The LLaVA Training Pipeline

LLaVA gets built in three stages:

Pre-training: Teach the model basic vision-language connections using image-text pairs. Like learning the alphabet before writing sentences.
Supervised Fine-Tuning (SFT): Drill the model with high-quality examples of following multimodal instructions (e.g., “Describe this image,” “Count the hats”). Teaching basic manners.
Factually-Augmented RLHF: The critical alignment phase, involving:
- Generating multiple possible responses for image-query pairs.
- Having humans rate responses not just for helpfulness, but explicitly for factual grounding.
- Training a reward model to mimic these human judgments (including the fact-checking).
- Using Reinforcement Learning (PPO) to tune the main model to maximize this fact-aware reward.

The crucial tweak is making factual correctness a non-negotiable part of the human feedback and reward signal. Annotators were explicitly told to penalize outputs that:

Invent objects not present in the image.
Misidentify things clearly visible.
Make bogus claims about where things are.
Hallucinate text that isn’t actually in the image.

This attempts to force the model’s words to more closely match the visual world presented to it.

Inside the RLHF Fine-Tuning Pipeline

Let’s break down the stages of this fact-infused RLHF pipeline:

1. Base Vision-Language Model Preparation

The starting point is the pre-trained VLM, connecting the vision encoder to the LLM via the projection layer:

Vision Encoder (CLIP ViT-L/14) → Projection Layer → LLM (Vicuna/LLaMA)

This base model has learned rudimentary image-text correlations from large datasets.

2. Supervised Fine-Tuning (SFT)

Before RLHF, the model is fine-tuned on examples of good behavior. For LLaVA, this meant datasets containing:

Visual instruction examples: Curated image-instruction-response sets covering description, reasoning, Q&A.
Multimodal chat logs: Conversations where the model answers questions about images helpfully.

This SFT phase provides a decent starting point – the model generates plausible responses that can then be steered by RLHF.

3. Human Feedback Collection

This is where the “factual augmentation” really happens:

Generate options: The SFT model spits out several answers for each image+query.
Apply fact-checking rules: Human annotators use specific guidelines to evaluate:
- Is the mentioned object really there?
- Are spatial descriptions accurate?
- Are attributes like color/size correct?
- Is the model inventing details beyond the image content?
Collect preferences: Annotators rank responses or choose the better one in pairs.

This generates a preference dataset that explicitly prioritizes factual grounding, not just fluency or helpfulness.

4. Reward Model Training

A separate model – the reward model – is trained to predict these human preferences, effectively becoming an automated judge of factual accuracy (as defined by the annotators).

# Concept: A model to predict the 'goodness' score based on human prefs
class RewardModel(nn.Module):
    def __init__(self, vision_encoder, language_model, projection_layer):
        super().__init__()
        self.vision_encoder = vision_encoder
        self.language_model = language_model
        self.projection_layer = projection_layer
        # Simple linear head on top of LLM's final state
        self.reward_head = nn.Linear(language_model.config.hidden_size, 1)

    def forward(self, images, text_responses, attention_masks): # Added attention_masks
        # Process images
        image_features = self.vision_encoder(images).pooler_output # Use pooled output
        projected_features = self.projection_layer(image_features)

        # Process text conditioned on image (conceptual)
        # Actual integration depends heavily on the base LLM architecture
        # This needs modification for a real implementation
        # Assume image features are somehow prepended or injected
        lm_outputs = self.language_model(
            input_ids=text_responses,
            attention_mask=attention_masks,
            # image_features=projected_features # How features integrate is model-specific
        )

        # Extract final hidden state (e.g., last token) and predict reward
        final_hidden = lm_outputs.last_hidden_state[:, -1, :]
        reward = self.reward_head(final_hidden)
        return reward

The reward model learns via a pairwise preference loss. Given an input $x$ (image+query), a preferred response $y_w$ , and a less preferred response $y_l$ , the goal is to maximize the difference in their predicted rewards:

$\mathcal{L}(r_\theta) = -\mathbb{E}_{(x,y_w,y_l) \sim \mathcal{D}} \left[ \log\sigma(r_\theta(x,y_w) - r_\theta(x,y_l)) \right]$

Where:

$r_\theta$ is the reward model parameterized by $\theta$ .
$\sigma$ is the sigmoid function.
$\mathcal{D}$ is the dataset of human preferences emphasizing factual accuracy.

This trains the reward model to assign higher scores to responses humans deemed more factually accurate and helpful.

5. Reinforcement Learning Optimization

The main VLM (the “policy” model) is then fine-tuned using an RL algorithm like PPO to generate responses that score highly according to the trained reward model.

# Simplified PPO-like training loop concept
# Requires libraries like TRL (Transformers Reinforcement Learning)

# Assume policy_model is the VLM to be trained
# Assume reward_model is the trained reward predictor
# Assume training_data yields batches of (images, queries)

for batch in training_data:
    images, queries = batch["images"], batch["queries"]

    # 1. Generate responses from the current policy model
    # Need to get responses AND log probabilities of those responses
    # response_tensors, log_probs_old = policy_model.generate(images, queries, return_log_probs=True)

    # For TRL, you'd prepare query tensors
    query_tensors = tokenizer(queries, return_tensors="pt").input_ids

    # Generate responses using PPOTrainer or similar
    # response_tensors = ppo_trainer.generate(query_tensors, **generation_kwargs)
    # Need to map responses back to text if needed, keep tensors for reward

    # 2. Get rewards for the generated responses using the reward model
    # Assume response_tensors contain the generated token IDs
    # rewards = reward_model(images, response_tensors, attention_mask_for_responses)
    # TRL's PPOTrainer often expects a list of reward tensors

    # 3. Perform PPO update step
    # This step uses the query, response, log_probs, and rewards
    # stats = ppo_trainer.step(query_tensors, response_tensors, rewards)

    # (The actual TRL implementation handles the PPO loss calculation internally)
    # Calculation involves:
    #   - Recomputing log_probs under the current policy (log_probs_new)
    #   - Calculating probability ratios (exp(log_probs_new - log_probs_old))
    #   - Computing advantages (often using KL penalty against reference model)
    #   - Calculating the clipped surrogate objective PPO loss
    #   - Backpropagating the loss and updating the policy_model parameters

PPO aims to maximize the expected reward, constrained by a KL divergence penalty that prevents the policy from straying too far, too fast from a reference model (often the initial SFT model). The core idea is captured by the clipped objective:

$\mathcal{L}^{CLIP}(\theta) = \mathbb{E}_t \left[ \min(r_t(\theta) A_t, \text{clip}(r_t(\theta), 1-\epsilon, 1+\epsilon) A_t) \right]$

Here:

$r_t(\theta) = \frac{\pi_\theta(a_t|s_t)}{\pi_{\theta_{old}}(a_t|s_t)}$ is the probability ratio of the action under new vs old policy.
$A_t$ is the estimated advantage (derived from the reward).
$\epsilon$ is the clipping hyperparameter (e.g., 0.2).

Crucially, because the reward model ( $A_t$ is based on it) was trained on fact-checked preferences, the PPO optimization inherently steers the model towards generating visually grounded, factually accurate outputs, not just plausible ones.

Measuring Success: Evaluating Factual Correctness

Does this convoluted process actually work? Measuring multimodal factual accuracy is tricky, but researchers use benchmarks like:

POPE (Precision of Object Presence Evaluation): Specifically tests if the model hallucinates objects that aren’t there. Can it correctly say “no ball” when there’s no ball?
MME (Multimodal Evaluation): A broader benchmark assessing perception, reasoning, and grounding.
Human evaluation: Still the gold standard, having people directly rate outputs for factual errors.

The results generally show that adding the factual constraint to RLHF significantly cuts down hallucination rates compared to baseline models or standard RLHF.

Here’s an illustrative comparison (numbers are conceptual, based on LLaVA’s findings):

Model Stage	POPE False Positive Rate	MME Factual Score	Human Eval Hallucination %
Base VLM	~24%	~62	~19%
After SFT	~18%	~69	~12%
After regular RLHF	~10%	~76	~5%
After Factual RLHF	~4%	~83	~3%

The trend is clear: explicit factual feedback makes a measurable difference.

Practical Applications and Real-World Impact

Taming hallucinations makes VLMs less likely to be dangerously misleading, potentially unlocking applications like:

Assistive Tech: More reliable descriptions for visually impaired users.
Content Moderation: Better tools for spotting manipulated images (though this is an arms race).
Education: Image-based learning tools that don’t invent facts.
Visual Documentation: More trustworthy automated descriptions for insurance, real estate, etc.

Of course, “more reliable” doesn’t mean perfect. These are still complex systems prone to failure.

Implementing Your Own Factually-Augmented RLHF Pipeline

Thinking of trying this yourself? Brace yourself. Key hurdles include:

Data Nightmare

Diverse Data: You need tons of varied image-text data.
Annotation Hell: Defining and consistently applying “factual accuracy” guidelines is hard and expensive. Needs clear rules and well-trained annotators.
Preference Collection: Designing interfaces and processes to efficiently capture fact-focused preferences is non-trivial.

Technical Stack & Compute

The reference code below sketches the idea of the pipeline components. A real implementation is far more complex, involving careful handling of tokenization, model architectures, distributed training, and RL libraries.

import torch
import torch.nn as nn
import torch.nn.functional as F
# Assume necessary imports from transformers, datasets, trl etc.
# from transformers import CLIPVisionModel, AutoModelForCausalLM, Trainer, AutoTokenizer
# from datasets import load_dataset
# from trl import PPOTrainer, PPOConfig, AutoModelForCausalLMWithValueHead

# --- Model Definitions (Conceptual) ---

# 1. Vision Language Model (Policy Model)
class VisionLanguageModelWithValueHead(nn.Module): # Modified for PPO value head
    def __init__(self, vision_encoder_id, language_model_id):
        super().__init__()
        self.vision_encoder = CLIPVisionModel.from_pretrained(vision_encoder_id)
        self.language_model = AutoModelForCausalLM.from_pretrained(language_model_id)
        # Projection layer (adapt dimensions)
        self.projection = nn.Linear(
            self.vision_encoder.config.hidden_size,
            self.language_model.config.hidden_size
        )
        # Value head for PPO critic
        self.value_head = nn.Linear(self.language_model.config.hidden_size, 1)

    def encode_images(self, pixel_values):
        # ... (image encoding logic as before) ...
        vision_outputs = self.vision_encoder(pixel_values=pixel_values)
        # Using pooler output or CLS token might be typical
        image_features = vision_outputs.pooler_output
        projected_features = self.projection(image_features)
        return projected_features # Shape: [batch_size, language_hidden_size]

    def forward(self, pixel_values, input_ids, attention_mask, return_value=False):
        image_features_proj = self.encode_images(pixel_values) # [batch, lang_hidden]

        # --- How image features integrate is CRITICAL and model-specific ---
        # Example: Prepend image features as soft prompts
        # Requires modifying input_ids and attention_mask accordingly
        # This is highly simplified; real implementations (LLaVA, etc.) use more complex methods.
        inputs_embeds = self.language_model.get_input_embeddings()(input_ids) # [batch, seq_len, lang_hidden]
        # Prepend image features (unsqueezed) to token embeddings
        image_features_embeds = image_features_proj.unsqueeze(1) # [batch, 1, lang_hidden]
        combined_embeds = torch.cat([image_features_embeds, inputs_embeds], dim=1) # [batch, 1+seq_len, lang_hidden]

        # Adjust attention mask
        image_attention = torch.ones(combined_embeds.size(0), 1, device=attention_mask.device)
        combined_attention_mask = torch.cat([image_attention, attention_mask], dim=1)

        outputs = self.language_model(
            inputs_embeds=combined_embeds,
            attention_mask=combined_attention_mask,
            output_hidden_states=True # Needed for value head potentially
        )

        logits = outputs.logits[:, 1:, :] # Shift logits to align with original input_ids

        if return_value:
             # Value prediction, often from last hidden state of sequence
            last_hidden_state = outputs.hidden_states[-1][:, -1, :]
            value = self.value_head(last_hidden_state).squeeze(-1)
            return logits, value
        else:
            return logits

# 2. Reward Model (Separate model usually)
class RewardModel(nn.Module):
     def __init__(self, vision_encoder_id, language_model_id):
         super().__init__()
         # Typically uses the same base architecture but trained differently
         self.vision_encoder = CLIPVisionModel.from_pretrained(vision_encoder_id)
         self.language_model = AutoModelForCausalLM.from_pretrained(language_model_id)
         self.projection = nn.Linear(
             self.vision_encoder.config.hidden_size,
             self.language_model.config.hidden_size
         )
         self.reward_head = nn.Linear(self.language_model.config.hidden_size, 1)

     def forward(self, pixel_values, input_ids, attention_mask):
        # Similar forward pass to VLM, but ends with reward head
        image_features_proj = self.encode_images(pixel_values)
        inputs_embeds = self.language_model.get_input_embeddings()(input_ids)
        image_features_embeds = image_features_proj.unsqueeze(1)
        combined_embeds = torch.cat([image_features_embeds, inputs_embeds], dim=1)
        image_attention = torch.ones(combined_embeds.size(0), 1, device=attention_mask.device)
        combined_attention_mask = torch.cat([image_attention, attention_mask], dim=1)

        outputs = self.language_model(
            inputs_embeds=combined_embeds,
            attention_mask=combined_attention_mask,
            output_hidden_states=True
        )
        # Reward often based on the final token's hidden state
        last_hidden = outputs.hidden_states[-1][:, -1, :]
        reward = self.reward_head(last_hidden)
        return reward

     def encode_images(self, pixel_values):
         # (Identical image encoding logic as VLM)
         vision_outputs = self.vision_encoder(pixel_values=pixel_values)
         image_features = vision_outputs.pooler_output
         projected_features = self.projection(image_features)
         return projected_features

# --- Dataset Preparation (Conceptual) ---
def prepare_datasets():
    print("Loading datasets (placeholders)...")
    # Placeholder: Load actual datasets here
    # pretrain_dataset = load_dataset(...)
    # sft_dataset = load_dataset(...) # Requires image, instruction, response
    # preference_dataset = load_dataset(...) # Requires image, query, chosen_response, rejected_response
    print("Datasets loaded.")
    # Preprocessing (tokenization, image transforms) would happen here
    return None, None, None # Return processed datasets

# --- Training Functions (Conceptual) ---
def pretrain_or_sft(model, dataset, training_args):
    print("Starting Pretraining/SFT (placeholder)...")
    # trainer = Trainer(model=model, args=training_args, train_dataset=dataset)
    # trainer.train()
    print("Pretraining/SFT finished.")
    return model

def train_reward_model(reward_model_instance, preference_dataset, training_args):
    print("Starting Reward Model training (placeholder)...")
    # Custom training loop needed for pairwise loss
    # optimizer = torch.optim.AdamW(reward_model_instance.parameters(), lr=training_args.learning_rate)
    # for epoch in range(int(training_args.num_train_epochs)):
    #     for batch in preference_dataset: # Assuming dataloader handles batching
    #         # pixel_values, chosen_ids, chosen_mask, rejected_ids, rejected_mask = batch
    #         chosen_rewards = reward_model_instance(pixel_values, chosen_ids, chosen_mask)
    #         rejected_rewards = reward_model_instance(pixel_values, rejected_ids, rejected_mask)
    #         loss = -F.logsigmoid(chosen_rewards - rejected_rewards).mean()
    #         # Backpropagate and optimize
    #         loss.backward()
    #         optimizer.step()
    #         optimizer.zero_grad()
    print("Reward Model training finished.")
    return reward_model_instance

def ppo_training(policy_model_with_value_head, reward_model_instance, dataset, ppo_config):
    print("Starting PPO training (placeholder)...")
    # Requires PPOTrainer setup from TRL
    # tokenizer = AutoTokenizer.from_pretrained(...) # Need tokenizer
    # ppo_trainer = PPOTrainer(config=ppo_config, model=policy_model_with_value_head, ...)
    # for batch in dataset: # Dataset yields (pixel_values, query_ids, query_mask)
    #    query_tensors = batch['query_ids']
    #    # 1. Generate response tensors + log_probs from policy
    #    # response_tensors = ppo_trainer.generate(query_tensors, ...)
    #    # 2. Compute rewards using reward_model
    #    # reward_outputs = reward_model_instance(batch['pixel_values'], response_tensors, ...)
    #    # 3. PPO step
    #    # stats = ppo_trainer.step(query_tensors, response_tensors, reward_outputs)
    print("PPO training finished.")
    return policy_model_with_value_head

# --- Main Pipeline Execution (Conceptual) ---
def train_factually_augmented_vlm():
    # Configs (replace with actual paths/IDs)
    vision_encoder_id = "openai/clip-vit-large-patch14"
    language_model_id = "meta-llama/Llama-2-7b-hf" # Or appropriate base

    # Initialize policy model (VLM + Value Head)
    print("Initializing models...")
    policy_model = VisionLanguageModelWithValueHead(vision_encoder_id, language_model_id)
    # Initialize reward model (separate instance)
    reward_model_instance = RewardModel(vision_encoder_id, language_model_id)
    print("Models initialized.")

    # Load and preprocess datasets
    pretrain_dataset, sft_dataset, preference_dataset = prepare_datasets()

    # Define Training Arguments (example)
    training_args = {"output_dir": "./results", "num_train_epochs": 1, "learning_rate": 2e-5} # Add many more
    ppo_config = {"learning_rate": 1e-6, "batch_size": 4} # Add many more

    # --- Execute Training Stages ---
    # Assumes datasets are ready
    # model = pretrain_or_sft(policy_model, pretrain_dataset, training_args) # Optional Pretrain
    # model = pretrain_or_sft(policy_model, sft_dataset, training_args)      # SFT
    # trained_reward_model = train_reward_model(reward_model_instance, preference_dataset, training_args) # Train RM
    # final_model = ppo_training(policy_model, trained_reward_model, sft_dataset, ppo_config) # PPO

    print("Pipeline simulation complete.")
    # In reality, you'd save the final PPO-tuned model
    # final_model.save_pretrained("your_factually_aligned_vlm")
    return policy_model # Return the (conceptually) trained model

if __name__ == "__main__":
    final_model = train_factually_augmented_vlm()

Brutal Compute Needs

RLHF, especially multimodal, burns through GPUs. Forget doing this on a laptop.
Prototyping on smaller models is wise before scaling up.
Techniques like RLAIF (RL from AI Feedback) might ease the pain, but don’t expect miracles.
Distributed training across multiple machines is almost mandatory for serious attempts.

Future Directions and Open Challenges

Despite progress like LLaVA, multimodal alignment is far from solved. Persistent headaches include:

Cross-modal Grounding: Reliably connecting specific visual details to precise language remains hard, especially for complex or abstract stuff.
Evaluation Metrics: We need better automated ways to measure factual accuracy at scale. Human evaluation is slow and costly. Current benchmarks are still blunt instruments.
Domain Adaptation: Making this work well for specialized areas (medicine, science) is another beast entirely.
Efficiency: The sheer cost and complexity of RLHF remain a major bottleneck. Finding cheaper, faster ways to achieve similar results is crucial.

$Future Research \/ Pain Points$

Conclusion

Factually-augmented RLHF, as demonstrated by LLaVA, is a serious attempt to beat hallucinations in multimodal AI. Forcing factual correctness into the human feedback loop yields models that lie less about what they see. It’s a step towards grounding these systems in reality, however imperfectly.

As VLMs creep into more applications, making them factually reliable becomes paramount. LLaVA offers a useful case study, showing one plausible path: explicitly target factual accuracy during alignment, don’t just optimize for helpfulness or style.

For those building these systems, LLaVA’s lesson is clear: reducing hallucinations requires deliberate, costly intervention in the training process. This principle – injecting factual checks into the feedback – is adaptable, but demands resources and rigor.

The field is grinding forward. Expect more refinements, but the road to VLMs that genuinely understand what they see – assuming “understanding” is even the right word – remains long and expensive.

Posted in AI / ML, LLM Advanced, LLM Research by Rakshit Kalra