Residual Gradient Scale Pre Vs Post Norm

Problem

Toy model of how residual depth scales gradients. For each block scale s, the pre-norm multiplier is (1 + s) because the identity path is untouched, while the post-norm multiplier is s / (1 + s) because the block output is renormalized before it is added back. Multiply the multipliers across all blocks.

The signature is residual_grad_scale(scales). Return [pre_norm_scale, post_norm_scale].

Hint

An empty stack gives the identity scale of 1.0 for both orderings.

Test cases

case 1

in: [[0.1,0.1,0.1]]

expected: [1.331,0.0007513148]

case 2

in: [[0.5]]

expected: [1.5,0.3333333333]

case 3

in: [[1,1]]

expected: [4,0.25]

case 4

in: [[]]

expected: [1,1]

Starter codePython

def residual_grad_scale(scales):
    # Your code here
    pass

4 test cases · run your code against them in the editor