Package: gato-torch
Gato: A Generalist Agent[Deepmind Publication][arXiv Paper]aper.Installation$pipinstallgato-torchimporttorchfromgatoimportGato#create model instancegato=Gato(input_dim=768,img_patch_size=16,token_sequence_length=1024,vocabulary_size=32000,actions_size=1024,continuous_values_size=1024,num_transformer_blocks=8,num_attention_heads=24,layer_width=768,feedforward_hidden_size=3072,key_value_size=32,dropout_rate=0.1,num_group_norm_groups=32,discretize_depth=128,local_position_encoding_size=512,max_seq_len=8192)#fake inputs for Gatoinput_dim=config.input_diminput_ids=torch.cat([torch.rand((1,1,input_dim))for_inrange(20)]+# 20 image patches[torch.full((1,1,input_dim),0.25),#continous value]torch.full((1,1,input_dim),624.0)]+#discrete (actions, texts)[torch.rand((1,1,input_dim))for_inrange(20)]+#20 image patches[torch.full((1,1,input_dim),0.12),#continous valuetorch.full((1,1,input_dim),295.0)],#discrete( actions, text)dim=1)encoding=torch.tensor([[0,0,0,0,1,2,0,0,0,0,1,2]])row_pos=(torch.tensor([[0.00,0.25,0.50,0.75,0,0,0.00,0.25,0.50,0.75,0,0]]),# pos_fromtorch.tensor([[0.25,0.50,0.75,1.00,0,0,0.25,0.50,0.75,1.00,0,0]])# pos_to)col_pos=(torch.tensor([[0.00,0.00,0.00,0.80,0,0,0.00,0.00,0.00,0.80,0,0]]),# pos_fromtorch.tensor([[0.20,0.20,0.20,1.00,0,0,0.20,0.20,0.20,1.00,0,0]])# pos_to)obs=(torch.tensor([[0,1,2,19,20,21,0,1,2,19,20,21]]),# obs tokentorch.tensor([[1,1,1,1,1,0,1,1,1,1,1,0]])# obs token masking (for action tokens))hidden_states=gato((input_ids,(encoding,row_pos,col_pos),obs))Dataset and Model ArchitecturePaper ReviewsFull Episode SequenceArchitecture VariantsAppendix C.1. Transformer HyperparametersIn the paper, Deepmind tested Gato with 3 architecture variants,1.18B,364M, and79M.I have named them aslarge(),baseline()andsmall()respectively inGatoConfig.HyperparametersLarge(1.18B)Baseline(364M)Small(79M)Transformer blocks24128Attention heads161224Layer width20481536768Feedforward hidden size819261443072Key/value size12812832Residual EmbeddingAppendix C.2. Embedding FunctionThere are no mentions that how many residual networks must be stacked for token embeddings.Therefore, I remain configurable inGatoConfig.Whatever how many residual layers are existing, full-preactivation is a key.The blocks are consisted of:Version 2 ResNet architecture (based on ResNet50V2)GroupNorm (instead of LayerNorm)GeLU (instead of ReLU)Position EncodingsAppendix C.3. Position EncodingsPatch Position EncodingsLikeVision Transformer (ViT)by Google, Gato takes the input images as raster-ordered 16x16 patches.Unlike the Vision Transformer model, however, Gato divides its patch encoding strategy into 2 phases, training and evaluation.For high-performance computation in TensorFlow, I have used the following expressions.$C$ and $R$ mean column and row-wise, and $F$ and $T$ meanfromandtorespectively.$$
\begin{align}
v^R_F &= \begin{bmatrix}
0 & 32 & 64 & 96
\end{bmatrix} \
v^R_T &= \begin{bmatrix}
32 & 64 & 96 & 128
\end{bmatrix} \
v^C_F &= \begin{bmatrix}
0 & 26 & 51 & 77 & 102
\end{bmatrix} \
v^C_T &= \begin{bmatrix}
26 & 51 & 77 & 102 & 128
\end{bmatrix} \
\
P_R &= \begin{cases}
\mathsf{if} \ \mathsf{training} & v^R_F + \mathsf{uniform}(v^R_T - v^R_F) \
\mathsf{otherwise} & \mathsf{round}(\frac{v^R_F + v^R_T}{2})
\end{cases} \
P_C &= \begin{cases}
\mathsf{if} \ \mathsf{training} & v^C_F + \mathsf{uniform}(v^C_T - v^C_F) \
\mathsf{otherwise} & \mathsf{round}(\frac{v^C_F + v^C_T}{2})
\end{cases} \
\
E^R_P &= P_R \cdot 1^{\mathsf{T}}_C \
E^C_P &= 1^{\mathsf{T}}_R \cdot P_C \
\
\therefore E &= E_I + E^R_P + E^C_P
\end{align}
$$Local Observation Position EncodingsIn the definition of Appendix B., text tokens, image patch tokens, and discrete & continuous values are observation tokensWhen Gato receives those values, they must be encoded with their own (local) time steps.ContributingWe welcome all contributions, please either submit a pull request or submit issues in the Agora discordLicenseLicensed under theMIT license.Roadmap:Get functional prototypeIntegrate ALIBI, multi query, qk norm and other SOTA stuffintegrate action tokens