A system for just-in-time (JIT) LoRA training that modifies a running language model's weights mid-conversation on consumer Apple Silicon hardware. Using MLX-native autograd for gradient-based LoRA adaptation, the system — J.A.R.V.I.S., a voice-enabled AI assistant — updates its own weights after every response via background backpropagation.
Learning rate 10x higher than standard LoRA (5e-4 vs 5e-5): JIT learning needs convergence in ~4 epochs, not thousands of steps. Gradient clipping (1.0) prevents instability.
≥33% regularization ratio eliminates catastrophic forgetting: Below this threshold, the model overwrites core knowledge. At ≥33%, general knowledge is preserved at 100% (CI: [94.0%, 100.0%]).
mx.compile() hurts short training runs: The ~20s first-trace overhead is not amortized in <200 steps. Per-step time is ~390ms without compilation.
Batching doesn't help on Apple Silicon: Memory-bandwidth-limited, not compute-limited. Batch=8 takes 2.5s/step vs 0.42s/step for batch=1.
Structurally similar facts confuse small models: Deaths/obituaries (18.2%) all follow "[Person] died on [Date]" pattern. The model learns the category but fabricates dates. Distinctive patterns (Sports, Awards) achieve 85-100%.
Architecture
The training engine is pure MLX — nn.value_and_grad() for real autograd, Adam optimizer, cosine LR with early stopping. LoRA adapters are injected in-place into the model, so mlx_lm.stream_generate() automatically uses the updated weights with no special handling.
The ane_*.py files and bridge/ directory are experimental and not used for training. The initial approach attempted to run LoRA kernels directly on Apple's Neural Engine via the private AppleNeuralEngine.framework. While the forward kernels compile and run, ANE produces IOSurface-backed tensors that are opaque to any autograd system — making gradient-based training impossible through ANE alone.
All training in this project uses MLX autograd on GPU. The ANE code remains in the repo for a potential future hybrid inference path (see Section 8.2 of the paper), where ANE could accelerate LoRA forward passes during multi-agent inference while the GPU handles the base model. This path is speculative and has not been benchmarked.
If you're interested in ANE internals, the bridge is based on maderix/ANE (MIT License) and requires macOS 15+ on Apple Silicon. Build with cd src/bridge && make. But this is not required to run any of the experiments or use the training system.
Citation
bibtex
1@article{elbaz2026jitlora,
2 title={JIT LoRA: Real-Time Conversational Knowledge Injection on Apple Silicon via MLX},
3 author={Elbaz, E.},
4 year={2026},
5 url={https://github.com/eelbaz/jit-lora}
6}
