Package: kosmosx
Kosmos-X: Advanced Multi-Modality AI Model 🚀🌌InstallationpipinstallkosmosxUsageimporttorchfromkosmosx.modelimportKosmos# Create a sample text token tensortext_tokens=torch.randint(0,32002,(1,50),dtype=torch.long)# Create a sample image tensorimages=torch.randn(1,3,224,224)# Instantiate the modelmodel=Kosmos()text_tokens=text_tokens.long()# Pass the sample tensors to the model's forward functionoutput=model.forward(text_tokens=text_tokens,images=images)# Print the output from the modelprint(f"Output:{output}")Trainingaccelerate configthen:accelerate launch train_distributed.pyGet InvolvedWe're just at the beginning of our journey. As we continue to develop and refine Kosmos-X, we invite you to join us. Whether you're a developer, researcher, or simply an enthusiast, your insights and contributions can help shape the future of Kosmos-X.Contributing to Kosmos-XWe are thrilled to invite you to be a part of the Kosmos-X project. This is not just an open source project but a community initiative, and we value your expertise and creativity. To show our appreciation, we have instituted a unique rewards system that directly compensates contributors from the revenue generated by the Kosmos-X API.Why ContributeContributing to Kosmos-X not only enhances your skills and profile but also comes with financial rewards. When you contribute code, documentation, or any form of improvement to the Kosmos-X project, you are adding value. As such, we believe it's only fair that you share in the rewards.The modelKOSMOS-1 uses a decoder-only Transformer architecture based onMagneto (Foundation Transformers), i.e. an architecture that employs a so called sub-LN approach where layer normilization is added both before the attention module (pre-ln) and afterwards (post-ln) combining the advantages that either approaches have for language modelling and image understanding respectively. The model is also initialized according to a specific metric also described in the paper, allowing for more stable training at higher learning rates.They encode images to image features using a CLIP VIT-L/14 model and use aperceiver resamplerintroduced inFlamingoto pool the image features from256 -> 64tokens. The image features are combined with the token embeddings by adding them to the input sequence surrounded by special tokens<image>and</image>. An example is
 <image> image_features </image> text 
. This allows image(s) to be interwoven with text in the same sequence.We follow the hyperparameters described in the paper visible in the following image:DetailsModel (decoder)We use the torchscale implementation of the decoder-only Transformer architecture from Foundation Transformers:fromtorchscale.architecture.configimportDecoderConfigfromtorchscale.architecture.decoderimportDecoderconfig=DecoderConfig(decoder_layers=24,decoder_embed_dim=2048,decoder_ffn_embed_dim=8192,decoder_attention_heads=32,dropout=0.1,activation_fn="gelu",attention_dropout=0.1,vocab_size=32002,subln=True,# sub-LN approachxpos_rel_pos=True,# rotary positional embeddingsmax_rel_pos=2048)decoder=Decoder(config,embed_tokens=embed,embed_positions=embed_positions,output_projection=output_projection)CLIP VIT-L/14For the image model (CLIP VIT-L/14) we use a pretrained OpenClip model:fromtransformersimportCLIPModelclip_model=CLIPModel.from_pretrained("laion/CLIP-ViT-L-14-laion2B-s32B-b82K").vision_model# projects image to [batch_size, 256, 1024]features=clip_model(pixel_values=images)["last_hidden_state"]Perceiver ResamplerWe follow the default hyperparams for the perceiver resampler as no hyperparams are given in the paper:fromflamingo_pytorchimportPerceiverResamplerperceiver=PerceiverResampler(dim=1024,depth=2,dim_head=64,heads=8,num_latents=64,num_media_embeds=256)# projects image features to [batch_size, 64, 1024]self.perceive(images).squeeze(1)Because the model expects a hidden dimension of2048, we use ann.Linearlayer to project the image features to the correct dimension and initialize it according to Magneto's initialization scheme:image_proj=torch.nn.Linear(1024,2048,bias=False)torch.nn.init.normal_(image_proj.weight,mean=0,std=2048**-0.5)scaled_image_features=image_proj(image_features)TokenizerThe paper describes aSentencePiecewith a vocabulary of64007tokens. For simplicity (as we don't have the training corpus available), we use the next best open-source alternative which is the pretrainedT5-large tokenizerfrom HuggingFace. This tokenizer has a vocabulary of32002tokens.fromtransformersimportT5Tokenizertokenizer=T5Tokenizer.from_pretrained("t5-large",additional_special_tokens=["<image>","</image>"],extra_ids=0,model_max_length=1984# 2048 - 64 (image features))We then embed the tokens with ann.Embeddinglayer. We actually use abnb.nn.Embeddingfrombitandbyteswhich allows us to use 8-bit AdamW later.importbitsandbytesasbnbembed=bnb.nn.Embedding(32002,# Num embeddings2048,# Embedding dimpadding_idx)For positional embeddings, we use:fromtorchscale.component.embeddingimportPositionalEmbeddingembed_positions=PositionalEmbedding(2048,# Num embeddings2048,# Embedding dimpadding_idx)Also, we add an output projection layer to project the hidden dimension to the vocabulary size and initialize it according to Magneto's initialization scheme:output_projection=torch.nn.Linear(2048,32002,bias=False)torch.nn.init.normal_(output_projection.weight,mean=0,std=2048**-0.5)Decoder changesI had to make some slight changes to the decoder to allow it to accept already embedded features in the forward pass. This was necessary to allow the more complex input sequence described above. The changes are visible in the following diff in line 391 oftorchscale/architecture/decoder.py:+if kwargs.get("passed_x", None) is None:+    x, _ = self.forward_embedding(+        prev_output_tokens, token_embeddings, incremental_state+    )+else:+    x = kwargs["passed_x"]-x, _ = self.forward_embedding(-    prev_output_tokens, token_embeddings, incremental_state-)Dataset StrategyHere is a markdown table with metadata for the datasets mentioned in the paper:DatasetDescriptionSizeLinkThe PileDiverse English text corpus800 GBHuggingfaceCommon CrawlWeb crawl data-Common CrawlLAION-400MImage-text pairs from Common Crawl400M pairsHuggingfaceLAION-2BImage-text pairs from Common Crawl2B pairsArXivCOYOImage-text pairs from Common Crawl700M pairsGithubConceptual CaptionsImage-alt text pairs15M pairsArXivInterleaved CC DataText and images from Common Crawl71M docsCustom datasetStoryClozeCommonsense reasoning16k examplesACL AnthologyHellaSwagCommonsense NLI70k examplesArXivWinograd SchemaWord ambiguity273 examplesPKRR 2012WinograndeWord ambiguity1.7k examplesAAAI 2020PIQAPhysical commonsense QA16k examplesAAAI 2020BoolQQA15k examplesACL 2019CBNatural language inference250 examplesSinn und Bedeutung 2019COPACausal reasoning1k examplesAAAI Spring Symposium 2011RelativeSizeCommonsense reasoning486 pairsArXiv 2016MemoryColorCommonsense reasoning720 examplesArXiv 2021ColorTermsCommonsense reasoning320 examplesACL 2012IQ TestNonverbal reasoning50 examplesCustom datasetCOCO CaptionsImage captioning413k imagesPAMI 2015Flickr30kImage captioning31k imagesTACL 2014VQAv2Visual QA1M QA pairsCVPR 2017VizWizVisual QA31k QA pairsCVPR 2018WebSRCWeb QA1.4k examplesEMNLP 2021ImageNetImage classification1.28M imagesCVPR 2009CUBImage classification200 bird speciesTOG 2011TodoImplement tokenizer for multi-modal processingRefactor training scriptTrain 7BLicense