def get_differential_lr_param_group(param_groups, lrs):  
  """Assigns different learning rates to different parameter groups.

  Discriminative fine-tuning, where different layers of the network have different learning rates, is first proposed in
  `Jeremy Howard and Sebastian Ruder. 2018. Universal language model fine-tuning for text classification. 
  https://arxiv.org/pdf/1801.06146.pdf.` It has been found to stabilize training and speed up convergence.

  Args: 
    param_groups: a list of parameter groups (each of which is a list of parameters)
        param group should look like:
        [ 
          [param1a, param1b, ..]  <-- parameter group 1
          [param2a, param2b, ..]  <-- parameter group 2
          ...
        ]
    lrs: a list of learning rates you want to assign to each of the parameter groups
        lrs should look like
        [
          lr1, <-- learning rate for parameter group 1
          lr2, <-- learning rate for parameter group 2
          ...
        ]

  Returns: 
    parameter groups with different learning rates that you can then pass into an optimizer
  """
  assert len(param_groups) == len(lrs), f"expect the learning rates to have the same lengths as the param_group length, instead got {len(param_groups)} and {len(lrs)} as lengths respectively"

  param_groups_for_optimizer = []
  for i in range(len(param_groups)): 
    param_groups_for_optimizer.append({
      'params': param_groups[i],
      'lr': lrs[i]
    })
  return param_groups_for_optimizer