Skip to content

kyegomez/HRTX

Repository files navigation

Multi-Modality

HRTX

Hivemind Multi-Modality Transformer (HMMT) Model Architecture. Multi-Modality Model that can ingest text and video modalities from an x amount of models at the same time and then output instructions for each robot or any N number of robots.

Install

pip install hrtx

Usage

EarlyExitTransformer

There are various output heads after every transformer forward pass.

import torch 
from hrtx.ee_network import EarlyExitTransformer

# Input tensor - tokens int
token = torch.randint(0, 1000, (1, 10))

# Create the model
model = EarlyExitTransformer(
    dim=512, depth=6,
    num_tokens=1000, seq_len=10,
    heads=8, dim_head=64, num_robots=3
)


# Forward pass
output = model(token)

# Print the output shape
print(output.shape)

# Output: torch.Size([1, 10, 512])

MIMOTransformer

Multi-Input Multi-Output Transformer model for parallel robotic execution.

import torch 
from hrtx.mimo import MIMOTransformer

# Input tensor
x = torch.randn(1, 10, 512)
x = [x, x, x]

# Create the model
model = MIMOTransformer(
    dim=512, depth=6, heads=8, dim_head=64, num_robots=3
)


# Forward pass
output = model(x)

# Print the output shape
print(output)

# Output: torch.Size([1, 10, 512])

SAETransformer

SAE: Multiple Inputs for every after then an output head.

import torch
from hrtx.sae_transformer import SAETransformer

# Input tensor
x = torch.randn(1, 10, 512)
x = [x, x, x]

# Create the model
model = SAETransformer(
    dim=512, depth=6, heads=8, dim_head=64, num_robots=3
)


# Forward pass
output = model(x)

# Print the output shape
print(output)

# Output: torch.Size([1, 10, 512])

MIMMO

import torch
from hrtx.mimmo import MIMMO


# Usage of the MIMMO module
x = [torch.randint(0, 1000, (1, 10)) for _ in range(3)]


# Create the model
model = MIMMO(
    dim=512,
    depth=6,
    num_tokens=1000,
    seq_len=10,
    heads=8,
    dim_head=64,
    num_robots=3,
)

# Forward pass
output = model(x)

# Print the output shape
print(output[0].shape)

Hivemind Multi-Modality Transformer (HMMT) Model Architecture Specification

Objective:

Design the model architecture for a Hivemind Multi-Modality Transformer that can accept multi-modal inputs from 'x' number of robots and send instructions to all robots, a single robot, or any selected number of robots.

Features:

  1. Multi-Modal Embedding Layers:

    • Distinct embedding layers tailored for each modality (e.g., vision, audio, sensor data).
    • Fusion mechanisms to cohesively merge embeddings into a comprehensive representation.
  2. Dynamic Input Channels:

    • Ability to automatically adapt the number of input channels based on 'x' robots.
    • Channel attention mechanisms that assign weights to the significance of each robot's input.
  3. Hierarchical Attention Mechanisms:

    • Multilayer attention that can concentrate on specific modalities or individual robots.
    • Global attention modules for a holistic scene or context comprehension.
  4. Adaptive Computation:

    • Layers designed to modulate computations based on input intricacy.
    • Streamlined processing pathways for rudimentary tasks, with deeper computations allocated for intricate tasks.
  5. Output Decoders:

    • Multiple decoders tailored for various types of instructions (e.g., navigation, task-specific commands).
    • Multi-head output configuration for concurrent instruction formulation for diverse robots.
  6. Latency Optimization:

    • Fast-track routes for immediate instruction delivery.
    • Asynchronous processing units to handle non-immediate tasks.
  7. Robustness & Generalization:

    • Embedded mechanisms to ensure model resilience against diverse input types.
    • Capacity to handle and process noisy or unexpected inputs without faltering.
  8. Model Parallelism & Scalability:

    • Distributed model design to cater to a vast number of robot inputs efficiently.
    • Individual micro-models for each robot that operate concurrently.

Transformers, since their introduction in the "Attention is All You Need" paper, have revolutionized the field of deep learning. Initially designed for sequence-to-sequence tasks like language translation, they've now found relevance in a variety of domains, including computer vision, where they're termed as Vision Transformers (ViTs).

In the context of our HMMT, the transformer architecture serves as the backbone. Its self-attention mechanism allows it to weigh the significance of various inputs relative to each other, making it perfect for multi-modal data and inputs from multiple robots.

1. Transformers in Multi-Modal Embeddings:

For each modality (e.g., vision, audio, sensor data), we first transform raw inputs into embeddings. Transformers can be employed here in two main ways:

  • Sequential Transformers: Each modality's data, which can often be sequential (like a series of sensor readings or words in a command), is fed into a transformer. This transformer learns the inherent sequence patterns and produces a contextual embedding for the entire sequence.

  • Cross-Modal Attention: Once individual modalities have their embeddings, a higher-order transformer can be used to establish attention across modalities. This means the model can understand, for instance, that a visual input of a red light might be highly relevant to an audio input of a siren.

2. Dynamic Channel Adjuster with Transformers:

The ability to dynamically adjust to varying robot counts is crucial. Here, transformers play a pivotal role:

  • Channel-wise Self-Attention: For each robot input channel, a transformer layer assesses the importance of that channel in the context of all other channels. It provides a weighted representation, emphasizing more crucial channels and dampening less relevant ones.

3. Hierarchical Attention Mechanisms:

The strength of transformers lies in their ability to handle various scales of attention:

  • Local Attention: For tasks like image recognition in a visual modality, transformers can focus on local patterns (like the shape of an object).

  • Global Attention: For understanding the broader context (like the overall scene in a visual feed or the overarching command in a textual instruction), transformers can spread their attention globally.

By stacking multiple transformer layers, we form a hierarchy. The initial layers focus on local patterns, while deeper layers capture broader contexts. This hierarchical structure is beneficial for multi-modal data as it helps in bridging local features of one modality with global features of another.

4. ViTs in Multi-Modal Fusion Blocks:

Vision Transformers (ViTs) divide an image into fixed-size non-overlapping patches, linearly embed them, and then feed them into a standard transformer. In the context of HMMT:

  • Patch-based Embedding: For each robot's visual feed, ViTs can extract crucial visual patches, allowing the model to focus on significant parts of the visual data, like an object of interest or a particular gesture.

  • Fused Visual Representation: The output of the ViT for each visual feed is a rich representation that can be fused with embeddings from other modalities using subsequent transformer layers.

5. Latency Optimized Fast-Track Routes with Transformers:

For scenarios demanding immediate response, we introduce fast-track transformer layers:

  • Shallow Transformers: Instead of passing data through the entire depth of the model, shallow transformer layers can quickly process and produce outputs. These layers are trained to handle frequent and time-critical scenarios.

6. Multi-Headed Decoders with Transformer Heads:

The concept of multi-headed attention in transformers can be extended to our decoders:

  • Task-specific Heads: Each head can be tailored for a specific type of instruction. For instance, one head can focus on navigation, while another can handle task-specific commands.

  • Conditional Parallelism in Heads: Depending on the input, certain heads can be activated while others remain dormant. This dynamic activation ensures efficient computation.


Incorporating Model Parallelism with Transformers:

Given the inherent parallel nature of transformers, where each token or patch attends to every other token or patch, we can leverage model parallelism:

  • Micro-Transformers for Each Robot: Each robot's data is first processed by a dedicated micro-transformer. These micro-models run concurrently, ensuring scalability.

  • Distributed Attention Computation: The self-attention computation, which is quadratic with respect to the input length, can be distributed across multiple GPUs or TPUs.


Conclusion:

The Hivemind Multi-Modality Transformer, with its deep integration of transformers and vision transformers, stands poised to redefine multi-robot control systems. The architecture leverages the strengths of transformers to handle diverse modalities, dynamically adjust to varying robot inputs, and produce precise instructions. With the added benefits of model parallelism, the HMMT ensures scalability, making it a promising solution for future swarm robotics and large-scale multi-robot systems.

The architecture of the Hivemind Multi-Modality Transformer is conceptualized to stand at the forefront of multi-robot control, promising efficient, adaptive, and secure interactions across varied scenarios.

License

MIT