{"id":17141090,"url":"https://github.com/zphang/architecture-objective","last_synced_at":"2026-06-23T18:32:04.727Z","repository":{"id":116145647,"uuid":"504728307","full_name":"zphang/architecture-objective","owner":"zphang","description":null,"archived":false,"fork":false,"pushed_at":"2022-06-29T07:29:41.000Z","size":5987,"stargazers_count":0,"open_issues_count":0,"forks_count":0,"subscribers_count":2,"default_branch":"main","last_synced_at":"2025-09-02T13:39:42.058Z","etag":null,"topics":[],"latest_commit_sha":null,"homepage":null,"language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"apache-2.0","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/zphang.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":"CONTRIBUTING.md","funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null}},"created_at":"2022-06-18T02:54:21.000Z","updated_at":"2022-06-18T04:19:13.000Z","dependencies_parsed_at":null,"dependency_job_id":"485cfd69-b815-4a9e-a1d6-673bb356a4ad","html_url":"https://github.com/zphang/architecture-objective","commit_stats":null,"previous_names":[],"tags_count":0,"template":false,"template_full_name":null,"purl":"pkg:github/zphang/architecture-objective","repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/zphang%2Farchitecture-objective","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/zphang%2Farchitecture-objective/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/zphang%2Farchitecture-objective/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/zphang%2Farchitecture-objective/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/zphang","download_url":"https://codeload.github.com/zphang/architecture-objective/tar.gz/refs/heads/main","sbom_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/zphang%2Farchitecture-objective/sbom","scorecard":null,"host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":286080680,"owners_count":34702911,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2026-05-26T15:22:16.424Z","status":"online","status_checked_at":"2026-06-23T02:00:07.161Z","response_time":65,"last_error":null,"robots_txt_status":"success","robots_txt_updated_at":"2025-07-24T06:49:26.215Z","robots_txt_url":"https://github.com/robots.txt","online":true,"can_crawl_api":true,"host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":[],"created_at":"2024-10-14T20:24:20.258Z","updated_at":"2026-06-23T18:32:04.679Z","avatar_url":"https://github.com/zphang.png","language":"Python","funding_links":[],"categories":[],"sub_categories":[],"readme":"# What Language Model Architecture and Pretraining Objective Work Best for Zero-Shot Generalization?\n\nLarge pretrained Transformer language models have been shown to exhibit zero-shot generalization, i.e. they can perform a wide variety of tasks that they were not explicitly trained on. However, the architectures and pretraining objectives used across state-of-the-art models differ significantly, and there has been limited systematic comparison of these factors. In this work, we present a large-scale evaluation of modeling choices and their impact on zero-shot generalization. In particular, we focus on text-to-text models and experiment with three model architectures (causal/non-causal decoder-only and encoder-decoder), trained with two different pretraining objectives (autoregressive and masked language modeling), and evaluated with and without multitask prompted finetuning. We train models with over 5 billion parameters for more than 170 billion tokens, thereby increasing the likelihood that our conclusions will transfer to even larger scales. Our experiments show that causal decoder-only models trained on an autoregressive language modeling objective exhibit the strongest zero-shot generalization after purely unsupervised pretraining. However, models with non-causal visibility on their input trained with a masked language modeling objective followed by multitask finetuning perform the best among our experiments. We therefore consider the adaptation of pretrained models across architectures and objectives. We find that pretrained non-causal decoder models can be adapted into performant generative causal decoder models, using autoregressive language modeling as a downstream task. Furthermore, we find that pretrained causal decoder models can be efficiently adapted into non-causal decoder models, ultimately achieving competitive performance after multitask finetuning. \n\nFull paper is available at: https://arxiv.org/abs/2204.05832\n\n## Checkpoints\n// TODO (@thomasw21):  link to checkpoints.\n\n## How to cite\n\n    @article{wang2022language,\n      title={What Language Model Architecture and Pretraining Objective Work Best for Zero-Shot Generalization?},\n      author={Wang, Thomas and Roberts, Adam and Hesslow, Daniel and Scao, Teven Le and Chung, Hyung Won and Beltagy, Iz and Launay, Julien and Raffel, Colin},\n      journal={arXiv preprint arXiv:2204.05832},\n      year={2022}\n    }\n\n\n\n\n# T5X\n\nT5X is a modular, composable, research-friendly framework for high-performance,\nconfigurable, self-service training, evaluation, and inference of sequence\nmodels (starting with language) at many scales.\n\nIt is essentially a new and improved implementation of the\n[T5 codebase](https://github.com/google-research/text-to-text-transfer-transformer)\n(based on [Mesh TensorFlow](https://github.com/tensorflow/mesh)) in [JAX](https://github.com/google/jax) and [Flax](https://github.com/google/flax).\n\n## Installation\n\nNote that all the commands in this document should be run in the commandline of\nthe TPU VM instance unless otherwise stated.\n\n1.  Follow the\n    [instructions](https://cloud.google.com/tpu/docs/jax-quickstart-tpu-vm#install_the_google_cloud_sdk)\n    to set up a Google Cloud Platform (GCP) account and enable the Cloud TPU\n    API.\n\n    **Note:** While T5X works with GPU as well, we haven't heavily tested the\n    GPU usage.\n\n2.  Create a\n    [Cloud TPU VM instance](https://cloud.google.com/blog/products/compute/introducing-cloud-tpu-vms)\n    following\n    [this instruction](https://cloud.google.com/tpu/docs/jax-quickstart-tpu-vm#create-vm).\n    We recommend that you develop your workflow in a single v3-8 TPU (i.e.,\n    `--accelerator-type=v3-8`) and scale up to pod slices once the pipeline is\n    ready. In this README, we focus on using a single v3-8 TPU. See\n    [here](https://cloud.google.com/tpu/docs/system-architecture-tpu-vm) to\n    learn more about TPU architectures.\n\n3.  With Cloud TPU VMs, you ssh directly into the host machine of the TPU VM.\n    You can install packages, run your code run, etc. in the host machine. Once\n    the TPU instance is created, ssh into it with\n\n    ```sh\n    gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} --zone=${ZONE}\n    ```\n\n    where `TPU_NAME` and `ZONE` are the name and the zone used in step 2.\n\n4.  Install T5X and the dependencies. JAX and Gin-config need to be installed\n    from the source.\n\n    ```sh\n    git clone --branch=main https://github.com/google-research/t5x\n    cd t5x\n\n    python3 -m pip install -e . -f \\\n      https://storage.googleapis.com/jax-releases/libtpu_releases.html\n\n    ```\n\n\n5.  Create toogle Cloud Storage (GCS) bucket to store the dataset and model\n    checkpoints. To create a GCS bucket, see these\n    [instructions](https://cloud.google.com/storage/docs/creating-buckets).\n\n## Example: English to German translation\n\nAs a running example, we use the WMT14 En-De translation. The raw dataset is\navailable in TensorFlow Datasets as\n[\"wmt_t2t_translate\"](https://www.tensorflow.org/datasets/catalog/wmt_t2t_translate).\n\nT5 casts the translation task such as the following\n\n```py\n{'en': 'That is good.', 'de': 'Das ist gut.'}\n```\n\nto the form called \"text-to-text\":\n\n```py\n{'inputs': 'translate English to German: That is good.', 'targets': 'Das ist gut.'}\n```\n\nThis formulation allows many different classes of language tasks to be expressed\nin a uniform manner and a single encoder-decoder architecture can handle them\nwithout any task-specific parameters. For more detail, refer to the [T5 paper\n(Raffel et al. 2019)][t5_paper].\n\nFor a scalable data pipeline and an evaluation framework, we use\n[`SeqIO`](https://github.com/google/seqio), which was factored out of the [T5\nlibrary][t5_github]. A `seqio.Task` packages together the raw dataset, vocabulary,\npreprocessing such as tokenization and evaluation metrics such as\n[BLEU](https://aclanthology.org/P02-1040.pdf) and provides a\n[`tf.data`](https://www.tensorflow.org/guide/data) instance.\n\n[The T5 library][t5_github] provides a number of `seqio.Task`s that were used in the\n[T5 paper][t5_paper]. In this example, we use [wmt_t2t_ende_v003](https://github.com/google-research/text-to-text-transfer-transformer/blob/d81c0bab2a41b4d5dfbe4971de32f7d67df65f31/t5/data/tasks.py#L212).\n\n\n### Training\n\nTo run a training job, we use the `t5x/train.py` script.\n\n```sh\n# Model dir to save logs, ckpts, etc. in \"gs://model_dir\" format.\nMODEL_DIR=\"...\"\n\n# Data dir to save the processed dataset in \"gs://data_dir\" format.\nTFDS_DATA_DIR=\"...\"\nT5X_DIR=\"...\"  # directory where the T5X repo is cloned.\n\npython3 ${T5X_DIR}/t5x/train.py \\\n  --gin_file=\"t5x/examples/t5/t5_1_1/examples/t5_1_1_base_wmt_from_scratch.gin\" \\\n  --gin.MODEL_DIR=\"'${MODEL_DIR}'\" \\\n  --tfds_data_dir=${TFDS_DATA_DIR}\n```\n\nThe configuration for this training run is defined in the Gin file\n[t5_1_1_base_wmt_from_scratch.gin](t5x/examples/t5/t5_1_1/examples/t5_1_1_base_wmt_from_scratch.gin).\n[Gin-config](https://github.com/google/gin-config) is a library to handle\nconfigurations based on dependency injection. Among many benefits, Gin allows\nusers to pass custom components such as a custom model to the T5X library\nwithout having to modify the core library. The [custom\ncomponents](#custom-components) section shows how this is done.\n\nWhile the core library is independent of Gin, it is central to the examples we\nprovide. Therefore, we provide a short [introduction][gin-primer] to Gin in the\ncontext of T5X.  All the configurations are written to a file \"config.gin\" in\n`MODEL_DIR`. This makes debugging as well as reproducing the experiment much\neasier.\n\nIn addition to the `config.json`, `model-info.txt` file summarizes the model\nparameters (shape, names of the axes, partitioning info) as well as the\noptimizer states.\n\n\n\n#### TensorBoard\n\nTo monitor the training in [TensorBoard](https://www.tensorflow.org/tensorboard), it is much easier (due to\nauthentification issues) to launch the TensorBoard on your own machine and _not_ in\nthe TPU VM. So in the commandline where you ssh'ed into the TPU VM, launch the\nTensorBoard with the `logdir` pointing to the `MODEL_DIR`.\n\n```sh\n# NB: run this on your machine not TPU VM!\nMODEL_DIR=\"...\"  # Copy from the TPU VM.\ntensorboard --logdir=${MODEL_DIR}\n```\n\nOr you can launch the TensorBoard inside a Colab. In a Colab cell, run\n\n```python\nfrom google.colab import auth\nauth.authenticate_user()\n```\n\nto authorize the Colab to access the GCS bucket and launch the TensorBoard.\n\n```python\n%load_ext tensorboard\nmodel_dir = \"...\"  # Copy from the TPU VM.\n%tensorboard --logdir=model_dir\n```\n\nTODO(hwchung): Add tfds preparation instruction\n\n\n### Fine-tuning\n\nWe can leverage the benefits of self-supervised pre-training by initializing\nfrom one of our pre-trained models. Here we use the T5.1.1 Base checkpoint.\n\n```sh\n# Model dir to save logs, ckpts, etc. in \"gs://model_dir\" format.\nMODEL_DIR=\"...\"\n\n# Data dir to save the processed dataset in \"gs://data_dir\" format.\nTFDS_DATA_DIR=\"...\"\nT5X_DIR=\"...\"  # directory where the T5X repo is cloned.\n\npython3 ${T5X_DIR}/t5x/train.py \\\n  --gin_file=\"t5x/examples/t5/t5_1_1/examples/t5_1_1_base_wmt_finetune.gin\" \\\n  --gin.MODEL_DIR=\"'${MODEL_DIR}'\" \\\n  --tfds_data_dir=${TFDS_DATA_DIR}\n```\n\n**Note:** when supplying a string, dict, list, tuple value, or a bash variable\nvia a flag, you must put it in quotes. In the case of strings, it requires\n\"triple quotes\" (`\"'\u003cstring\u003e'\"`). For example:\n`--gin.utils.DatasetConfig.split=\"'validation'\"` or\n`--gin.MODEL_DIR=\"'${MODEL_DIR}'\"`.\n\nGin makes it easy to change a number of configurations. For example, you can\nchange the `partitioning.ModelBasedPjitPartitioner.num_partitions` (overriding\nthe value in\n[t5_1_1_base_wmt_from_scratch.gin](t5x/examples/t5/t5_1_1/examples/t5_1_1_base_wmt_from_scratch.gin))\nto chanage the parallelism strategy and pass it as a commandline arg.\n\n```sh\n--gin.partitioning.ModelBasedPjitPartitioner.num_partitions=8\n```\n\n\n\n\n\n\n\n\n### Evaluation\n\nTo run the offline (i.e. without training) evaluation, you can use `t5x/eval.py`\nscript.\n\n```sh\nEVAL_OUTPUT_DIR=\"...\"  # directory to write eval output\nT5X_DIR=\"...\"  # directory where the t5x is cloned, e.g., ${HOME}\"/t5x\".\nTFDS_DATA_DIR=\"...\"\nCHECKPOINT_PATH=\"...\"\n\npython3 ${T5X_DIR}/t5x/eval.py \\\n  --gin_file=\"t5x/examples/t5/t5_1_1/examples/t5_1_1_base_wmt_eval.gin\" \\\n  --gin.CHECKPOINT_PATH=\"'${CHECKPOINT_PATH}'\" \\\n  --gin.EVAL_OUTPUT_DIR=\"'${EVAL_OUTPUT_DIR}'\" \\\n  --tfds_data_dir=${TFDS_DATA_DIR}\n```\n\n\n### Inference\n\nTo run inference, you can use `t5x/infer.py` script. Here we use the same\n`seqio.Task`, but for inference we do not use the targets features other than\nlogging them alongside the prediction in a JSON file.\n\n```sh\nINFER_OUTPUT_DIR=\"...\"  # directory to write infer output\nT5X_DIR=\"...\"  # directory where the t5x is cloned, e.g., ${HOME}\"/t5x\".\nTFDS_DATA_DIR=\"...\"\nCHECKPOINT_PATH=\"...\"\n\npython3 ${T5X_DIR}/t5x/infer.py \\\n  --gin_file=\"t5x/examples/t5/t5_1_1/examples/t5_1_1_base_wmt_infer.gin\" \\\n  --gin.CHECKPOINT_PATH=\"'${CHECKPOINT_PATH}'\" \\\n  --gin.INFER_OUTPUT_DIR=\"'${INFER_OUTPUT_DIR}'\" \\\n  --tfds_data_dir=${TFDS_DATA_DIR}\n```\n\n\n\n## Custom components\n\n[The translation example](#example-english-to-german-translation) uses the\nencoder-decoder model that T5X provides as well as the dataset from the T5\nlibrary. This section shows how you can use your own dataset and a model and\npass via Gin.\n\n### Example: custom dataset in a user directory\n\nFor this example, we have the following directory structure with\n`${HOME}/dir1/user_dir` representing a user directory with custom components.\n\n```\n${HOME}\n└── dir1\n    └── user_dir\n        ├── t5_1_1_base_de_en.gin\n        └── tasks.py\n```\n\nAs an example, let's define a new dataset. Here we use the same Translation\ndataset but we define the translation task in the opposite direction, i.e.,\nGerman to English intead of English to German. We define this task in `tasks.py`\n\n```py\n# ${HOME}/dir1/user_dir/tasks.py\n\nimport functools\nimport seqio\nimport tensorflow_datasets as tfds\nfrom t5.evaluation import metrics\nfrom t5.data import preprocessors\n\nvocabulary = seqio.SentencePieceVocabulary(\n    'gs://t5-data/vocabs/cc_all.32000/sentencepiece.model', extra_ids=100)\noutput_features = {\n    'inputs': seqio.Feature(vocabulary=vocabulary),\n    'targets': seqio.Feature(vocabulary=vocabulary)\n}\n\nseqio.TaskRegistry.add(\n    'wmt_t2t_de_en_v003',\n    source=seqio.TfdsDataSource(tfds_name='wmt_t2t_translate/de-en:1.0.0'),\n    preprocessors=[\n        functools.partial(\n            preprocessors.translate,\n            source_language='de', target_language='en'),\n        seqio.preprocessors.tokenize,\n        seqio.CacheDatasetPlaceholder(),\n        seqio.preprocessors.append_eos_after_trim,\n    ],\n    metric_fns=[metrics.bleu],\n    output_features=output_features)\n```\n\nIn the Gin file, most of the settings are equivalent to those used in the\n[En-\u003eDe example](#example-english-to-german-translation). So we include the Gin\nfile from that example. To use \"wmt_t2t_de_en_v003\" task we just defined, we\nneed to import the task module \"tasks.py\". Note that we use a relative path\ndefined with respect to the user directory. This will be specified as a\nflag.\n\n```py\n# ${HOME}/dir1/user_dir/t5_1_1_base_de_en.gin\nfrom __gin__ import dynamic_registration\nimport tasks  # This imports the task defined in dir1/user_dir/tasks.py.\n\ninclude \"t5x-tmp/t5x/examples/t5/t5_1_1/examples/t5_1_1_base_wmt_from_scratch.gin\"\nMIXTURE_OR_TASK_NAME = \"wmt_t2t_de_en_v003\"\n```\n\nFinally, we launch training passing the user directory as a flag\n`gin_search_paths` such that the Gin file and python modules can be specified\nwith relative paths.\n\n```sh\nPROJECT_DIR=${HOME}\"/dir1/user_dir\"\nT5X_DIR=\"...\"  # directory where the t5x is cloned.\nTFDS_DATA_DIR=\"...\"\nMODEL_DIR=\"...\"\nexport PYTHONPATH=${PROJECT_DIR}\n\npython3 ${T5X_DIR}/t5x/train.py \\\n  --gin_search_paths=${PROJECT_DIR} \\\n  --gin_file=\"t5_1_1_base_de_en.gin\" \\\n  --gin.MODEL_DIR=\"'${MODEL_DIR}'\" \\\n  --tfds_data_dir=${TFDS_DATA_DIR}\n```\n\n\n\n## Released Checkpoints\n\nWe release the checkpoints for the T5.1.1 models in a native T5X format.\n\n* **t5.1.1.small** (~77 million parameters): [gs://t5-data/pretrained_models/t5x/t5_1_1_small/checkpoint_1000000](https://console.cloud.google.com/storage/browser/t5-data/pretrained_models/t5x/t5_1_1_small/checkpoint_1000000)\n* **t5.1.1.base** (~250 million parameters): [gs://t5-data/pretrained_models/t5x/t5_1_1_base/checkpoint_1000000](https://console.cloud.google.com/storage/browser/t5-data/pretrained_models/t5x/t5_1_1_base/checkpoint_1000000)\n* **t5.1.1.large** (~800 million parameters): [gs://t5-data/pretrained_models/t5x/t5_1_1_large/checkpoint_1000000](https://console.cloud.google.com/storage/browser/t5-data/pretrained_models/t5x/t5_1_1_large/checkpoint_1000000)\n* **t5.1.1.xl** (~3 billion parameters): [gs://t5-data/pretrained_models/t5x/t5_1_1_xl/checkpoint_1000000](https://console.cloud.google.com/storage/browser/t5-data/pretrained_models/t5x/t5_1_1_xl/checkpoint_1000000)\n* **t5.1.1.xxl** (~11 billion parameters): [gs://t5-data/pretrained_models/t5x/t5_1_1_xxl/checkpoint_1000000](https://console.cloud.google.com/storage/browser/t5-data/pretrained_models/t5x/t5_1_1_xxl/checkpoint_1000000)\n\nThese are converted from the public [Mesh TensorFlow\ncheckpoints](https://github.com/google-research/text-to-text-transfer-transformer/blob/main/released_checkpoints.md#t511)\n.\n\n\n\n## Compatibility with the Mesh TensorFlow checkpoints\nThe Mesh TensorFlow checkpoints trained using the [T5 library][t5_github] can be\ndirectly loaded into T5X. For example, we can rerun the fine-tuning example\ninitializing from the MTF checkpoint by changing the `INIT_CHECKPOINT` Gin\nmacro.\n\n```sh\n# Model dir to save logs, ckpts, etc. in \"gs://model_dir\" format.\nMODEL_DIR=\"...\"\n\n# Data dir to save the processed dataset in \"gs://data_dir\" format.\nTFDS_DATA_DIR=\"...\"\nT5X_DIR=\"...\"  # directory where the T5X repo is cloned.\n\npython3 ${T5X_DIR}/t5x/train.py \\\n  --gin_file=\"t5x/examples/t5/t5_1_1/examples/wmt19_ende_from_scratch.gin\" \\\n  --gin.MODEL_DIR=\"'${MODEL_DIR}'\" \\\n  --gin.MIXTURE_OR_TASK_NAME=\"'wmt_t2t_ende_v003'\" \\\n  --gin.INIT_CHECKPOINT=\"'gs://t5-data/pretrained_models/t5.1.1.base/model.ckpt-1000000'\" \\\n  --tfds_data_dir=${TFDS_DATA_DIR}\n```\n\nNote that restoring directly from the Mesh TensorFlow checkpoints can be\ninefficient if heavy model parallelism is used for large models. This is\nbecause each host loads the entire copy of the model first and then keep only\nthe relevant slices dictated by the model parallelism specification. If you have\nMesh TensorFlow checkpoints that you run often, we recommend converting the\ncheckpoints to T5X native format using\n[`Checkpointer.convert_from_tf_checkpoint`](https://github.com/google-research/t5x/blob/fba685d1d49bfb1000f37b5952a9a0533f24ed36/t5x/checkpoints.py#L886).\n\nTODO(hwchung): Add a conversion script.\n\n\n\n## Note\nThis is not an officially supported Google product\n\n[t5_paper]: https://arxiv.org/abs/1910.10683\n[t5_github]: https://github.com/google-research/text-to-text-transfer-transformer\n[gin-primer]: gin-primer.md\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fzphang%2Farchitecture-objective","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fzphang%2Farchitecture-objective","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fzphang%2Farchitecture-objective/lists"}