google-deepmind / google-deepmind/optax
[Feature Request] Add NEAT conflict-aware gradient transformation to optax.contrib
- Dominant language
- Python
- Stars
- 2.3k
- Forks
- 369
- Avg merge
- 10h 15m
- Merged PRs (30d)
- 7
Description
## Summary
Hi Optax team,
I'd like to propose adding NEAT, Nash-Equilibrium Adaptive Training, as a conflict-aware gradient transformation and contributed optimizer in `optax.contrib`.
NEAT detects directional conflict between the current gradient and a stateful opponent proxy, such as the previous momentum direction, and partially removes the conflicting component before momentum accumulation.
The reference implementation is available here:
https://github.com/ItCodinTime/NEAT
NEAT is also listed in the official Keras ecosystem as a Keras-first optimizer library for conflict-aware neural network training:
https://keras.io/getting_started/ecosystem/
I would be happy to implement and maintain a JAX/Optax-native version if this fits the scope of `optax.contrib`.
## Motivation
Most standard optimizers determine update magnitude and momentum but do not explicitly reason about temporal directional conflict between the current gradient and an accumulated optimization signal.
NEAT targets settings where gradient direction changes are meaningful, including noisy optimization, oscillatory objectives, multi-domain training, and objectives with genuine gradient disagreement.
The goal is not to claim universal superiority over AdamW, SGD, or other established optimizers. Instead, NEAT provides an explicit mechanism for detecting and selectively correcting conflicting gradient components.
This also exposes useful optimization diagnostics such as conflict rate and correction magnitude.
## Core update
Given current gradient `g_t` and opponent proxy `o_t`:
```python
conflict = relu(-cosine_similarity(g_t, o_t))
projection = (
dot(g_t, o_t) /
(squared_norm(o_t) + eps)
) * o_t
correction = -alpha * conflict * projection
g_corrected = g_t + correction
Contributor guide
Research direction
Start by reviewing the existing optax.contrib entry points and the linked NEAT reference implementation. Compare its conflict detection, projection, correction, opponent state, and diagnostics with Optax conventions; done means a JAX/Optax-native contributed optimizer and gradient transformation are defined with appropriate validation and tests.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- machine-learning
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100