search

saprmarks / dictionary_learning

MIT License
142 stars 37 forks source link

add an implementation of ghost gradients #3

Closed g-w1 closed 9 months ago

g-w1 commented 9 months ago

Implements Ghost Gradients from https://transformer-circuits.pub/2024/jan-update/index.html#dict-learning-resampling

Thanks @woog97 for pair programming with me on this!

g-w1 commented 9 months ago

Sounds good; this matches what I'm seeing in my experiments

On Fri, Feb 2, 2024 at 6:29 PM saprmarks @.***> wrote:

@.**** commented on this pull request.

In training.py https://github.com/saprmarks/dictionary_learning/pull/3#discussion_r1476870992 :

 # set up optimizer and scheduler

optimizer = ConstrainedAdam(ae.parameters(), ae.decoder.parameters(), lr=lr) def warmup_fn(step):

  • if step % resample_steps < warmup_steps:
  • if resample_steps and step % resample_steps < warmup_steps:

I think we should do a warmup at the start of training -- even with ghost grads, I expect that the first few gradient steps will kill a bunch of neurons without warming up first.

— Reply to this email directly, view it on GitHub https://github.com/saprmarks/dictionary_learning/pull/3#discussion_r1476870992, or unsubscribe https://github.com/notifications/unsubscribe-auth/AOA23ZNFYZ6YSFOO56GUFL3YRVZGBAVCNFSM6AAAAABCM2M46GVHI2DSMVQWIX3LMV43YUDVNRWFEZLROVSXG5CSMV3GSZLXHMYTQNRQGYZDMNRWGE . You are receiving this because you authored the thread.Message ID: @.***>