Make visualizing_derivatives snappy again

Test plan: local

Auditors: alpert