Straightening Out the Straight-Through Estimator: Overcoming Optimization Challenges in Vector Quantized Networks.