search

center-for-humans-and-machines / transformer-heads

Toolkit for attaching, training, saving and loading of new heads for transformer models
https://transformer-heads.readthedocs.io/en/latest/
MIT License
245 stars 21 forks source link

Question regarding the MLP Regression head architecture diagram provided #4

Closed ArchchanaKugathasan closed 3 months ago

ArchchanaKugathasan commented 3 months ago
image

  1. In the above image for MLP regression head, it has been shown there are two linear layers inside the regression head. Is this predefined? or can we modify the number of linear layers? if the number of linear layers can be defined how can we specify it in the head_configs?

  2. The value 4096 x 128 in the first linear layer, does 4096 mean the hidden size pre-defined? or is it the hidden size we define in head_configs?

  3. The value 128 x 5 in the second linear layer, does 5 mean the number of outputs from the regression head? if we expect just one output does that need to be 128 x 1 ??

yannikkellerde commented 3 months ago
  1. This is completely custom. The image on the readme only shows a random example of how you could design your heads. It is defined via the head config. Specify the depth of your head via num_layers (1 for linear) and the width using hidden_size. See the docs here https://github.com/center-for-humans-and-machines/transformer-heads/blob/2515241b9420360ae5843c96b05d947bb8a3d1fb/transformer_heads/config.py#L19-L40).
  2. The 4096 depends on the base transformer architecture. In this case, the transformer used has a hidden size of 4096 for each token (You can find the hidden size of your preferred LLM in it's config. E.g. here for llama 3). The hidden size of the head is 128 and specified in the HeadConfig.
  3. Yes. Specify the number of outputs using num_outputs.

Also check this notebook to get a better understanding of HeadConfigs.

ArchchanaKugathasan commented 2 months ago

Thank you for the detailed explanation :)