Post

Have you ever wondered how training dynamics differ between LLMs 🖋️ and Vision 👁️ models? We explore this and close the gap between VMs and LLMs in our #NeurIPS2024 paper "TrAct: Making First-layer Pre-Activations Trainable". Paper link 📜: arxiv.org/abs/2410.23970 Video link 🎥: youtu.be/ZjTAjjxbkRY 🧵

The big difference between LLMs and Vision models lies in the first layer: * in LLMs we update Embeddings (/activations) directly * but in Vision models we update the *weights* of the first layer, which causes indirect updates to the Activations (/embeddings)

We close this gap by proposing TrAct: we conceptually *Tr*ain *Act*ivations. While we can't train activations directly bc only weights are trainable, we formulate an optimization problem to find the optimal weights to match a GD step on the activations, and in closed-form modify the gradients resp.

Thus, we can effectively train the first-layer activations of a Vision model, with updates similar to those in the LLM Embedding layer.

Our implementation is efficient, only modifies the gradient in the backward, and is compatible with various optimizers. To use *TrAct*, just wrap your first layer in a "TrAct" module...