What's the role of PRNGKey in Haiku transforms?

[EngineerKhan]

I have seen Haiku transforms where we have the provision of adding an optional PRNGKey. I wonder what job does it serve?

[tomhennigan]

In JAX prng keys are used to create random samples. You can find out more about them in this guide.

In Haiku, hk.transform (and the variants of it) produce two functions: init and apply. Both of these functions accept random keys, and Haiku then makes new keys available inside your function via hk.next_rng_key().

You can haiku given rng keys for any purpose you prefer, but typically rng keys are required for init since initial parameter values are random samples, and occasionally needed in apply when you need random samples as part of your forward pass (e.g. for dropout layers).

PS feel free to open haiku specific questions on the dm-haiku github page.

[EngineerKhan]

typically rng keys are required for init since initial parameter values are random samples

I have one follow-up question: Which distribution is used by default, and can we override it? I can't seem to see any such a thing, sorry.

[tomhennigan]

By default Haiku uses a truncated normal distribution with stddev=1/sqrt(prod(fan_in)) for linear/conv. This should be documented for modules that ship with Haiku (e.g. see the doc for w_init on hk.Linear).

Most modules take a w_init (and b_init etc) argument in their constructor allowing you to override with your own initializer for each parameter. You can see an example of us setting a different initializer in our transformer example.

Haiku ships with a number of common initialisers that you can use (or you can write your own).

[EngineerKhan]

Thank you so much. One last question: does Flax (Linen) also work the same way? I have seen its init() also using the PRNG keys. I am especially curious about the Module class's init().

[tomhennigan]

I am not very familiar with Flax (so I might be wrong here), but I suspect the RNG argument to init is for the same purpose.

FYI It is quite possible that they have chosen different default initialisers vs. Haiku (e.g. they may not use truncated normal for nn.Dense).

Flax modules seem to also support {name}_init args (e.g. nn.Dense(.., kernel_init=X)) and you can see one of their transformer examples overriding the default initializer.

[EngineerKhan]

I am really grateful for your detailed answers.