object_argsort_autoregressive

Experiments evaluating learning curves of Abstractor and Transformer models on an autoregressive object sorting task. N = 48 'objects' are generated as 12-dimensional gaussian vectors. They are generated via two sets of 'attribute objects' which each have an associated ordering. The first set of objects is 4-dimensional and has 4 possible values (each of which is a gaussian vector). The second set of objects is 8-dimensional and has 12 possible values (each of which is a gaussian vector). The 48 objects are formed by a cartesian product of the two sets of objects/attributes, where the first set of attributes forms the primary key and the second set of attributes forms the secondary key. This forms a strict ordering on the 48 objects. Models are trained to sort random permutations of the objects autoregressively. The task is to predict the `argsort` of sequences of 10 randomly permuted objects. For each model, learning curves are evaluated by training on subsets of the training data of varying size and evaluating on a hold-out test set. A RelationalAbstractor, SimpleAbstractor, and a Transformer model are evaluated. A second sorting task is formed by shuffling the order underlying the first attribute, and keeping the second attribute's ordering fixed. We evaluate pre-training on this task. When pre-training, the pre-training dataset contains 1000 sequences. The transformer and abstractor have hyper-parameters chosen such that they have a roughly equal parameter count (in fact, the transformer is larger): transformer_kwargs = dict( num_layers=4, num_heads=2, dff=64, input_vocab='vector', target_vocab=seqs_length+1, output_dim=seqs_length, embedding_dim=64) param count = 469,898 ---------- rel_abstracter_kwargs = dict( num_layers=2, num_heads=2, dff=64, input_vocab='vector', target_vocab=seqs_length+1, output_dim=seqs_length, embedding_dim=64, rel_attention_activation='softmax' ) param count = 386,954 ---------- simple_abstractor_kwargs = dict( embedding_dim=64, input_vocab='vector', target_vocab=seqs_length+1, output_dim=seqs_length, abstractor_kwargs=dict(num_layers=1, num_heads=4, dff=64, use_pos_embedding=False, mha_activation_type='softmax'), decoder_kwargs=dict(num_layers=1, num_heads=4, dff=64, dropout_rate=0.1)) param count = 219,210 ---------- symbolic_abstracter_kwargs = dict( num_layers=2, num_heads=2, dff=64, input_vocab='vector', target_vocab=seqs_length+1, output_dim=seqs_length, embedding_dim=64, rel_attention_activation='softmax' ) param count = 386,954 ---------- ablation_abstractor_kwargs = dict( num_layers=2, num_heads=2, dff=64, input_vocab='vector', target_vocab=seqs_length+1, output_dim=seqs_length, embedding_dim=64, use_self_attn=True, use_encoder=True, mha_activation_type='softmax' ) param count = 386,954