??? Expanding the Scope of LLMs: Multimodal and Task-Enhanced AI

Multimodal Large Language Models (LLMs) that understand both text and images (or other media formats) are becoming crucial in applications like visual question answering, captioning, and cross-modal retrieval. In this article, we’ll walk through building a multimodal LLM using PyTorch, focusing on text and image inputs. The concepts can be extended to other modalities, like audio or video.

Part 1 - Building a Multimodal LLM in Pytorch

Step 1: Setting up the Environment

First, ensure your environment has the necessary dependencies:

pip install torch torchvision transformers        

We'll use Hugging Face's transformers for the text model and torchvision for handling images.

Step 2: Model Architecture

We will combine a pre-trained text model with a vision model. The text model will encode the text input, while the vision model encodes the image. We'll then fuse these embeddings for downstream tasks like classification or captioning.

Text Encoder: BERT

We'll use BERT for text encoding:

from transformers import BertModel, BertTokenizer

tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
text_model = BertModel.from_pretrained('bert-base-uncased')

def encode_text(text):
    inputs = tokenizer(text, return_tensors='pt', padding=True, truncation=True)
    outputs = text_model(**inputs)
    return outputs.last_hidden_state.mean(dim=1)  # [batch_size, hidden_size]        

Vision Encoder: ResNet

We'll use a pre-trained ResNet for image encoding:

from torchvision import models, transforms
from PIL import Image

vision_model = models.resnet50(pretrained=True)
vision_model.fc = torch.nn.Identity()  # Remove classification head

preprocess = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

def encode_image(image_path):
    image = Image.open(image_path).convert('RGB')
    image_tensor = preprocess(image).unsqueeze(0)  # [batch_size, channels, height, width]
    return vision_model(image_tensor)  # [batch_size, 2048]        

Step 3: Multimodal Fusion

We now need to combine the text and image embeddings into a single representation. A simple approach is concatenation followed by a linear layer:

import torch

class MultimodalModel(torch.nn.Module):
    def __init__(self, text_dim, image_dim, hidden_dim, output_dim):
        super(MultimodalModel, self).__init__()
        self.fc1 = torch.nn.Linear(text_dim + image_dim, hidden_dim)
        self.fc2 = torch.nn.Linear(hidden_dim, output_dim)

    def forward(self, text_emb, image_emb):
        combined = torch.cat((text_emb, image_emb), dim=1)  # [batch_size, text_dim + image_dim]
        x = torch.relu(self.fc1(combined))
        return self.fc2(x)

# Define model dimensions
text_dim = 768
image_dim = 2048
hidden_dim = 512
output_dim = 10  # e.g., for a classification task

model = MultimodalModel(text_dim, image_dim, hidden_dim, output_dim)        

Step 4: Training the Model

Here’s a simple training loop for classification, assuming you have a dataset of text-image pairs:

optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
criterion = torch.nn.CrossEntropyLoss()

for epoch in range(num_epochs):
    model.train()
    total_loss = 0

    for text, image, label in dataloader:  # Assuming dataloader yields text, image, label batches
        text_emb = encode_text(text)
        image_emb = encode_image(image)

        outputs = model(text_emb, image_emb)
        loss = criterion(outputs, label)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        total_loss += loss.item()

    print(f"Epoch {epoch + 1}/{num_epochs}, Loss: {total_loss / len(dataloader)}")        

Step 5: Inference

Once trained, using the model for inference is simple:

text = "A cat sitting on a couch."
image_path = "cat.jpg"

with torch.no_grad():
    text_emb = encode_text(text)
    image_emb = encode_image(image_path)
    prediction = model(text_emb, image_emb)

print(f"Predicted class: {torch.argmax(prediction)}")        

Just incase you wanted to introduce more modalities into the LLM -

Audio Encoder: Wav2Vec2

For audio, we can use a pre-trained Wav2Vec2 model from Hugging Face’s transformers library. Wav2Vec2 converts raw audio waveforms into embeddings, similar to how BERT handles text.

from transformers import Wav2Vec2Processor, Wav2Vec2Model
import torchaudio

processor = Wav2Vec2Processor.from_pretrained("facebook/wav2vec2-base-960h")
audio_model = Wav2Vec2Model.from_pretrained("facebook/wav2vec2-base-960h")

def encode_audio(audio_path):
    waveform, sample_rate = torchaudio.load(audio_path)
    inputs = processor(waveform, sampling_rate=sample_rate, return_tensors="pt", padding=True)
    outputs = audio_model(**inputs)
    return outputs.last_hidden_state.mean(dim=1)  # [batch_size, hidden_size]        

Video Encoder: 3D CNN (ResNet3D)

For video data, we can use a 3D convolutional neural network (CNN) like ResNet3D, which captures temporal information from video frames.

import torchvision

# Load a pre-trained 3D ResNet model
video_model = torchvision.models.video.r3d_18(pretrained=True)
video_model.fc = torch.nn.Identity()  # Remove classification head

def encode_video(video_path):
    frames = extract_frames(video_path)  # Function to extract frames from the video
    frame_tensor = preprocess(frames).unsqueeze(0)  # [batch_size, num_frames, channels, height, width]
    return video_model(frame_tensor)  # [batch_size, 512]        

For extract_frames, use a video processing library like opencv or torchaudio.

Tabular Data Encoder: MLP

For tabular data (e.g., numerical or categorical data), we can use a simple Multilayer Perceptron (MLP) as the encoder:

class TabularEncoder(torch.nn.Module):
    def __init__(self, input_dim, hidden_dim):
        super(TabularEncoder, self).__init__()
        self.fc1 = torch.nn.Linear(input_dim, hidden_dim)
        self.fc2 = torch.nn.Linear(hidden_dim, hidden_dim)

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        return self.fc2(x)

tabular_model = TabularEncoder(input_dim=20, hidden_dim=128)  # Adjust input_dim to your data        

Multimodal Fusion with Multiple Encoders

Now that we have additional encoders, you can modify the fusion step to include all modalities:

class ExtendedMultimodalModel(torch.nn.Module):
    def __init__(self, text_dim, image_dim, audio_dim, video_dim, tabular_dim, hidden_dim, output_dim):
        super(ExtendedMultimodalModel, self).__init__()
        self.fc1 = torch.nn.Linear(text_dim + image_dim + audio_dim + video_dim + tabular_dim, hidden_dim)
        self.fc2 = torch.nn.Linear(hidden_dim, output_dim)

    def forward(self, text_emb, image_emb, audio_emb, video_emb, tabular_emb):
        combined = torch.cat((text_emb, image_emb, audio_emb, video_emb, tabular_emb), dim=1)
        x = torch.relu(self.fc1(combined))
        return self.fc2(x)        

Part 2 - Integrating External Tools for Specific Tasks

Neural networks, including LLMs, are not inherently designed for tasks like precise arithmetic calculations or fetching real-time data from external sources (e.g., the web). To handle these tasks within a multimodal LLM, we can augment the model with external tools or specialized modules that handle these operations more efficiently. Below are two strategies to integrate such capabilities:

1. Arithmetic Calculations via External Functions

Instead of relying on the LLM for arithmetic tasks, we can explicitly route arithmetic requests to an external function or tool that can handle the computation with higher precision. This can be achieved through prompt-engineering-based task recognition, where the LLM identifies an arithmetic request and invokes the appropriate function.

Here’s a conceptual workflow for arithmetic operations:

1. Task Recognition: The model identifies when a query involves arithmetic (e.g., addition, subtraction).
2. External Calculation Function: The query is routed to a specific Python function or an API to handle the calculation.
3. Response Integration: The result is passed back to the model for generating the final output.

def perform_calculation(query):
    # This function extracts numbers and operators from the query and performs the operation.
    if "add" in query or "+" in query:
        numbers = [int(num) for num in query.split() if num.isdigit()]
        return sum(numbers)
    elif "subtract" in query or "-" in query:
        numbers = [int(num) for num in query.split() if num.isdigit()]
        return numbers[0] - numbers[1]

# Extend model logic to invoke external calculator
def extended_forward(text_emb, image_emb, query):
    if "add" in query or "subtract" in query:  # Task identification
        result = perform_calculation(query)    # Perform calculation
        return f"The result of the operation is {result}"  # Return response
    else:
        # Process normally
        combined = torch.cat((text_emb, image_emb), dim=1)
        x = torch.relu(self.fc1(combined))
        return self.fc2(x)        

In practice, you can build a more sophisticated module to handle various arithmetic operations, allowing the LLM to offload this task to a more accurate calculation engine.

2. Fetching Real-Time Data from the Web

For tasks that involve fetching real-time data (e.g., stock prices, weather, or web scraping), we can integrate an external API or tool to retrieve the data and pass it back into the model. The LLM acts as a controller, directing queries to appropriate external systems for execution.

Here’s a sample workflow for integrating web data retrieval:

1. Task Recognition: The LLM identifies that the query involves fetching real-time data.
2. API or Web Scraper Call: Based on the task (e.g., weather query), an external API (e.g., OpenWeather API) or a web scraper is invoked.
3. Response Integration: The data fetched is passed back to the LLM for generating the final response.

Example:

import requests

def fetch_web_data(query):
    if "weather" in query:
        city = query.split()[-1]
        response = requests.get(f"https://api.openweathermap.org/data/2.5/weather?q={city}&appid=your_api_key")
        data = response.json()
        return f"The current temperature in {city} is {data['main']['temp']}°C."
    # Add more APIs for other data-fetching tasks

# Extend model logic to handle web queries
def extended_forward_with_web(text_emb, image_emb, query):
    if "weather" in query:
        result = fetch_web_data(query)  # Fetch real-time data
        return result  # Return response
    else:
        # Process normally
        combined = torch.cat((text_emb, image_emb), dim=1)
        x = torch.relu(self.fc1(combined))
        return self.fc2(x)        

3. Hybrid Models with Specialized Modules

In a more advanced setup, you can use a hybrid model approach where the LLM identifies specific tasks and dynamically routes the request to a module suited for that task, such as:

• Symbolic reasoning engines for arithmetic and logic tasks.
• Database connectors for retrieving information from structured data.
• Web scraping or API query modules for real-time data fetching.

This modular architecture ensures that each task is handled by the most efficient and accurate system while keeping the LLM focused on language understanding and generation.

Part 3 - LLMs with RAG & Vector Stores

Retrieval-Augmented Generation (RAG) in PyTorch

RAG integrates retrieval and generation by leveraging a pre-trained language model (like BERT) for document retrieval and a generative model (like GPT-2 or BART) for response generation. Below is a simplified PyTorch implementation that demonstrates how to set up a RAG pipeline with a vector store.

1. Query Encoding (BERT Encoder for Retrieval)

We’ll use a BERT model to encode the query into a vector, which we can compare to vectors in the vector store.

from transformers import BertTokenizer, BertModel
import torch

# Load a pre-trained BERT tokenizer and model
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
bert_model = BertModel.from_pretrained('bert-base-uncased')

def encode_query(query):
    inputs = tokenizer(query, return_tensors='pt', truncation=True, padding=True)
    with torch.no_grad():
        outputs = bert_model(**inputs)
    return outputs.last_hidden_state.mean(dim=1)  # [batch_size, hidden_size]        

2. Vector Store and Document Retrieval

For simplicity, let’s assume we have a list of pre-encoded documents in a vector store. We can use cosine similarity to find the most relevant document based on the query.

import numpy as np
from sklearn.metrics.pairwise import cosine_similarity

# Example vector store with pre-encoded documents
vector_store = np.random.rand(10, 768)  # Assume we have 10 documents encoded into vectors

def retrieve_document(query_vector, vector_store):
    similarities = cosine_similarity(query_vector, vector_store)
    top_doc_idx = np.argmax(similarities)
    return top_doc_idx, similarities[0][top_doc_idx]  # Return the most similar document index and its score        

2.1 For an advanced Vector Store

A vector store is a specialized data structure that stores high-dimensional vectors representing documents or other data. In a multimodal setting, vector stores can store embeddings for text, images, or other modalities.

To make retrieval efficient at scale, vector stores often use techniques like Approximate Nearest Neighbors (ANN). Libraries like FAISS (Facebook AI Similarity Search) can help with efficient retrieval over millions of vectors.

Using FAISS for Efficient Document Retrieval

Here’s an example of integrating FAISS to handle large-scale vector retrieval.

import faiss

# Build the FAISS index
dimension = 768  # Vector dimension for BERT embeddings
index = faiss.IndexFlatL2(dimension)

# Add document vectors to the FAISS index
index.add(vector_store.astype(np.float32))

# Function to retrieve document using FAISS
def faiss_retrieve(query_vector, index, document_texts):
    query_vector = query_vector.cpu().numpy().astype(np.float32)
    distances, indices = index.search(query_vector, k=1)  # Retrieve top-1 document
    top_doc_idx = indices[0][0]
    return document_texts[top_doc_idx], distances[0][0]

# Example usage
retrieved_document, distance = faiss_retrieve(encode_query(query), index, document_texts)
print("Retrieved document:", retrieved_document)        

3. Fusion with Generative Model (BART for Generation)

Once we retrieve the relevant document, we pass it along with the query to a generative model (like BART) to generate a response.

from transformers import BartTokenizer, BartForConditionalGeneration

# Load a pre-trained BART tokenizer and model
bart_tokenizer = BartTokenizer.from_pretrained('facebook/bart-large')
bart_model = BartForConditionalGeneration.from_pretrained('facebook/bart-large')

def generate_response(query, retrieved_document):
    input_text = query + " " + retrieved_document  # Concatenate query and retrieved doc
    inputs = bart_tokenizer(input_text, return_tensors='pt', truncation=True, padding=True)
    with torch.no_grad():
        outputs = bart_model.generate(**inputs, max_length=50)
    return bart_tokenizer.decode(outputs[0], skip_special_tokens=True)        

4. Complete RAG Process

Now, we can integrate the query encoding, document retrieval, and response generation into a complete RAG process.

def rag_pipeline(query, vector_store, document_texts):
    # Step 1: Encode the query
    query_vector = encode_query(query)

    # Step 2: Retrieve the most relevant document
    doc_idx, similarity = retrieve_document(query_vector, vector_store)
    retrieved_document = document_texts[doc_idx]

    # Step 3: Generate a response based on the query and retrieved document
    response = generate_response(query, retrieved_document)
    
    return response, retrieved_document, similarity

# Example usage
document_texts = ["Document 1 content...", "Document 2 content...", "Document 3 content..."]  # Example docs
query = "Explain multimodal models"
response, doc, sim = rag_pipeline(query, vector_store, document_texts)

print("Retrieved Document:", doc)
print("Generated Response:", response)        

Conclusion

In this article, we've explored how to build a multimodal LLM capable of understanding and reasoning across various data types like text, images, and even audio. By integrating external tools for tasks like arithmetic calculations and real-time data retrieval, we extend the model's capabilities beyond traditional language generation. Furthermore, incorporating Retrieval-Augmented Generation (RAG) and vector stores enhances the model's reasoning ability within specific domains by allowing it to access relevant information dynamically, ensuring more precise and contextually aware responses.

This multimodal and hybrid approach opens up exciting new possibilities for LLMs in various applications—whether for handling visual data, integrating deterministic operations, or retrieving and utilizing up-to-date knowledge in real-time. As LLMs continue to evolve, integrating these techniques will be crucial for building robust, versatile AI systems that can understand and respond to a wide range of inputs.