Add support to dropout. #29
Conversation
…ble dropout during evaluation.
aliciafmachado
changed the title
There was a problem hiding this comment.
There are some commits going back and forth on some things that I did not fully understood at first, so feel free to squash the commits before merging to main to avoid confusion. Otherwise I can recreate the pull request and fix the commit history.
I also have a few questions / discussion topics:
- I added support to dropout but we need something to manage random seeds so that we can seed properly. Should we create an issue for that?
- I tried to add dropout based on T5 architecture, but I decided to not add it after the FF layer and in the output. For the FF, I don't think it makes sense since we have a single layer and we apply dropout before the residual connection after the FF network. For the output, I don't see any additional computations after getting out of the stack, so I think it would only increase noise if we were to add another dropout there (I also did not see an additional dropout on the output for the haiku implementation linked in the issue to add dropout).
| @@ -225,13 +229,20 @@ function gelu(x: tf.Tensor) { | |||
| export function computeAttnHead( | |||
| spec: AttnHeadComputeSpec, | |||
| params: AttnHeadParams<TensorKind>, | |||
| seqInput: GTensor<'batch' | 'pos' | 'inputRep'> | |||
| seqInput: GTensor<'batch' | 'pos' | 'inputRep'>, | |||
| evalMode: boolean = false | |||
There was a problem hiding this comment.
lets drop evalMode flag, and just depend on the spec having dropoutRate set different at eval vs inference time.
| export function dropout<G extends string, D extends G>( | ||
| dropoutRate: number, | ||
| g: GTensor<G>, | ||
| deterministic: boolean, |
There was a problem hiding this comment.
Lets remove deterministic, and just check if rate is 0.
| let unNormedSeqOuput = inputToFF | ||
| .contract(ff.w, ['inputRepToFF']) | ||
| .pointwiseAdd(ff.bIn) | ||
| .applyPointWiseTfFn(gelu) | ||
| .pointwiseAdd(ff.bOut); | ||
|
|
||
| // Dropout before layer norm and residual connection. | ||
| let unNormedSeqOuputAfterDropout = unNormedSeqOuput; |
There was a problem hiding this comment.
Lets use this: https://github.com/Shivanandroy/simpleT5 as the reference for where to put it for T5. And maybe name this function computeT5AttnHead, and then later we can make a gpt2 one.
There was a problem hiding this comment.
I suggest we use the T5 implementation in Jax (T5X): https://github.com/google-research/t5x/blob/705247b743d26a33d0c058b41c72ad030e51891b/t5x/examples/t5/network.py#L222
| const layerSpec: transformer.TransformerParamLayerSpec = { | ||
| nHeads: 1, | ||
| hasPosEncoding: true, | ||
| computeSpec: { residuals: true, dropoutRate: 0.1 }, |
There was a problem hiding this comment.
Maybe add one test also for dropout rate of 1, and then test that loss doesn't decrease.
…to follow the T5X implementation.
Intended to resolve Issue: #1