Junior+ ML Engineer (Model Architecture, Training and Edge Optimisation)

Full-time
Salary: $800–$1500/month, reviewed after 3 months based on delivery

About CubicMind AI
CubicMind AI Ltd is a UK deep tech startup building a small plug-in device that operates any computer like a person: it reads the screen, listens to the user, and clicks and types to complete tasks. Everything runs on the device itself, with no cloud.

https://cubicmind-ai.com
 

We design our own models rather than wrapping someone else's. That includes transformer and non-transformer architectures, built around what runs fast on edge hardware, and sometimes trained from scratch.

About the role
You will work on designing, training and optimising the models that power the device: screen understanding, action prediction and efficient on-chip inference.

This is research-heavy engineering. You will work on model architecture, training runs, attention and memory bottlenecks, low-bit compression, and getting models to run correctly on real hardware.
 

What you will do
- Design and implement new model architectures in PyTorch, both transformer-based and beyond
- Train models, sometimes from scratch, and run experiments to compare architectures
- Build training and evaluation pipelines on screen-understanding and computer-control data
- Profile latency, memory use and KV-cache behaviour on long, high-resolution screen inputs
- Compress models for edge chips using quantisation (INT8/INT4), pruning and distillation
- Export models to low-level runtimes and fix unsupported operators
- Work with our C++ team to get models running efficiently on the device
 

Requirements

- Deep architectural understanding of Transformers and Attention: you know the mathematical and operational differences between MHA, MQA, GQA, FlashAttention, and linear attention variants; you understand KV-cache scaling and memory access bottlenecks. 

- Low-level PyTorch fluency: writing custom PyTorch modules, implementing non-standard forward/backward logic, vectorised tensor operations, and profiling memory/compute bottlenecks using PyTorch Profiler / Nsight.

- Experience with model training & convergence: stability tricks, learning rate schedules, gradient clipping, mixed-precision training (BF16/FP16), and diagnosing loss spikes or exploding/vanishing gradients.

- Rigorous mathematical foundation: comfortable with high-dimensional linear algebra, matrix calculus, multivariable optimization, and probability theory.

- Practical understanding of model compression: Post-Training Quantization (PTQ) vs Quantization-Aware Training (QAT), outlier activations, weight-only vs weight-and-activation quantization (INT8, INT4, FP8).

- You read papers, implement ideas yourself, and debug by inspecting real activations rather than guessing

- Сommit full-time
 

Nice to have
- Experience with vision transformers or vision-language models
- Interest in non-transformer architectures (state-space models, linear attention, recurrent designs)
- Experience exporting to ONNX, TensorRT, OpenVINO or NPU toolchains
- Basic C++, enough to read inference code
- Degree in Applied Mathematics, Computer Science or Physics
 

What we offer
- Real ownership of the core models in a hardware product built from the ground up
- Work directly with the founder, with no bureaucracy
- $800–$1,200/month depending on experience, reviewed after 3 months based on results
- Fully remote, judged on results, not hours logged