A. Causal Modeling Foundations

flowchart LR
    U1(("U₁")) --> X1(("X₁"))
    U2(("U₂")) --> X2(("X₂"))
    X1 --> X2
    U3(("U₃")) --> Y(("Y"))
    X1 --> Y
    X2 --> Y

B. Towards Causal Deep Learning

Classical causal models are powerful precisely because their assumptions are explicit. Linear SEMs expose interpretable path coefficients, Bayesian networks support exact of approximate inference on compact graphs, and additive-noise models can identify directions under suitable distributional assumptions. Their common weakness is that the mechanisms are usually linear, table-based, or hand-designed. As the variables become high-dimensional and heterogeneous, these mechanisms struggle to represent nonlinear dependencies, multimodal uncertainty, latent factors, and complex observation process. Inference can also become expensive: exact graphical-model algorithms scale poorly with graph width, while direct inversion of a high-dimensional mechanism is rarely available.

Deep learning offers a natural response to these limitations. Neural networks provide flexible approximations to nonlinear structural assignments. Variational inference (Jordan et al. 1999) turns hidden-state estimation into an amortized encoder, and hierarchical generative models distribute information across multiple scales (…). Instead of specifying eery conditional density manually, a model can learn a mechanism from data while retaining a causal graph that determines which mechanisms may respond to an intervention. This is especially important when the outcome is not a single scalar but a rich observation whose individual-specific details must be preserved in a counterfactual.

Deep learning does not automatically make a model causal. Unsupervised representations are generally non-identifiable (Locatello et al. 2019), and a predictor that is accurate under observation may still fail under intervention. The goal of causal deep learning is therefore more precise: combine expressive learned mechanisms with explicit graph structure, identifiable or supervised anchors where possible, and an inference procedure that recovers the factual latent state before changing a cause. These requirements motivated the early neural causal-generative models discussed next and eventually led to tractable deep structural causal models.

Algorithm 1: Deep Structural Causal Model

Input:
    DAG G = (V, E)
    Observational samples D = {x^(n)}
    Mechanism type M_i for every variable X_i

Training:
    for every variable X_i, define X_i = f_i(epsilon_i ; pa_i)
    choose the mechanism-specific contribution ell_i
        (exact likelihood, ELBO, or adversarial objective)
    jointly maximize the sum of all ell_i terms

Counterfactual query: do(A := a_tilde) for factual x
    # Abduction
    for each variable X_i in topological order do
        if f_i is invertible then
            epsilon_i <- f_i^{-1}(x_i ; pa_i)
        else if f_i is amortised explicit then
            sample or infer z_i ~ Q_eta_i(. | x_i, pa_i)
            u_i <- h_i^{-1}(x_i ; g_i(z_i, pa_i), pa_i)
        else
            epsilon_i <- e_i(x_i, pa_i)       # approximate abduction
        end if
    end for
        # Action and prediction
    replace the assignment for A by A := a_tilde
    for each descendant X_i in topological order do
        x_i_tilde <- f_i(epsilon_i ; pa_i_tilde)
    end for

Output: counterfactual sample x_tilde

C. Causal-Gen: High-fidelity Image Counterfactuals

DSCM makes neural structural mechanisms counterfactual-capable: it provides a graph of learned assignments, tractable or amortised abduction, and reuse of factual exogenous noise after an intervention. That solves the problem of counterfactual validity, but it does not guarantee counterfactual effectiveness for high-fidelity images.

There are three practical failure modes. First, the image likelihood can be optimized mainly through latent variables or pixel residuals, so changing a causal parent produces a plausible reconstruction but barely changes the requested attribute. Second, a high-dimensional decoder can bypass the intended causal pathway: the graph says that a parent should change, yet the generated image hides that change in an unconstrained latent or residual channel. Third, forcing the parent change strongly enough to become visible can damage identity, texture, or other factual details. The generic abduction–action–prediction recipe tells us how to construct a valid counterfactual, but it supplies no separate observable loss that measures whether the requested visual effect is actually present.

Causal-Gen (Ribeiro et al. 2023) is needed precisely to close this gap. It retains the DSCM semantics while adding a practical training organization through modular staged optimization: a low-dimensional causal parent PGM makes interventions explicit, supervised image-to-parent predictors measure the attributes expressed by generated images, and a fourth, counterfactual fine-tuning process improves those measured effects without abandoning factual fidelity.

Below are the detailed training stages; each can run independently / in parallel, except the finetuning.

D. Causal-Gen Implementation on MorphoMNIST

Now let us dive deeper into how Causal-Gen is applied to the MorphoMNIST image dataset. We will move from the dataset’s measured attributes to the four Causal-Gen components, and then ask how an intervention can change one attribute while preserving the rest of the observed digit.

MorphoMNIST is an MNIST-derived grayscale-digit benchmark augmented with morphometric measurements (Castro et al. 2019). Each example combines a raster digit, its class label, and measurements that describe how that digit is drawn. This pairing matters: the class label identifies the digit, whereas thickness and intensity provide continuous attributes on which Causal-Gen can intervene and condition the image mechanism.

The raw observation is a single-channel digit image. The benchmark is deliberately small enough to inspect directly, yet it separates a full pixel realization from interpretable, low-dimensional attributes. That separation is what lets Causal-Gen distinguish “draw another plausible digit” from “change this particular digit under an intervention.”

flowchart LR
    UY(("digit noise U_Y")) --> Y(("digit Y"))
    UT(("thickness noise U_T")) --> T(("thickness T"))
    UI(("intensity noise U_I")) --> I(("intensity I"))
    T --> I
    Y --> X(("image X"))
    T --> X
    I --> X

D.2. Code Implementation with JAX

The code implementation of Causal Gen is provided by the original authors, written in PyTorch and Pyro: https://github.com/biomedia-mira/causal-gen. It already provides the complete recipe: the parent PGM, supervised image-to-attribute predictors, conditional image model, counterfactual fine-tuning, and inference.

We port the code into a pure JAX version to make the same workflow practical across CPUs, GPUs, and TPUs through one compiled array-programming stack: https://github.com/ghif/causal-genx. JAX/XLA can compile the numerically intensive training steps, while explicit parameter trees, optimizer state, random keys, moving averages, and checkpoint state make accelerator exectuion and multi-device replication easier to reason about. The rewrite also removes the Pyro runtime dependency from deployment and creates a native Flax/Optax/Orbax implementation that can be tested against the original numerical behaviour.

Stage	Artifact	Purpose on MorphoMNIST
1	Parent SCM	Fits the thickness and intensity mechanisms used to make intervened parent vectors.
2	Image-to-parent predictor	Reads thickness, intensity, and digit from an image; used as an effectiveness critic.
3	Conditional HVAE	Generates  from the 12-value parent vector while retaining image-level variation in latents.
4	Counterfactual fine-tuning	Updates only the HVAE so edited images satisfy the frozen predictor without losing reconstruction quality.

Stage 1 - Fit the parent SCM

For observed , the PGM minimizes the negative log likelihood of the implemented mechanisms:

def scm_loss(graphdef, params, batch):
    model = materialize_nnx(graphdef, params)
    logp = model.log_prob(
        batch["thickness"], batch["intensity"], batch["digit"]
    )
    loss = -jnp.mean(logp["joint"])
    return loss, {f"logp({name})": jnp.mean(logp[name])
                  for name in ("digit", "thickness", "intensity")}

Architecturally, the parent model is a hybrid rather than a pure neural network: digit uses learned categorical logits, thickness uses a monotonic rational-quadratic spline flow, and intensity uses an affine conditional flow whose location and scale are predicted by a small two-hidden-layer dense network (an MLP). The Flax NNX module keeps these pieces in one parameter graph. A compact architecture sketch is:

class ParentSCM(nnx.Module):
    def __init__(self, rngs):
        self.digit_logits = nnx.Param(jnp.zeros((1, 10)))
        self.unnormalized_widths = nnx.Param(jax.random.normal(rngs.params(), (1, 4)))
        self.unnormalized_heights = nnx.Param(jax.random.normal(rngs.params(), (1, 4)))
        self.unnormalized_derivatives = nnx.Param(jax.random.normal(rngs.params(), (1, 3)))
        self.context0 = nnx.Linear(1, 32, rngs=rngs)
        self.context1 = nnx.Linear(32, 32, rngs=rngs)
        self.context_out = nnx.Linear(32, 2, rngs=rngs)

    def intensity_params(self, thickness):
        h = jax.nn.gelu(self.context0(thickness[..., None]))
        h = jax.nn.gelu(self.context1(h))
        return jnp.split(self.context_out(h), 2, axis=-1)

At counterfactual time the SCM first recovers the factual noise for thickness and intensity. Reusing intensity’s noise after produces the descendant intensity associated with the requested thickness, rather than copying the factual intensity.

Stage 2 — Train the predictor (image critic)

The predictor is not a second causal model. It learns to estimate the MorphoMNIST variables from pixels: , , and . The conditioning in the thickness head follows the dataset’s relationship. Its supervised negative log likelihood makes it a differentiable check that a generated image visibly has the requested thickness, intensity, and digit.

The predictor is a convolutional neural network: each CNNEncoder uses six strided convolutional blocks, BatchNorm, global spatial averaging, and dense output layers. The thickness head receives intensity as a one-dimensional context; the intensity and digit heads use the image alone. It minimizes the supervised negative log likelihood of the three attributes. There is no KL term here; the predicted scales are made positive and the CNN’s BatchNorm statistics are carried as mutable state:

def predictor_loss(model, batch):
    logp = model.anticausal_log_probs(
        x=batch["x"], thickness=batch["thickness"],
        intensity=batch["intensity"], digit=batch["digit"],
    )
    return -jnp.mean(logp["joint"]), logp

The three heads share the convolutional encoder pattern but expose the graph-aware inputs required by the predictors:

class MorphoPredictor(nnx.Module):
    def __init__(self, rngs):
        self.thickness = CNNEncoder(num_outputs=2, context_dim=1, rngs=rngs)
        self.intensity = CNNEncoder(num_outputs=2, context_dim=0, rngs=rngs)
        self.digit = CNNEncoder(num_outputs=10, context_dim=0, rngs=rngs)

    def __call__(self, x, intensity):
        t_loc, t_raw = jnp.split(self.thickness(x, intensity), 2, -1)
        i_loc, i_raw = jnp.split(self.intensity(x), 2, -1)
        return {"thickness": (jnp.tanh(t_loc), t_raw),
                "intensity": (jnp.tanh(i_loc), i_raw),
                "digit_logits": self.digit(x)}

Stage 3 — Train the conditional image HVAE

The HVAE is a hierarchical convolutional variational autoencoder. Its encoder and decoder are built from convolutional residual Blocks at multiple resolutions; decoder blocks parameterize the Gaussian latent priors and posteriors, while DGaussNet is a convolutional output head for the discretized-Gaussian pixel likelihood. It models , conditioning on measured morphology and digit while latent variables retain stroke shape, placement, and other detail not represented by the three parents. Its objective is the usual reconstruction-plus-rate objective. The Flax module is composed from these three network families:

class ConditionalHVAE(nnx.Module):
    def __init__(self, rngs, context_dim=12):
        self.encoder = Encoder(input_channels=1, input_res=32, rngs=rngs)
        self.decoder = Decoder(context_dim=context_dim, rngs=rngs)
        self.likelihood = DGaussNet(input_channels=1, rngs=rngs)

    def loss(self, x, parents, beta=1.0, rng=None):
        acts = self.encoder(x)
        features, stats = self.decoder(
            parents=parents, x=acts, rng=rng, training=True
        )
        nll = self.likelihood.nll(features, x).mean()
        kl = sum(jnp.sum(s["kl"], axis=(1, 2, 3)).mean()
                 for s in stats)
        kl = kl / np.prod(x.shape[1:])
        return nll + beta * kl, {"nll": nll, "kl": kl}

The HVAE regularizer is the hierarchical KL term, weighted by . With kl_free_bits > 0, each latent block applies a floor to its KL contribution before summation. The supplied MorphoMNIST configuration uses cond_prior: true, so the prior networks also receive the parent state.

Stage 4 — Fine-tune for an effective edit

Stage 4 introduces no new architecture: it reuses the spline-plus-MLP PGM, convolutional predictor, and hierarchical convolutional HVAE. The SCM and predictor are frozen. For a factual batch, the code creates an intervened parent vector with the SCM, abducts the HVAE latents, renders with the intervened parents, and asks the predictor to recognize those new MorphoMNIST attributes. Only HVAE parameters and a non-negative Lagrange multiplier are optimized:

def counterfactual_loss(vae, pgm, predictor, batch, intervention,
                        lmbda, beta, elbo_constraint, damping, rng):
    vae_out, cf = counterfactual_forward(
        vae, pgm, batch, intervention, beta=beta, rng=rng
    )
    aux_loss = -jnp.mean(predictor.model_anticausal(**cf)["joint"])
    constraint = elbo_constraint - vae_out["elbo"]
    loss = aux_loss - (
        lmbda - damping * jax.lax.stop_gradient(constraint)
    ) * constraint
    return loss, {"aux_loss": aux_loss, "constraint": constraint}

The auxiliary loss asks the frozen predictor to recognize the intervened attributes. The damped Lagrangian keeps the HVAE negative ELBO near its factual constraint; only the HVAE parameters and non-negative multiplier are optimized, while the PGM and predictor remain frozen. AdamW weight decay is applied in the optimizer transformation rather than added to any of these loss expressions.

In JAX, freezing the model weights is expressed by the arguments that a loss closes over and the arguments passed to value_and_grad. Stage 4 captures the PGM and predictor bundles as constants, so their parameters receive no gradient and no optimizer state:

def frozen_cf_loss(vae_params, lmbda, batch, intervention, rng):
    vae = Bundle(vae_graphdef, vae_params)       # trainable argument
    # pgm_bundle and predictor_bundle are captured, frozen values.
    out = _cf_forward(
        args, vae, pgm_bundle, predictor_bundle, batch, intervention, rng,
        beta=args.beta, alpha=args.alpha, lmbda=lmbda,
        cf_particles=args.cf_particles, training=True,
    )
    return out["loss"], out

(loss, out), (vae_grads, lmbda_grads) = jax.value_and_grad(
    frozen_cf_loss, argnums=(0, 1), has_aux=True
)(vae_params, lmbda, batch, intervention, rng)

Example of Counterfactual Query

For a thickness edit, the operational sequence is simple:

1. Read factual and infer the SCM noise plus HVAE latents.
2. Apply .
3. Recompute through the implemented thickness-to-intensity mechanism with the factual intensity noise held fixed.
4. Decode the factual latents under and reuse the factual pixel residual.

The result changes the morphology selected by the intervention while aiming to keep digit identity and incidental visual detail stable. A direct intervention on intensity, , instead bypasses that descendant calculation. A digit edit changes the digit parent supplied to the image model; as noted above, the current PGM does not automatically derive a new thickness from digit.

cd causal-genx

# Parent SCM/PGM
python scripts/run.py train-scm --config configs/morphomnist_scm.yaml dataset.root=/path/to/morphomnist

# Predictor
python scripts/run.py train-predictor --config configs/morphomnist_predictor.yaml dataset.root=/path/to/morphomnist

# HVAE image generator
python scripts/run.py train-image-model --config configs/morphomnist_image_model.yaml dataset.root=/path/to/morphomnist

# Counterfactual finetuning
python scripts/run.py finetune-counterfactual --config configs/morphomnist_counterfactual.yaml dataset.root=/path/to/morphomnist

use_tpu_pmap = (
    args.accelerator == "tpu"
    and jax.local_device_count() > 1
    and args.execution_mode != "single_device"
)
devices = jax.local_devices() if use_tpu_pmap else []
device_count = len(devices) if use_tpu_pmap else 1
if use_tpu_pmap and args.bs % device_count:
    raise ValueError("global batch must be divisible by local device count")

mesh = jax.sharding.Mesh(np.asarray(devices), ("devices",))
sharding = jax.sharding.NamedSharding(
    mesh, jax.sharding.PartitionSpec("devices")
)

def shard_batch(x):
    per_device = x.shape[0] // len(devices)
    return jax.device_put(
        x.reshape((len(devices), per_device) + x.shape[1:]), sharding
    )

def pmap_step(params, opt_state, batch, rng):
    (loss, aux), grads = jax.value_and_grad(loss_fn, has_aux=True)(
        params, batch, rng
    )
    grads = jax.lax.pmean(grads, axis_name="devices")
    aux = jax.tree_util.tree_map(
        lambda value: jax.lax.pmean(value, axis_name="devices"), aux
    )
    updates, opt_state = tx.update(grads, opt_state, params)
    return optax.apply_updates(params, updates), opt_state, aux

train_step = jax.pmap(
    pmap_step, axis_name="devices", in_axes=(0, 0, 0, 0), devices=devices
)

flowchart LR
    User["web browser"] --> Pages["GitHub Pages<br/>static frontend"]
    Pages -->|HTTPS + CORS| API["Cloud Run<br/>FastAPI service"]
    API --> GCS["private GCS<br/>model checkpoints"]
    API --> Result["PNG +<br/>factual/CF parents"]
    Result --> Pages

Conclusion and Acknowledgement

Our attempt on implementing a counterfactual image generation with a pure JAX approach (causal-genx) keeps the original causal-gen scientific recipe intact: an invertible parent SCM, supervised image-to-parent predictors, a hierarchical image model, and constrained counterfactual fine-tuning, but turns that recipe into an accelerator-ready, inspectable system. Relative to the original PyTorch/Pyro implementation, this work contributes a native JAX/Flax/Optax/Orbax codebase with explicit parameter, optimizer, EMA, random-key, and checkpoint state; a single workflow that runs on CPU, GPU, single-core TPU, and replicated multi-core TPU; and a documented four-stage MorphoMNIST implementation whose losses, architectures, intervention semantics, and freeze/unfreeze boundaries can be read directly from the code.

The result is more than a framework translation. Compilation and batch-sharding choices in JAX make the workload practical at scale. The same artifacts can be served through the interactive GitHub Pages and Google Cloud Run demo, allowing a user to inspect factual parents, apply an intervention, and compare the resulting image with its source. These additions improve portability, throughput, and usability.

We thank the authors of the original Causal-Gen work for releasing the PyTorch/Pyro reference implementation that made this port and its numerical checks possible. We also gratefully acknowledge Google for providing the cloud credits used for TPU experiments and the public demonstration. The Causal-GenX repository extends that foundation as an open, reproducible engineering companion: the original work supplies the causal blueprint, while this implementation supplies a modern accelerator and deployment path for exploring it.

About Tensor Processing Units (TPUs)

Before diving into the optimization strategy, let’s briefly discuss the hardware accelerator itself. A basic understanding of TPU architecture will help us plan the optimization more effectively.

TPUs are Google’s custom-developed application-specific integrated circuits (ASICs) designed specifically to accelerate machine learning workloads. Unlike general-purpose GPUs, TPUs are architected around the requirements of deep learning, prioritizing high-throughput matrix multiplications and low-latency interconnects.

gcloud alpha compute tpus queued-resources create tpu-southamerica-queue \
    --zone=southamerica-east1-c \
    --accelerator-type=v6e-4 \
    --runtime-version=v2-alpha-tpuv6e \
    --node-id=my-tpu-node \
    --provisioning-model=flex-start \
    --max-run-duration=96h \
    --valid-until-duration=96h \
    --labels=purpose=flex-start

Optimization Strategy

The initial implementation used a traditional deep learning training loop structure:

for i in pbar:
    idx = random.randin(0, len(jax_cameras)-1)
    cam = jax_cameras[idx]
    target = jax_targets[idx]
    
    # Single step dispatched to device
    state, loss, metrics = train_step(state, target, cam.W2C, camera_static, optimizer)
    
    if i % 100 == 0:
        img = render(state[0], jax_cameras[0]) # Synchronous render
        save_ply(...) # Symchronous I/O
    

While functional, this approach has a few drawbacks or fails to capitalize on JAX’s core strengths.

1. Dispatch Overhead: Python has to tell the TPU what to do at every step.
2. Host-Device Communication: Sampling images on the CPU and sending them to the TPU creates a bottleneck.
3. Synchronous I/O: The entire training process pauses to wait for disk writes.

To achieve faster performance, we must transition to a JAX-native architecture designed for TPU efficiency. The primary goal of jax-gs is to minimize host-accelerator communication and maximize the utilization of TPU systolic arrays (MXUs). We achieve this through several core optimization approaches that restructure the training and rendering pipeline into high-throughput, block-compile execution units.

# From train.py
@partial(jax.jit, static_argnums=(4, 5, 6, 7, 8))
def train_block(state, rng_key, all_targets, all_w2cs, steps_per_block, camera_static, optimizer, fast_tpu_rasterizer, sh_degree):
    def one_step(carry, _):
        state, key = carry
        key, subkey = jax.random.split(key)
        
        # Sample camera index
        idx = jax.random.randint(subkey, (), 0, all_targets.shape[0])
        target = all_targets[idx]
        w2c = all_w2cs[idx]
        
        # Perform training step
        state, loss, metrics = train_step(state, target, w2c, camera_static, optimizer, fast_tpu_rasterizer=fast_tpu_rasterizer, sh_degree=sh_degree)
        
        return (state, key), loss

    # The entire loop runs on the TPU as a single XLA program
    (state, rng_key), losses = jax.lax.scan(one_step, (state, rng_key), None, length=steps_per_block)
    return state, rng_key, losses

# From train.py
# 1. Prepare data on device (all images and matrices loaded once)
all_targets = jnp.stack(jax_targets)
all_w2cs = jnp.stack([c.W2C for c in jax_cameras])

# 2. Inside the JIT-compiled train_block:
# Sampling happens entirely on-device
idx = jax.random.randint(subkey, (), 0, all_targets.shape[0])
target = all_targets[idx]
w2c = all_w2cs[idx]

# Background executor for I/O and rendering
executor = concurrent.futures.ThreadPoolExecutor(max_workers=1)

# Inside training loop:
if curr_iter % 1000 == 0:
    snap_gaussians_dict = get_active_gaussians(curr_state)
    fut = executor.submit(
        save_artifacts_task, 
        snap_gaussians_dict, curr_iter, progress_dir, ply_path, 
        jax_cameras[0], fast_tpu_rasterizer, render, save_ply, sh_degree
    )

# From jax_gs/renderer/rasterizer.py
def rasterize_single_tile(tile_idx):
    # ... logic for one tile ...
    def blend_pixel(p_coord, p_valid):
        # ... blending logic for one pixel ...
        return final_color

    # Parallelize over pixels in the tile
    tile_image = jax.vmap(blend_pixel)(pixel_coords, pixel_valid)
    return tile_image.reshape(tile_size, tile_size, 3)

# Parallelize over all tiles in the image
all_tiles = jax.vmap(rasterize_single_tile)(jnp.arange(num_tiles))

# From jax_gs/renderer/rasterizer_tpu.py
# Pre-calculate global pixel coordinates for EVERY pixel in the image, 
# grouped by tile. Shape: [num_tiles, 256]
px = (tx[:, None] * TILE_SIZE).astype(jnp.float32) + (idx % TILE_SIZE)[None, :].astype(jnp.float32) + 0.5
py = (ty[:, None] * TILE_SIZE).astype(jnp.float32) + (idx // TILE_SIZE)[None, :].astype(jnp.float32) + 0.5

@jax.checkpoint
def scan_fn(carry, i):
    # Vectorized across [num_tiles, 256] flat dimension
    dx = px - mu_x 
    dy = py - mu_y
    # ... blending logic ...

# From jax_gs/renderer/rasterizer_tpu.py
# BROADCASTED GATHER: Construct indices for all Gaussians across all tiles.
# Resulting shape: [num_tiles, BLOCK_SIZE]
all_tile_indices = tile_starts[:, None] + jnp.arange(BLOCK_SIZE)[None, :]

# Prefetch Gaussian data for all tiles at once
tile_gids = valid_ids[all_tile_indices] 
g_means = means2D[tile_gids]            # [num_tiles, BLOCK_SIZE, 2]
g_cols = colors[tile_gids]              # [num_tiles, BLOCK_SIZE, 3]

@jax.checkpoint
def scan_fn(carry, i):
    # Processes the i-th Gaussian for ALL pixels in ALL tiles simultaneously
    # No random access lookups here!
    mu_x = g_means[:, i, 0][:, None] # [num_tiles, 1]
    dx = px - mu_x                   # [num_tiles, 256]
    # ... blending logic ...

# From train_parallel.py
@partial(jax.pmap, axis_name="batch", ...)
def train_block(state, ...):
    # ... computation ...
    def one_step(carry, inputs):
        # ... train_step_internal performs pmean on gradients ...
        return state, (loss, metrics)

    state, (losses, metrics) = jax.lax.scan(one_step, state, (batch_targets, batch_w2cs))
    avg_metrics = jax.tree_util.tree_map(lambda x: jnp.mean(x, axis=0), metrics)
    return state, losses, avg_metrics

# From jax_gs/training/trainer.py
# Gradients and metrics are averaged across devices
grads = jax.tree_util.tree_map(lambda x: jax.lax.pmean(x, axis_name='batch'), grads)
loss = jax.lax.pmean(loss, axis_name='batch')
metrics = jax.tree_util.tree_map(lambda x: jax.lax.pmean(x, axis_name='batch'), metrics)

Benchmark Results

Our benchmarks on the LLFF room dataset (504x378 resolution) demonstrate the massive performance gains achieved by our JAX-native optimizations. On a Google Cloud TPU v6e-4 (Trillium), we observed the following:

Conclusion

The transformation of 3D Gaussian Splatting into a static-shape, JIT-compatible architecture within JAX enables efficient training and rendering on modern accelerators. By prioritizing systolic array saturation, minimizing host-device communication, and leveraging the low-latency TPU interconnects, jax-gs provides a robust and scalable foundation for the next generation of radiance field research.

Acknowledgments

Google Cloud credits are provided for this project. #TPUSprint

1. The Rise of On-Device Medical Intelligence

Traditional medical AI often relies on sending high-resolution scans and patient records to powerful GPU clusters in the cloud. While effective, this approach introduces significant bottlenecks:

• Privacy Risks: Every byte of data leaving the hospital network is a potential point of failure.
• Latency: In critical care, waiting for a round-trip to a data center can be the difference between a prompt diagnosis and a delayed one.
• Connectivity: Many clinical environments (ORs, remote clinics, or mobile health units) suffer from inconsistent internet access.

On-device intelligence solves these by performing inference where the data is born. With the release of Google’s Gemma 3 architecture and its medical sibling MedGemma 1.5, the “edge” is now powerful enough to handle 4-billion parameter multimodal models.

2. MedGemma 1.5 4B: A Multimodal Leap

MedGemma 1.5 4B ****represents a significant architectural shift over its predecessors. While MedGemma 1.0 was a pioneer in clinical text understanding, the 1.5 iteration—built on the Gemma 3 foundation—is a true multimodal powerhouse.

Some key advancements of MedGemma 1.5 are as follows:

• From 2D to 3D: Previous models focused on 2D images like X-rays or dermoscopy. MedGemma 1.5 natively interprets 3D medical volumes from CT and MRI scans.
• Longitudinal Reasoning: One of the model’s strongest features is its ability to track disease progression by comparing historical scans against current ones.
• Clinical Accuracy: EHR Question Answering accuracy has jumped from ~68% in version 1 to a staggering ~90% in version 1.5.
• Anatomy Localization: Precise identification of anatomical structures and abnormalities saw an improvement from ~3% IoU to ~38% IoU.

3. Deep Dive into the .litertlm Format

To run MedGemma efficiently on the edge, we leverage the .litertlm format. This isn’t just another file extension; it is LiteRT’s (formerly TensorFlow Lite) specialized bundle for Generative AI.

The .litertlm offers several benefits.

1. Stateful Optimization: Unlike standard .tflite graphs, a .litertlm bundle is designed for the iterative nature of LLMs. It contains separate, optimized graphs for prefill (processing the prompt) and decode (generating tokens one-by-one), while natively managing the KV-cache.
2. Multimodal Synergy: MedGemma requires a vision encoder and a language head to work tandem. .litertlm bundles these disparate components into a single, self-describing artifact, ensuring the vision-language projection layers are always synchronized.
3. Hardware Native: The format is built to leverage the LiteRT GenAI API, which provides highly optimized kernels for mobile GPUs and NPUs, significantly outperforming generic graph execution.

4. The Conversion Workflow

The transformation from raw PyTorch weights to a production-ready LiteRT bundle is an intricate process handled via the litert-torch export pipeline. It starts with the original MedGemma 1.5 4B IT weights from Google and re-architects them for high-performance edge execution.

Instead of a simple “file format save,” the conversion logic performs several critical architectural bridges and optimizations:

1. Structural Alignment (Architecture Bridging): MedGemma 1.5 is built on the Gemma 3 architecture. In some versions of the Hugging Face transformers library, the vision modules (tower and projector) are nested deeply within the model structure. The workflow includes a structural patch that maps these nested components to the top level of the model class. This ensures the export engine can accurately “see” and trace the multimodal connection points during the graph-generation phase.
2. Multimodal Graph Tracing: The conversion initiates a image_text_to_text export task. This process traces the mathematical flow of data through both the vision encoder and the language head. It effectively captures how an MRI image is transformed into tokens and how those tokens are processed by the LLM to generate a clinical description.
3. Prefill Bucketing: To optimize the “time-to-first-token” on edge devices, the workflow generates specialized graphs for different prefill lengths (e.g., 128, 256, 512 tokens). This allows the runtime to use the most efficient computation path based on the size of the user’s initial prompt or image metadata.
4. Specialized Dual-Quantization: To compress the 4B parameter model to a browser-friendly ~3GB, we apply distinct quantization strategies to different components:
   1. LLM Core: Uses a dynamic_wi8_afp32 recipe (8-bit weights with 32-bit activations), balancing reasoning depth with memory footprint.
   2. Vision Encoder: Uses a weight_only_wi8_afp32 recipe, ensuring that the high-dimensional features required for medical imaging are preserved while still reducing the storage overhead.
5. KV-Cache Architecture: The workflow configures a fixed-length Key-Value (KV) cache (typically 4096 tokens). This is embedded into the LiteRT graph, enabling the model to “remember” the context of a long medical conversation without re-processing the entire history for every new word generated.
6. Unified Bundling: The final step packages the optimized graphs, the token embedder, the tokenizer configuration, and essential model metadata into a single, self-describing .litertlm container. This eliminates the need for external configuration files and ensures the model is “plug-and-play” for the edge runtime.

5. Model Deployment on Edge Web Browsers

Deploying a 4-billion parameter multimodal model like MedGemma in a browser tab is a feat of modern web engineering. It transforms the browser from a simple document viewer into a secure, hardware-accelerated sandbox for private clinical intelligence. To verify the effectiveness of the converted model, we implement a simple GUI prototype written in JavaScript/TypeScript that runs on a local web browser, in which the model deployment is optimized through a multi-layered architectural approach.

medgemma_edge_demo.mp4

5. Conclusion

The integration of MedGemma 1.5 4B with LiteRT represents a significant milestone in medical AI accessibility, delivering expert-level multimodal intelligence directly to the browser. By enabling private, always-available assistants on-device, we are bridging the gap between cloud-scale performance and edge efficiency, ensuring clinicians can access critical decision support regardless of their connectivity.

To transition this from a research prototype to a global production tool for resource-constrained environments, our future roadmap focuses on three key optimizations:

• Int4 QAT Quantization: Implementing Quantization-Aware Training to drastically reduce memory requirements for entry-level mobile and laptop hardware.
• Speculative Decoding: Integrating draft models to achieve a 2–3× increase in inference speed, essential for high-volume clinical workflows.
• Localized Fine-Tuning: Adapting the model to regional medical terminology and local languages to ensure real-world clinical utility.

A primary example of the need for these advancements is found in Indonesia’s 3T regions (underdeveloped, frontier, and outermost). In these settings, healthcare workers require high-accuracy tools that operate fully offline, speak Bahasa Indonesia, and run on the affordable, low-spec hardware already available in community health centers (Puskesmas).

Resources:

• Model (.litertlm): huggingface.co/ai4med-id/medgemma-1.5-4b-it-litertlm
• Implementation & GUI Prototype: github.com/AI4MedResearch/edge-ai

Why C++ Native Kernels over Pure Python?

While high-level frameworks like MLX and PyTorch offer excellent vectorized operations, custom algorithms like 3D Gaussian Splatting hit several architectural limitations when implemented in pure Python. Moving to native C++ and Metal kernels is not just about raw speed; it is about achieving the architectural freedom required for production-quality results.

Implementation: A Fully GPU-Resident MLX Metal Pipeline

Achieving 40+ iterations per second on a single M4 chip required a fundamental shift in how the rendering pipeline is structured. Instead of treating the GPU as a math accelerator for specific loops, we implemented a Fully GPU-Resident Pipeline. This includes the projection of 3D Gaussians into 2D space, where the 2D covariance matrix  is computed from the 3D covariance , the world-to-view transformation , and the Jacobian of the affine approximation of the projective transformation :

Here is the step-by-step breakdown of how the MLX Metal renderer is implemented:

# mlx_gs/renderer/rasterizer_metal.py
...
# Vectorized expansion using cumsum and marking masks
offsets = mx.cumsum(counts)
mark = mx.zeros((total,), dtype=mx.int32)
mark[active_starts] = 1
map_idx = mx.cumsum(mark) - 1
gaussian_ids = active_indices[map_idx]

# High-speed sorting using bitwise-packed 64-bit keys
keys = (tile_ids.astype(mx.uint64) << 32) | depth_quant
sort_indices = mx.argsort(keys)

// mlx_gs/csrc/rasterizer_metal.mm

id<MTLBuffer> wrap(void* ptr, size_t size){
    // Zero-copy wrapping of Python/MLX pointers into Metal buffers
    return [device newBufferWithBytes:ptr length:size options:MTLResourceStorageModeShared];
}

// Dispatching the kernel with 16x16 tile threadgroups
[enc setComputePipelineState:forwardPSO];
[enc setBuffer:b_m offset:0 atIndex:0]; // Means
[enc dispatchThreads:MTLSizeMake(ntx*16, nty*16, 1)
    threadsPerThreadgroup:MTLSizeMake(16, 16, 1)];
...
/**
 * Nanobind Module Definition.
 */
NB_MODULE(_rasterizer_metal, m) {
    m.def("init_metal", &init_metal, "Initialize Metal device and compile shaders");
    m.def("render_forward", &render_forward, "Perform forward rasterization on Metal");
    m.def("render_backward", &render_backward, "Compute gradients for rasterization on Metal");
}

// mlx_gs/csrc/rasterizer.metal

for (int i = s_idx; i < e_idx; ++i) {
    int g_idx = sorted_gaussian_ids[i];
    // ... compute Gaussian alpha ...
    float alpha = 1.0f - exp(-exp(power) * opacities[g_idx]);
    float test_T = T * (1.0f - alpha);

    // Early exit: Stop if the pixel is already opaque
    if (test_T < 0.0001f) break;

    C += color * alpha * T;
    T = test_T;
}

// mlx_gs/csrc/rasterizer.metal

void atomic_add_float(device atomic_uint* addr, float val){
    uint old_val = atomic_load_explicit(addr, memory_order_relaxed);
    uint new_val;
    do {
        new_val = as_type<uint>(as_type<float>(old_val) + val);
    } while (!atomic_compare_exchange_weak_explicit(addr, &old_val, new_val, ...));
}

// Accumulating gradients for means and colors
atomic_add_float(&grad_means2D[g_idx * 2 + 0], dL_dmean.x);
atomic_add_float(&grad_colors[g_idx * 3 + 0], dL_dPixel.x * w);

# mlx_gs/renderer/rasterizer_metal.py

@mx.custom_function
def forward(m, ic, s_o, c, sti, sgi, bg):
    # Call the C++ extension
    out_np = rasterizer_metal.render_forward(...)
    return mx.array(out_np)

@forward.vjp
def backward(primals, cotangents, outputs):
    # Pass gradients back through the Metal backward kernel
    gm_np, gic_np, go_np, gc_np = rasterizer_metal.render_backward(...)
    return mx.array(gm_np), mx.array(gic_np), mx.array(go_np), ...

The PyTorch Challenge: Why CPU-GCD was the Choice for Stability

Besides MLX, we could also use PyTorch with Python-C++ bindings to speed-up 3D Gaussian Splatting. Therefore, I also attempted to implement a similar “Zero-Sync” Metal extension for PyTorch on MPS.

However, I encountered several architectural hurdles. These challenges highlight the fundamental differences between frameworks designed for unified memory (MLX) and those with more complex, cross-platform backends (PyTorch).

Benchmark Analysis: The Performance Landscape

To evaluate the impact of these native optimizations, I benchmarked the training loop using the standard Fern dataset (10,091 Gaussians) at a resolution of 504x378 on an Apple M4 (10-core GPU, 16GB Unified Memory).

Here are some training visualization over a few datasets produced by the Native Metal MLX:

fern_training.mp4

pinecone_training.mp4

room_training.mp4

trex_training.mp4

Conclusion

The transition from high-level Python references to native Metal extensions has fundamentally redefined what is possible for 3D computer vision on Apple Silicon. For years, production-grade 3D Gaussian Splatting was considered synonymous with NVIDIA hardware and CUDA kernels. This simple project demonstrates the possibility of leveraging Apple Silicon as part of the “first-class citizens” in the production-grade 3D AI landscape.

By combining the MLX framework with a fully GPU-resident Metal pipeline, we can deliver a solution that bridges the gap between prototyping and production. With a steady-state performance of over 38 iterations per second, the Mac is no longer a platform restricted to slow experimentation or viewing-only roles.

The Core Idea: Coarse-to-Fine

DECA operates on a coarse-to-fine principle. It doesn’t try to solve everything at once. Instead, it breaks the problem down into two stages:

1. Coarse Reconstruction: Estimate the underlying head shape, head pose, and facial expression using a statistical 3D face model (FLAME).
2. Detail Reconstruction: Predict a person-specific detail map (displacement map) that adds high-frequency details (like forehead wrinkles or crow’s feet) to the coarse mesh.

The FLAME Model

At the heart of DECA is the FLAME (Faces Learned with Articulated Model and Expressions) model (Li et al. 2017). FLAME is a parametric model, meaning it generates a 3D mesh based on a set of low-dimensional parameters.

Mathematically, a 3D mesh with vertices is generated as:

where:

• (Shape parameters): Controls identity-specific shape (e.g., big nose, wide jaw).
• (Expression parameters): Controls facial expressions (e.g., smile, frown).
• (Pose parameters): Controls the rotation of the neck and jaw.

I discussed about FLAME in more details in my previous article.

In DECA, the goal of the neural network is to predict these parameters (, , ) along with camera parameters and lighting from a single input image.

Coarse Reconstruction Stream

The coarse stream is a straightforward encoder-decoder architecture.

class ResnetEncoder(nn.Module):
    def __init__(self, outsize, last_op=None):
      super(ResnetEncoder, self).__init__()
    feature_size = 2048
        self.encoder = resnet.load_ResNet50Model() # Backbone
        ### regressor
        self.layers = nn.Sequential(
        nn.Linear(feature_size, 1024),
            nn.ReLU(),
            nn.Linear(1024, outsize) # Projects to n_shape + n_exp + n_pose + ...
    )
        self.last_op = last_op

    def forward(self, inputs):
            features = self.encoder(inputs)
            parameters = self.layers(features)
            if self.last_op:
                parameters = self.last_op(parameters)
            return parameters

Detail Reconstruction Stream

This is where DECA shines. A standard parametric model is too smooth. To capture wrinkles, DECA also predicts a UV displacement map.

The detail reconstruction uses another encoder to extract a latent detail code . This code, combined with the expression and pose parameters, is fed into a generator.

Why include expression and pose? Because wrinkles are dynamic. A forehead wrinkle appears when we raise our eyebrows. By conditioning the generator on expression, DECA learns expression-dependent details.

Here is the implementation of the detail reconstruction through Generator class below.

class Generator(nn.Module):
    def __init__(self, latent_dim=100, out_channels=1, out_scale=0.01, sample_mode = 'bilinear'):
        super(Generator, self).__init__()
        self.out_scale = out_scale
        
        self.init_size = 32 // 4 
        self.l1 = nn.Sequential(nn.Linear(latent_dim, 128 * self.init_size ** 2))
        self.conv_blocks = nn.Sequential(
            # Upsampling layers to generate a high-res displacement map
            nn.BatchNorm2d(128),
            nn.Upsample(scale_factor=2, mode=sample_mode), 
            nn.Conv2d(128, 128, 3, stride=1, padding=1),
            # ... (more layers) ...
            nn.Conv2d(16, out_channels, 3, stride=1, padding=1),
            nn.Tanh(),
        )

    def forward(self, noise):
        out = self.l1(noise)
        out = out.view(out.shape[0], 128, self.init_size, self.init_size)
        img = self.conv_blocks(out)
        return img * self.out_scale

The output is a displacement map . The final detailed mesh vertices are calculated by displacing the coarse vertices along their normals :

Training: The Loss Functions

DECA uses a complex combination of loss functions to train these networks without requiring 3D ground truth (unsupervised learning).

Landmark Loss (): ****Ensures the projected 3D keypoints match the 2D facial landmarks detected in the image.

def landmark_loss(predicted_landmarks, landmarks_gt, weight=1.):
    if torch.is_tensor(landmarks_gt) is not True:
        real_2d = torch.cat(landmarks_gt).cuda()
    else:
        real_2d = torch.cat([landmarks_gt, torch.ones((landmarks_gt.shape[0], 68, 1)).cuda()], dim=-1)

    loss_lmk_2d = batch_kp_2d_l1_loss(real_2d, predicted_landmarks)
    return loss_lmk_2d * weight

Photometric Loss (): ****Computes the pixel-wise difference between the input image and the rendered 3D face. This forces the texture and shape to match the input.

Identity Loss (): ****Uses a pretrained Face Recognition network (like ResNet50 trained on VGGFace2) to extract feature vectors. The loss minimizes the cosine distance between the features of the input image and the rendered image. This ensures the reconstructed 3D face “looks like” the person.

class VGGFace2Loss(nn.Module):
    # ...
    def forward(self, gen, tar, is_crop=True):
        gen = self.transform(gen)
        tar = self.transform(tar)

        gen_out = self.reg_features(gen)
        tar_out = self.reg_features(tar)
        # Cosine similarity loss
        loss = self._cos_metric(gen_out, tar_out).mean()
        return loss

Detail Loss (): ****Since pixel-wise difference doesn’t work well for high-frequency details (due to slight misalignments), DECA uses an ID-MRF loss. This matches local patches of features from the input image to the rendered detail image, effectively transferring the “style” (details) of the photo to the 3D model.

Putting It All Together: The Inference Process

A critical distinction of DECA’s architecture is that it is a learning-based approach, as opposed to traditional optimization-based methods. While optimization-based methods solve an energy minimization problem for every image (seconds or minutes), DECA’s pipeline predicts all parameters in a single forward pass (milliseconds). This enables the real-time interaction and animation.

When DECA is run on a face image, the following sequence of operations occurs:

Custom C++ CPU Rasterizer

In standard DECA (and many modern 3D project), the rasterization step is typically handled by CUDA-based engines or heavy libraries like PyTorch3D. To achieve native, fast performance on a local machine (e.g., Apple Silicon) without these dependencies, I implemented a custom C++ CPU rasterizer that also leverages PyTorch C++ API.

//standard_rasterize_cpu.cpp
#include <torch/extension.h>
#include <vector>
#include <cmath>
#include <algorithm>
#include <iostream>

// Helper struct for 2D points
template <typename T>
struct Point {
    T x, y;
};

// Barycentric weight calculation
template <typename T>
void barycentric_weight(T* w, Point<T> p, Point<T> p0, Point<T> p1, Point<T> p2) {
    Point<T> v0 = {p2.x - p0.x, p2.y - p0.y};
    Point<T> v1 = {p1.x - p0.x, p1.y - p0.y};
    Point<T> v2 = {p.x - p0.x, p.y - p0.y};

    T dot00 = v0.x * v0.x + v0.y * v0.y;
    T dot01 = v0.x * v1.x + v0.y * v1.y;
    T dot02 = v0.x * v2.x + v0.y * v2.y;
    T dot11 = v1.x * v1.x + v1.y * v1.y;
    T dot12 = v1.x * v2.x + v1.y * v2.y;

    T denom = dot00 * dot11 - dot01 * dot01;
    T invDenom = (denom == 0) ? 0 : 1 / denom;

    T u = (dot11 * dot02 - dot01 * dot12) * invDenom;
    T v = (dot00 * dot12 - dot01 * dot02) * invDenom;

    w[0] = 1 - u - v;
    w[1] = v;
    w[2] = u;
}

void standard_rasterize_cpu_kernel(
    const torch::Tensor& face_vertices,
    torch::Tensor& depth_buffer,
    torch::Tensor& triangle_buffer,
    torch::Tensor& baryw_buffer,
    int h, int w
) {
    // face_vertices: [B, F, 3, 3]
    // depth_buffer: [B, H, W]
    // triangle_buffer: [B, H, W]
    // baryw_buffer: [B, H, W, 3]

    auto B = face_vertices.size(0);
    auto F = face_vertices.size(1);

    // Get accessors for efficient element access
    // Assuming float input as per standard renderer usage
    auto face_vertices_a = face_vertices.accessor<float, 4>();
    auto depth_buffer_a = depth_buffer.accessor<float, 3>();
    auto triangle_buffer_a = triangle_buffer.accessor<int, 3>();
    auto baryw_buffer_a = baryw_buffer.accessor<float, 4>();

    for (int b = 0; b < B; ++b) {
        for (int f = 0; f < F; ++f) {
            float p0x = face_vertices_a[b][f][0][0];
            float p0y = face_vertices_a[b][f][0][1];
            float p0z = face_vertices_a[b][f][0][2];

            float p1x = face_vertices_a[b][f][1][0];
            float p1y = face_vertices_a[b][f][1][1];
            float p1z = face_vertices_a[b][f][1][2];

            float p2x = face_vertices_a[b][f][2][0];
            float p2y = face_vertices_a[b][f][2][1];
            float p2z = face_vertices_a[b][f][2][2];

            // Calculate bounding box for the triangle
            int x_min = std::max(0, (int)std::ceil(std::min({p0x, p1x, p2x})));
            int x_max = std::min(w - 1, (int)std::floor(std::max({p0x, p1x, p2x})));
            int y_min = std::max(0, (int)std::ceil(std::min({p0y, p1y, p2y})));
            int y_max = std::min(h - 1, (int)std::floor(std::max({p0y, p1y, p2y})));

            if (x_min > x_max || y_min > y_max) continue;

            Point<float> p0 = {p0x, p0y};
            Point<float> p1 = {p1x, p1y};
            Point<float> p2 = {p2x, p2y};

            for (int y = y_min; y <= y_max; ++y) {
                for (int x = x_min; x <= x_max; ++x) {
                    Point<float> p = {(float)x, (float)y};
                    float bw[3];
                    barycentric_weight(bw, p, p0, p1, p2);

                    // Check if pixel is inside the triangle
                    if (bw[0] > 0 && bw[1] >= 0 && bw[2] >= 0) {
                        // Perspective correct depth interpolation
                        float zp = 1.0f / (bw[0] / p0z + bw[1] / p1z + bw[2] / p2z);
                        
                        // Z-buffer test
                        if (zp < depth_buffer_a[b][y][x]) {
                            depth_buffer_a[b][y][x] = zp;
                            triangle_buffer_a[b][y][x] = f;
                            baryw_buffer_a[b][y][x][0] = bw[0];
                            baryw_buffer_a[b][y][x][1] = bw[1];
                            baryw_buffer_a[b][y][x][2] = bw[2];
                        }
                    }
                }
            }
        }
    }
}

void standard_rasterize(
    torch::Tensor face_vertices,
    torch::Tensor depth_buffer,
    torch::Tensor triangle_buffer,
    torch::Tensor baryw_buffer,
    int h, int w
) {
    standard_rasterize_cpu_kernel(face_vertices, depth_buffer, triangle_buffer, baryw_buffer, h, w);
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("standard_rasterize", &standard_rasterize, "Standard Rasterize (CPU)");
}

from torch.utils.cpp_extension import load
...

try:
    standard_rasterize_cpu_module = load(
    name='standard_rasterize_cpu', 
    sources=[source_path],
    verbose=True
    )
  _standard_rasterize_cpu_impl = standard_rasterize_cpu_module.standard_rasterize
  print(f"[set_rasterizer] Successfully loaded standard rasterizer extension.")
except Exception as e:
  print(f"[set_rasterizer] Error loading C++ extension: {e}")
  raise
...

Conclusion

DECA represents a considerable milestone in 3D facial mesh reconstruction. By combining robust coarse shape estimation with detailed, expression-driven displacement mapping, it bridges the gap between static mesh reconstruction and realistic animation.

With the native Apple Silicon port leveraging PyTorch capabilities to run on the CPU/MPS devices, we can run DECA without requiring CUDA GPU.

For the full source code and implementation details, check out my forked DECA repository.

What is FLAME?

FLAME is a Linear Blend Skinning (LBS) model that captures the vast variety of human head shapes and expressions. Unlike older models that focused solely on the face, FLAME models the entire head — including the neck, jaw, and eyballs.

At its heart, FLAME represents a 3D mesh , which is a function of:

• Shape parameters : Identity-specific features (height, face width, etc.).
• Pose parameters : Rotations for the neck, jaw, and eyeballs.
• Expression parameters : Dynamic movements like smiles or frowns.

The final position of a vertex is calculated using the LBS formula:

where

• is the average head shape.
• , , are the Shape, Pose, and Expression blendshapes.
• is the global transformation matrix for point - are the skinning weights.

Recent Advancements

FLAME is not just a static model; it has become the core for modern facial research.

How is FLAME Trained?

Understanding the usage of FLAME is one thing, but how is such a model actually created? The training of FLAME is an iterative process of statistical learning from raw 3D data.

Implementing FLAME in MLX

Why MLX? As a specialized framework for Apple Silicon, MLX allows us to run these complex matrix operations directly on the GPU/NPU with unified memory. This means faster training and real-time inference on edge device with MacBook machines — the real reason is actually, I just want to maximize the utilization of my best local machine: Apple Silicon laptop 🙂

For those who are fluent with NumPy / JAX / PyTorch, MLX feels remarkably familiar. If you can write NumPy, you can write MLX. The API between NumPy and MLX is almost 1:1, with added power of composable function transformations like mx.grad and mx.vmap.

In MLX, memory is unified. One of the hurdles in PyTorch is managing .to(device) or .cuda() calls. There is no cost in MLX to move data between the CPU and GPU because they share the same physical memory, simplifying code and eliminates “device mismatch” errors.

Like JAX, MLS uses lazy evaluation. It builds on a computation graph and only executes it when needed (e.g., when we print a value or save a file). This allows for automatic graph optimizations without the complexity of manual JIT decorators in many cases.

Here is how we implement the Linear Blend Skinning forward pass with MLX:

def lbs(betas, pose, v_template, shapedirs, posedirs, J_regressor, parents, lbs_weights):
    # 1. Add shape, pose, and expression contributions
    v_shaped = v_template + blend_shapes(betas, shapedirs)
    
    # 2. Add pose-dependent blend shapes
    ident = mx.eye(3)
    rot_mats = batch_rodrigues(pose.reshape(-1, 3)).reshape(batch_size, -1, 3, 3)
    pose_feature = (rot_mats[:, 1:] - ident).reshape(batch_size, -1)
    pose_offsets = mx.matmul(pose_feature, posedirs).reshape(batch_size, -1, 3)
    
    v_posed = v_shaped + pose_offsets
    
    # 3. Apply the global joint transformations
    J_transformed, A = batch_rigid_transform(rot_mats, J, parents)
    
    # 4. Final skinning
    # ... matrix transforms ...
    return verts, J_transformed

Once we have the posed vertices from the lbs function, MLX makes it easy to implement downstream tasks. For instance, calculating 3D landmarks via barycentric interpolation, 3D joint locations from mesh vertices, and per-vertex displacements from blend shapes can be done using mx.einsum operation.

def vertices2landmarks(vertices, faces, lmk_faces_idx, lmk_bary_coords):
    """
    Calculates 3D landmarks by barycentric interpolation over specific mesh faces.
    
    Args:
        vertices (mx.array): Batch of mesh vertices of shape (B, V, 3).
        faces (mx.array): Mesh face indices of shape (F, 3).
        lmk_faces_idx (mx.array): Indices of faces where landmarks are located (B, L).
        lmk_bary_coords (mx.array): Barycentric coordinates of landmarks on those faces (B, L, 3).
        
    Returns:
        mx.array: Calculated 3D landmarks of shape (B, L, 3).
    """
    batch_size, num_verts = vertices.shape[:2]
    
    # Extract the indices of the vertices for each face
    lmk_faces = faces[lmk_faces_idx]
    
    # Add batch offset to indices for flat indexing
    batch_offset = mx.arange(batch_size)[:, None, None] * num_verts
    lmk_faces_absolute = lmk_faces + batch_offset
    
    # Gather vertices
    vertices_flat = vertices.reshape(-1, 3)
    lmk_vertices = vertices_flat[lmk_faces_absolute] # (B, L, 3, 3)
    
    # Compute landmarks using barycentric interpolation
    landmarks = mx.einsum('blfi,blf->bli', lmk_vertices, lmk_bary_coords)
    return landmarks

def vertices2joints(J_regressor, vertices):
    """
    Calculates 3D joint locations from mesh vertices using a regressor matrix.
    
    Args:
        J_regressor (mx.array): Joint regressor matrix of shape (J, V).
        vertices (mx.array): Batch of mesh vertices of shape (B, V, 3).
        
    Returns:
        mx.array: Calculated 3D joint locations of shape (B, J, 3).
    """
    return mx.einsum('bik,ji->bjk', vertices, J_regressor)

def blend_shapes(betas, shape_disps):
    """
    Calculates per-vertex displacements from blend shapes (identities or expressions).
    
    Args:
        betas (mx.array): Coefficients for the blend shapes of shape (B, K).
        shape_disps (mx.array): Blend shape bases of shape (V, 3, K).
        
    Returns:
        mx.array: Vertex displacements of shape (B, V, 3).
    """
    return mx.einsum('bl,mkl->bmk', betas, shape_disps)

Since MLX is fully differentiable, we can also easily compute gradient for fitting the model to 2D landmarks or 3D scans.

High-Performance Demos

With this implementation, we can achieve real-time interactivity. Below is a demo of our Viser-based 3D Viewer running with FLAME parameters:

Conclusion

FLAME remains the gold standard for controllable 3D face modeling. Its disentangled representation of identity and expression makes it perfect bridge between classic computer graphic and modern generative AI. By porting these models to high-performance frameworks on edge device like MLX, it’s bridging the digital human faces to every edge desktop computer.

Check out the full implementation on GitHub: https://github.com/ghif/FLAME_mlx.

Also, thanks to Google’s Antigravity for providing developer experience to the next level in assisting me to port the code.

From Rigid Triangles to Fuzzy Blobs

To understand why 3DGS is revolutionary, we must look at how it defines space.

In a traditional Triangle Mesh, an object is a hollow shell. To render a cat, you stretch a “skin” of texture over a wireframe skeleton. It works great for solid surfaces like walls or cars, but it fails at complex, thin, or semi-transparent structures. Have you ever noticed how video game hair often looks like stiff strips of paper? That is the limitation of the Triangle.

3DGS abandons the shell. Instead, it treats the world as a volumetric cloud.

Imagine a point cloud, but instead of tiny, single-pixel dots, every point is a 3D ellipsoid (a stretched sphere). We call these “Gaussians.” Each Gaussian is defined not just by a position, but by a set of learnable parameters that describe its existence in space.

If the Triangle is a piece of origami paper, the Gaussian is a soft, colored snowball.

Figure 1: Illustration of Multi-View Scene Reconstruction with 3D Gaussian Splatting taken from (Choi et al. 2025)

3DGS can enable a new class of product features that were previously impossible due to the “quality vs speed” trade-off. Previously, if a product needed photorealism, it used pre-rendered video (non-interactive). If it needed interactivity, it used polygon meshes (often lacking photorealism for complex materials). 3DGS bridges this gap.

Here are the specific functionalities in real-world products that can be enabled by 3DGS:

1. Next-Gen E-Commerce: “The Unscannable Product” Viewer

Traditional photogrammetry (Meshes) fails drastically when scanning objects with transparency, refraction, thin structures, or high reflectivity. This limits 3D product views to matte objects like shoes or furniture.

• The Enabled Functionality: An interactive web viewer for jewelry, perfume bottles, fluffy apparel, and complex electronics.
• Why 3DGS?
  • Refraction/Reflection: 3DGS captures view-dependent effects (shiny diamonds sparkle as you rotate the view) because it uses Spherical Harmonics.
  • Fuzzy Details: It captures the individual strands of a mohair sweater or the bristles of a toothbrush, which would otherwise be a flat texture on a blobby mesh.

2. Real Estate & Tourism: The “Cinematic” Virtual Tour

Current virtual tours (like Matterport) often rely on 360° panoramas. You can jump from point A to point B, but you cannot “walk” smoothly between them. Mesh-based navigations often look like low-poly video games.

• The Enabled Functionality: 6-Degrees-of-Freedom (6DoF) Walkthroughs on mobile devices. Users can fly a drone path through a hotel suite or walk through a house museum with photorealistic lighting preservation.
• Why 3DGS?
  • Lighting Baking: 3DGS bakes the complex bounce lighting of a room into the scene. Mirrors work correctly (reflecting the room), and windows show the actual view outside, rather than a flat white texture.
  • Mobile Performance: Unlike NeRFs, which kill mobile batteries, optimized 3DGS (like Luma AI or Spline) can render at 60fps on iPhones.

3. Telepresence: The Photorealistic Avatar

Video conferencing has looked the same for 15 years (2D grid). VR avatars (like Meta’s Horizon) look cartoonish because rendering realistic skin and hair on a mesh is computationally heavy.

• The Enabled Functionality: Holographic Telepresence. A system where you see the other person in full 3D volume, with realistic hair movement and clothing wrinkles, viewable from different angles in a VR/AR headset.
• Why 3DGS?
  • Hair & Eyes: These are the hardest parts of a human to render. 3DGS handles the thin geometry of hair and the wetness of eyes naturally.
  • GaussianAvatar: Technologies like GaussianAvatar bind splats to a skeleton, allowing for real-time animation of a photorealistic human scan.

4. VFX & Virtual Production: Instant Background Plates

In film production, creating a digital twin of a movie set for “Virtual Production” (LED walls) usually takes weeks of manual modeling and texturing by artists.

• The Enabled Functionality: Rapid Set Digitization. A drone flies over a location (e.g., a forest or city street), and within 20 minutes, a Director can scout that location in VR or use it as a background plate for a car chase scene.
• Why 3DGS?
  • Training Speed: A 3DGS scene trains in minutes, not days.
  • Unstructured Environments: It excels at organic nature (trees, bushes, dirt) which are notoriously difficult and expensive to model by hand with polygons.

5. Automotive: The Interactive Configurator

Car configurators usually rely on heavy ray-tracing (cloud streaming) or simplified WebGL meshes (loss of quality).

• The Enabled Functionality: Real-time Ray-Tracing Quality on Web. A user can inspect a car’s metallic paint flakes, the leather stitching, and the glass headlights in a browser without lag.
• Why 3DGS?
  • Anisotropic Splats: The “squashed” shape of Gaussians is perfect for representing the metallic sheen of car paint and the sharp specular highlights on the bodywork.

6. Robotics & Sim: Simulation Environments

Robots need to train in simulators. Mesh simulators are “too clean”—they don’t have the visual noise, dust, and lighting artifacts of the real world, leading to the “Sim-to-Real gap.”

• The Enabled Functionality: Digital Twin Simulation. A robot can train navigation inside a 3DGS scan of a warehouse that is visually indistinguishable from the real warehouse.
• Why 3DGS?
  • It captures the “messiness” of the real world (wires hanging down, reflective puddles on the floor) that confuses robot sensors, providing better training data than hand-made CAD models.

Anatomy of a Gaussian

The fundamental building block of the 3DGS representation is the 3D Gaussians / Ellipsoids. Unlike isotropic points or voxels used in previous explicit methods, 3D Gaussians are anisotropic, meaning they can be stretched and scaled along independent axes. This anisotropy is crucial for modeling surface geometry efficiently: a single flat Gaussian can represent a large patch of a wall, while thin, elongated Gaussians can represent wires or hair strands.

Mathematically, a 3D Gaussian is defined by a mean position and a covariance matrix . The influence of the Gaussian at any point in 3D space is given by the standard multivariate Gaussian distribution:

However, directly optimizing the covariance matrix is problematic because it must remain positive semi-definite to represent a valid physical ellipsoid. To enforce this constraint during gradient descent optimization, (Kerbl et al. 2023) proposes decomposing into a scaling matrix and a rotation matrix :

Here, is a diagonal matrix representing the scaling factors along the three axes, and is a rotation matrix constructed from a unit quaternion . This decomposition allows the optimizer to adjust the size and orientation of the splats independently, enabling the representation to “flatten” along surfaces and align with geometric structures effectively.

Quaternions are critical for efficient optimization. Unlike rotation matrices, which break and shear when modified by gradient descent, quaternions maintain valid rotations without requiring expensive correction algorithms like Gram-Schmidt.

So, the optimization utilizes quaternions representing the blob pose. To convert it into its corresponding rotation matrix, the following rule can be applied: for a unit quaternion , where is the scalar (real) part and are the vector (imaginary) parts, the corresponding rotation matrix is:

Technically, 3D Gaussian blobs can be generated randomly by initializing the mean (a single point coordinate) and covariances. In standard practice, however, 3D Gaussian blobs are almost never generated randomly for complex real-world scenes. Instead, the pipeline relies heavily on Structure-from-Motion (SfM) algorithms, most notably COLMAP (Schönberger et al. 2016). Before training begins, COLMAP analyzes the input images to calculate camera poses and generates a sparse point cloud of the scene. Each valid point in this cloud serves as the initial “seed” for a Gaussian blob — the point’s location becomes the Gaussian’s mean , while its scale and opacity are initialized to generic low values. This provides the optimizer with a rough geometric scaffold, preventing it from getting stuck in local minima which frequently happens if training starts from random noise.

The following is the JAX code snippet implementing the 3D Gaussian blob data structure.

import jax.numpy as jnp
import chex

@chex.dataclass
class Gaussians:
    means: jnp.ndarray  # (N, 3)
    scales: jnp.ndarray  # (N, 3)
    quaternions: jnp.ndarray  # (N, 4)
    opacities: jnp.ndarray  # (N, 1)
    sh_coeffs: jnp.ndarray  # (N, K, 3) where K is num SH coefficients

def init_gaussians_from_pcd(points: jnp.ndarray, colors: jnp.ndarray):
    """
    Initialize Gaussians from a point cloud.

    Args:
        points: (N, 3)
        colors: (N, 3) in [0, 1]
    Returns:
        gaussians: Gaussians dataclass
    """
    num_points = points.shape[0]
    
    # Position: mean of the point cloud
    means = points
    
    # Scales: log of the distance to the nearest neighbors
    # Initialized to a small value (approx 0.05m)
    scales = jnp.full((num_points, 3), -3.0) 
    
    # Rotations: identity quaternions [1, 0, 0, 0]
    quaternions = jnp.tile(jnp.array([1.0, 0.0, 0.0, 0.0]), (num_points, 1))
    
    # Opacities: inverse sigmoid of 0.5 = 0.0
    opacities = jnp.full((num_points, 1), 0.0) 
    
    # SH Coefficients (DC term only)
    # SH_DC = (R - 0.5) / 0.28209
    sh_dc = (colors - 0.5) / 0.28209479177387814
    sh_coeffs = jnp.zeros((num_points, 16, 3)) # Degree 3 SH -> 16 coefficients
    sh_coeffs = sh_coeffs.at[:, 0, :].set(sh_dc)
    
    return Gaussians(
        means=means,
        scales=scales,
        quaternions=quaternions,
        opacities=opacities,
        sh_coeffs=sh_coeffs
    )

Here we use the @chex.dataclass decorator for the structure due to several reasons:

• JIT compatibility (Pytree registration): In JAX, functions are often transformed using jax.jit, jax.vmap, or jax.grad. These transformations require that any data passed into them be a “Pytree” (a container of JAX arrays). A standard Python dataclass is not a Pytree by default. If we try to pass it to a JIT-compiled function, JAX won’t know hot to “look inside” it to find the arrays. @check.dataclass automatically registers the class as a JAX Pytree. This allows us to pass the Gaussians object directly into JIT-compiled functions as if it were a simple tuple or dictionary.
• Immutability: @check.dataclass are immutable by default (similar to frozen=True in standard dataclasses). JAX’s functional programming model relies on pure functions and immutable data.
• Ease of optimization: Because it’s a Pytree, we can use JAX’s gradient transformations directly on the object. For example:

def loss_fn(gaussians):
    # compute loss ...
    return loss

# This works because Gaussians is a chex.dataclass (Pytree). The resulting grads grads will have the same structure as the Gaussians class, containing gradients for means, scales. etc.
grads = jax.grad(loss_fn)(gaussians)

Note that there are two attributes that have not been introduced before: opacities and sh_coeefs. Those are needed from the rendering through rasterization that will be explained in the next section.

Below is the function to compute the covariance matrix from the quaternion of Gaussian blobs. Here we employ JAX vectorizing map (jax.vmap) to allow more efficient batched computation.

def get_covariance_3d(scales: jnp.ndarray, quaternions: jnp.ndarray):
    """
    Computes 3D covariance matrix from scales and quaternions.
    Σ = R S S^T R^T

    Args:
        scales: (N, 3)
        quaternions: (N, 4)
    Returns:
        covariance: (N, 3, 3)
    """
    # Normalize quaternions
    q = quaternions / jnp.linalg.norm(quaternions, axis=-1, keepdims=True)
    
    # Rotation matrix from quaternion
    r = q[:, 0]
    x = q[:, 1]
    y = q[:, 2]
    z = q[:, 3]
    
    R = jnp.stack([
        jnp.stack([1 - 2*y**2 - 2*z**2, 2*x*y - 2*r*z, 2*x*z + 2*r*y], axis=-1),
        jnp.stack([2*x*y + 2*r*z, 1 - 2*x**2 - 2*z**2, 2*y*z - 2*r*x], axis=-1),
        jnp.stack([2*x*z - 2*r*y, 2*y*z + 2*r*x, 1 - 2*x**2 - 2*y**2], axis=-1)
    ], axis=-2)
    
    # Scaling matrix
    s = jnp.exp(scales)
    S = jax.vmap(lambda x: jnp.diag(x))(s) # "Vectorize" operation. Since s has shape (N, 3), vmap applies the diag function to each of the N rows independently so that the operation becomes much more efficien
    
    # M = R S
    M = R @ S
    
    # Σ = M M^T
    Sigma = M @ M.transpose(0, 2, 1)
    
    return Sigma

The figures below illustrate the initialization of Gaussian blobs: random generation vs extraction from SfM.

Figure 2: 3D Gaussian splats/blobs generated randomly

Figure 3: 3D Gaussian splats induced by point clouds extracted from Structure-from-Motion (COLMAP)

Rendering through Rasterization: From 3D Cloud to 2D Image

A defining characteristic of photorealistic rendering is the ability to capture non-Lambertian effects—how the appearance of a surface changes based on the viewing angle (e.g., specular highlights on metal, the sheen of silk). Now that we have a scene populated with millions of initialized 3D Gaussians, we need to render them.

This is where 3DGS diverges sharply from NeRFs. Instead of shooting rays backwards from the camera into the scene (ray marching), 3DGS utilizes Forward Rasterization. It takes the existing 3D data and “pushes” it onto the 2D screen, similar to how standard game engines render triangles. 3DGS moves away from the neural decoding of color used in NeRF. Instead, each Gaussian stores a set of coefficients for Spherical Harmonics.

SHs serve as a basis function for representing functions on the surface of a sphere. By storing higher-order SH coefficients (typically up to degree 3, resulting in 16 coefficients per color channel), 3DGS can approximate complex view-dependent lighting effects directly. During rendering, the viewing direction vector is used to evaluate the SH functions, producing a specific RGB color for that angle. This explicit storage contributes to the high memory footprint of 3DGS but eliminates the need for expensive MLP inference during the render pass.

Here is the step-by-step mathematical flow.

1. The 2D Projection / Viewing Transformation (World —> Camera)

First, we transform the geometry relative to the camera. If the camera moves forward, the Gaussians effectively move backward.

For the position mean , this is a standard matrix multiplication using the camera view matrix (pose):

However, transforming the covariance matrix — the shape of the blob — is trickier. Because perspective projection is non-linear (objects shrink as the get further away), we cannot simply rotate the covariance.

To solve this, we use a linear approximation: Jacobian of the affine projection (). This matrix encodes how 3D space compresses into 2D space based on the camera’s focal length and the point’s depth .

The projection equation: we compute the new 2D covariance by “sandwiching” the 3D covariance between the transformation matrices:

• : The original shape of the blob
• : The rotation matrix of the camera
• : The Jacobian approximation

The result, , is a covariance matrix that describes a flat 2D Ellipse on your screen. If the calculated size of this 2D splat is smaller than a single pixel, we usually discard it to save performance.

The code snippet below implements the 2D projection of 3D Gaussian blobs given camera parameters using JAX:

def project_gaussians(gaussians: Gaussians, camera: Camera):
    """
    Project 3D Gaussians to 2D splats.

    Args:
        gaussians: Gaussians dataclass
        camera: Camera dataclass 
    Returns:
        means2D: 2D means of the projected splats
        cov2D: 2D covariance of the projected splats
        radii: Radii of the projected splats
        valid_mask: Valid mask for the projected splats
        z: Depth of the projected splats
    """
    means3D = gaussians.means
    scales = gaussians.scales
    quats = gaussians.quaternions
    
    # 1. Transform means
    means3D_homo = jnp.concatenate([means3D, jnp.ones((means3D.shape[0], 1))], axis=-1)
    means_cam = (means3D_homo @ camera.W2C.T)[:, :3]
    
    x, y, z = means_cam[:, 0], means_cam[:, 1], means_cam[:, 2]
    
    # 2. Filter 
    valid_mask = z > 0.01
    
    # 3. Covariance
    cov3D = get_covariance_3d(scales, quats)
    
    # 4. Project to 2D
    # Jacobian of the perspective transformation
    J = jnp.zeros((means3D.shape[0], 2, 3))
    J = J.at[:, 0, 0].set(camera.fx / z)
    J = J.at[:, 0, 2].set(-camera.fx * x / (z**2))
    J = J.at[:, 1, 1].set(camera.fy / z)
    J = J.at[:, 1, 2].set(-camera.fy * y / (z**2))
    
    W_rot = camera.W2C[:3, :3]
    
    def project_single_cov(c3d, j_mat):
        return j_mat @ W_rot @ c3d @ W_rot.T @ j_mat.T
    
    cov2D = jax.vmap(project_single_cov)(cov3D, J)
    # Add a small bias for numerical stability (low pass filter)
    cov2D = cov2D.at[:, 0, 0].add(0.3)
    cov2D = cov2D.at[:, 1, 1].add(0.3)
    
    # 5. Means 2D
    means2D = jnp.stack([
        camera.fx * x / z + camera.cx,
        camera.fy * y / z + camera.cy
    ], axis=-1)
    
    # 6. Radii for tile interaction
    det = cov2D[:, 0, 0] * cov2D[:, 1, 1] - cov2D[:, 0, 1]**2
    trace = cov2D[:, 0, 0] + cov2D[:, 1, 1]
    mid = trace / 2.0
    term = jnp.sqrt(jnp.maximum(mid**2 - det, 0.0))
    lambda1 = mid + term
    max_eigen = lambda1 
    radii = jnp.ceil(3.0 * jnp.sqrt(max_eigen))
    
    return means2D, cov2D, radii, valid_mask, z

1. Sorting (The Depth Problem)

Since we are blending transparent blobs, order matters, i.e., Red behind Blue Blue behind Red. Once all Gaussians are projected to 2D, the engine performs a Global Sort based on their depth ( -value) relative to the camera. In the official CUDA implementation, this uses a highly optimized GPU Radix Sort.

• Input: Unsorted list of 2D splats
• Output: List of indices sorted from closest to furthest.

1. Tile-based Rasterization

A 4K screen has 8 million pixels. A scene might have 4 million Gaussians. A naive loop (for every pixel —> check every Gaussian) would be 8M 4M operations. That is impossible.

3DGS uses a tiled approach, similar to modern GPU rasterizers:

• The screen is divided into pixel tiles.
• Identify which tiles a Gaussian overlaps with.
• Each tile creates a generic “list” of Gaussians that touch it.

Now, a pixel only needs to check the Gaussians in its specific tile list, reducing complexity drastically.

1. Alpha Blending (The Compositing Equation)

This is the final step of the rendering: color calculation. For a single pixel, we iterate through the sorted Gaussians in its tile.

The color of a pixel is computed using standard front-to-back composition:

Let’s break down the terms:

• : The color of the -th Gaussian (from Spherical Harmonics)
• : The effective opacity. This combines the learned opacity with the Gaussian falloff distance from the center of the blob:

(Pixels at the center of the blob are opaque; pixels at the edge are transparent)

• : The transmittance. This represents “how much light has not been blocked yet.”

The optimization (early stopping): Because it renders front-to-back, we can tract the transmittance :

• Start with (100% light passes).
• After Gaussian A (), drops to .
• After Gaussian B (), drops to - Stop condition: if , the pixel is fully saturated (opaque). We stop processing the list. This ignores objects hidden behind a wall, saving massive computation.

Training Loop: How to Optimize the Blob Cloud

Unlike a neural network, which optimizes “weights” inside a black box, 3DGS optimizes the scene geometry directly. The training process is a continuous cycle of rendering the scene, comparing it to a real photo, and nudging the blobs to match.

Learnable Parameters

Recall equation (1) where a Gaussian is expressed with In a 3DGS scene, we have a set of Gaussians, where starts around 100k and grows to millions. Each Gaussian is parameterized by a tuple of learnable variables :

where each represents:

• Position : World coordinates with learning rate scaled by scene extent.
• Scale : Log-scale vector stored as to enforce .
• Rotation : Unit quaternion normalized every step to ensure valid rotation.
• Opacity : Scalar probability stored as logit passed through Sigmoid.
• Color : Spherical harmonics, usually degree 3 ( coeffs).

As discussed earlier, we do not optimize the covariance matrix directly, because it must remain positive semi-definite. Instead, we optimize the scale and rotation, and construct the covariance matrix on the fly using equation (2).

Forward Pass (Differentiable Rasterization)

The forward pass process is basically the differentiable pipeline sequence of rendering through rasterization described in the previous section, from 2D projection to alpha blending. The final output of this pipeline produces a rendered RGB image .

For each pixel location in , the value is rendered by the compositing equation (3). We can rewrite the equation by introducing the pixel location and learnable parameters:

where

Denote by the entire image induced the Gaussian’s learnable parameters .

Gradient-based Optimization

The rendered image is then compared against the ground truth training image . 3DGS utilizes a loss function in the form of a combination of L1 distance (for raw pixel accuracy) and D-SSIM (Structural Similarity Index, for perceptual structure).

Standard implementation uses .

The objective of 3DGS is therefore minimizing the loss function in equation (5) with respect to the gradient , which guides the rendering quality closer to the ground truth. This tells every Gaussian blob exactly how to move, stretch, or change color to minimize the error. The gradient computation can be easily implemented through JAX autograd.

Here is the JAX code snippet in implementing the gradient-based training step, using Optax library for the optimization helper.

import jax
import optax
from functools import partial
from jax_gs.renderer.renderer import render
from jax_gs.training.losses import l1_loss
from jax_gs.core.camera import Camera
import jax.numpy as jnp

@partial(jax.jit, static_argnums=(3, 4))
def train_step(state, target_image, w2c, camera_static, optimizer):
    """
    Standard training step.
    Args:
        state: (params, opt_state)
        target_image: (H, W, 3)
        w2c: (4, 4)
        camera_static: (W, H, fx, fy, cx, cy)
        optimizer: optax optimizer
    Returns:
        (next_params, next_opt_state), loss 
    """
    params, opt_state = state
    W, H, fx, fy, cx, cy = camera_static
    
    # Reconstruct Camera object inside JIT    
    camera = Camera(W=W, H=H, fx=fx, fy=fy, cx=cx, cy=cy, W2C=w2c, full_proj=jnp.eye(4))
    
    def loss_fn(p):
        image = render(p, camera)
        return l1_loss(image, target_image)
    
    loss, grads = jax.value_and_grad(loss_fn)(params)
    updates, next_opt_state = optimizer.update(grads, opt_state, params)
    next_params = optax.apply_updates(params, updates)
    
    return (next_params, next_opt_state), loss

Adaptive Density Control

Standard Gradient Descent changes the values of parameters, but it cannot change the number of parameters. If a scene needs more geometry (e.g., a complex hair strand), moving existing blobs isn’t enough.

Every 100 iterations, 3DGS performs Density Control based on the gradients of the positions.

A. Densification (Growth): Look for Gaussians with high positional gradients (blobs that are “struggling” to fit the data).

1. Clone: if the Gaussian is small, we duplicate it — fills empty holes.
2. Split: if the Gaussian is large, we split it into 2 smaller Gaussians and reduce their scale by factor 1.6 — refines blurry details.

B. Pruning (Cleanup): Remove Gaussians that are redundant or artifacts:

1. Low Opacity: if (transparent), delete it.
2. Screen-Space Size: if a Gaussian is huge (covering the whole screen), delete it.

This dynamic growing and pruning allows the optimization to start with a sparse point cloud and organically evolve into a dense, detailed representation.

Minimal Implementation of 3DGS with JAX

We have covered the theory: the Gaussian structure, spherical harmonics, and differentiable rasterization pipeline. Now, it is time to see how these equations translate into actual, executable code.

I have prepared a minimal, educational implementation of this engine using JAX, available on Github: https://github.com/ghif/jax-gs. I use agentic AI workflow through Google’s Antigravity IDE for helping me writing and structuring the codebase.

This repository is designed to be a “mathematical playground.” It strips away the complex C++/CUDA kernels found in the original implementation, replacing them with pure Python and JAX. This makes the codebase significantly shorter and easier to debug.

The beauty of building with JAX is the hardware agnosticims provided by the XLA (Accelerated Linear Algebra) compiler. We don’t need to rewrite the rasterizer to switch between hardware backends.

Here’s an example of visualizing the training progress using the codebase on a Fern scene from LLFF dataset.

**# CPU (Apple M4)**
tests/test_benchmark_training.py
Loading Fern dataset for training benchmark from data/nerf_example_data/nerf_llff_data/fern...
Benchmarking training step with 10091 Gaussians at 504x378 resolution...
Warming up (JIT compilation of train_step)...
Warm-up/Compilation took 0.8526s
Running benchmark (10 training iterations)...
Iteration 1: 0.8092s (Loss: 0.4832)
Iteration 2: 0.3521s (Loss: 0.4755)
Iteration 3: 0.3724s (Loss: 0.4723)
Iteration 4: 0.3697s (Loss: 0.4645)
Iteration 5: 0.3752s (Loss: 0.4633)
Iteration 6: 0.3748s (Loss: 0.4597)
Iteration 7: 0.3683s (Loss: 0.4649)
Iteration 8: 0.3685s (Loss: 0.4717)
Iteration 9: 0.3726s (Loss: 0.4786)
Iteration 10: 0.3738s (Loss: 0.4835)

Training Benchmark Result (Fern):
Average Iteration Time: 0.3697s
Average Speed:          2.70 it/s

# **GPU (NVIDIA L4 through Vertex AI Workbench)
tests/test_benchmark_training.py
Loading Fern dataset for training benchmark from data/nerf_example_data/nerf_llff_data/fern...
Benchmarking training step with 10091 Gaussians at 504x378 resolution...
Warming up (JIT compilation of train_step)...
Warm-up/Compilation took 5.7437s
Running benchmark (10 training iterations)...
Iteration 1: 3.6011s (Loss: 0.4832)
Iteration 2: 0.0397s (Loss: 0.4754)
Iteration 3: 0.0397s (Loss: 0.4722)
Iteration 4: 0.0396s (Loss: 0.4645)
Iteration 5: 0.0396s (Loss: 0.4633)
Iteration 6: 0.0395s (Loss: 0.4597)
Iteration 7: 0.0395s (Loss: 0.4649)
Iteration 8: 0.0396s (Loss: 0.4717)
Iteration 9: 0.0396s (Loss: 0.4786)
Iteration 10: 0.0395s (Loss: 0.4835)

Training Benchmark Result (Fern):
Average Iteration Time: 0.0396s
Average Speed:          25.26 it/s**

State of Audio-Driven Talking Faces

Crossing the “Uncanny Valley”—generating a digital face indistinguishable from a human—has been the “boss level” of computer graphics for decades. However, the last few years have seen the industry make a massive leap from “significant progress” to “solved problem.”

For a long time, the standard approach was simple lip synchronization: essentially gluing a moving mouth onto a static face. If you followed the early explosion of “talking head” AI, you likely remember models like Wav2Lip (Prajwal et al. 2020). While they were technically impressive at matching phonemes to mouth shapes, the results were often unsettling — the mouth moved perfectly, but the eyes were dead and the head was frozen. It was accurate, but robotic.

Research from 2024 through 2025 has rewritten the playbook. We are no longer just syncing lips: we are generating holistic facial dynamics. Here is the breakdown of how AI avatars woke up.

The Workflow at a Glance

Creating a convincing talking head isn’t just one button press: it is a pipeline. Here is the workflow we will cover:

1. Speech script preparation: Crafting the narrative
2. Voice synthesis and cloning: Generating the audio driver
3. Initial video generation: Bringing the static image to life (blinking, head movement)
4. Voice dubbing/face reenactment: The final lip-syncing using MuseTalk

Stop calling them chatbots. That era is over. The new trend isn't just about Language Models anymore; it’s about World Models.

Here is the difference. An LLM predicts the next word in a sentence. A World Model predicts the next second of reality. It understands 3D space, gravity, and consequence. Here are the big players building the Matrix.

First, meet the Dreamer: Genie 3 from Google DeepMind. This isn't just a video generator; it is an interactive game engine. You can actually grab a controller and play inside the video it creates. It generates the world and the gameplay in real-time. It is essentially an infinite holodeck.

Second, the Architect: Marble from World Labs. Most AI video is like a hallucination; it warps and shifts. Marble creates solid, persistent 3D worlds. You can take a photo of a room, and Marble turns it into a 3D asset you can actually edit and walk around in. This gives AI true Spatial Intelligence.

Finally, the Engineers: NVIDIA Cosmos and PAN from MBZUAI. These are the brains for robots. They don't just see video; they understand physics. They allow robots to preview the future, simulating what happens if they drop a cup or crash a car before they actually do it.

Genie dreams it, Marble builds it, and NVIDIA simulates the physics. We are moving from AI that writes about the world, to AI that lives in it. Which is scarier to you: An infinite game, or a robot that can predict the future? Let me know in the comments.

Generate an animated face video of this portrait with natural and gentle eye blinking as well as minimal head movement. Don't make it talk.

# Step 1: Create a new conda environment in a persistent directory (under /home/jupyter/)
conda create -p /home/jupyter/envs/MuseTalk python==3.10 -y

# Step 2: Activate the new environment
conda activate /home/jupyter/envs/MuseTalk

# Step 3: Clone the MuseTalk's source code and go to the root directory
git clone https://github.com/TMElyralab/MuseTalk.git
cd MuseTalk

# Step 4: Install all required libraries
pip install -r requirements.txt
pip install --no-cache-dir -U openmim
mim install "mmcv==2.0.1"
mim install "mmdet==3.1.0"
mim install --no-build-isolation "mmpose==1.1.0"

# Step 5: Remove huggingface-hub and gdown libraries inside download_weights.sh to avoid library's version conflicts
# Open download_weights.sh file and delete lines associated with huggingface-hub and gdown.

# Step 6: Download all the model's weights
sh ./download_weights.sh

# Step 7: Configure the test.yaml to specify the inputs and run the inference script
sh ./inference.sh v1.5 normal

# Go to home directory
cd ~

# Create a "bin" directory
mkdir -p ~/bin

# Download the source
wget https://johnvansickle.com/ffmpeg/releases/ffmpeg-release-amd64-static.tar.xz -O ~/ffmpeg-static.tar.xz

# Extract the source files to "bin"
tar -xf ~/ffmpeg-static.tar.xz -C ~/bin/

# Embed the path to the PATH env variable
echo 'export PATH="$HOME/bin/ffmpeg-7.0.2-amd64-static:$PATH"' >> ~/.bashrc

# Activate the new PATH
source ~/.bashrc

# Go back to conda env
conda activate /home/jupyter/envs/MuseTalk

Final Content

By combining a thoughtful script, cloned voice, AI-generated video motion, and MuseTalk’s precision lip-syncing, we get a final video that crosses the uncanny valley.

An example of full generated video with minor post-editing can be seen below.

https://www.youtube.com/watch?v=GIsw39CROow

Another example with a different speech script:

https://www.youtube.com/watch?v=mrQnMZWh-3Q

Personal AI Doctor Web App

The personal AI doctor I built is just a simple web-app prototype featuring a conversational UI that supports both chat and voice commands. It runs entirely on a local machine without any Internet connection. Users can also upload images alongside particular instructions for the virtual doctor to perform medical image analysis. The system can respond to user prompts with text and, optionally, speech output, as illustrated in Figure 1.

Figure 1: Personal AI Doctor Demo

The prototype is fully written in Python, comprising frontend and backend services. The frontend implements the user interface components, while all relevant AI models will be served in backend services, comprising:

• Large multimodal model (MedGemma-4B)
• Voice-to-text model (Whisper)
• Text-to-Speech (VITS through Coqui TTS)

MedGemma as The Foundation Model

MedGemma is a collection of open-model variants developed by Google DeepMind specifically optimized for medical-domain tasks involving both text and images. Built on the underlying architecture of the Gemma 3 family, it brings healthcare-focused capabilities into a freely accessible model ecosystem.

MedGemma currently comes in two variants: a 4B-parameter multimodal version (capable of ingesting medical images + text) and a 27B-parameter text-only (and/or multimodal) version. The multimodal variant uses an image encoder, SigLIP (Zhai et al., 2023), which was pre-trained on a large corpus of de-identified medical imager (e.g., chest X-rays, dermatology photos, pathology pictures).

Because MedGemma inherits the capabilities of Gemma 3 (multimodality, long context, efficient architecture), developers can build healthcare applications that have strong baseline performance and benefit from the engineering work behind Gemma 3. Therefore, to further understand MedGemma, it’s worth understanding Gemma 3’s architecture in more details.

Personal Doctor System

# Pull the specific 6-bit quantized MedGemma model
ollama pull amsaravi/medgemma-4b-it:q6

# Run the model to start the local inference server
ollama run amsaravi/medgemma-4b-it:q6

from google.adk.agents import Agent
from google.adk.models.lite_llm import LiteLlm

model_name = "ollama_chat/amsaravi/medgemma-4b-it:q6"
model = LiteLlm(model=model_name)

chat_agent = Agent(
    model=model,
    name="medical_chat_service",
    description="An agent that answers medical queries, optionally with images.",
    instruction="""You are a helpful medical assistant. 
    Answer the user's questions accurately and concisely. Always answer in English.
    If an image is provided, analyze it in the context of the medical question.
    """
)

summary_agent = Agent(
    model=model,
    name="medical_summary_service",
    description="An agent that summarizes medical advice into concise paragraphs..",
    instruction="""
    You are an expert medical summarizer. 
    Your goal is to summarize the provided medical advice into a concise, single paragraph. 
        Ensure the summary flows naturally when read aloud. 
        YOU MUST USE periods (.) and commas (,) to create natural pauses. 
      DO NOT use other punctuation marks like question marks (?), exclamation marks (!), colons (:), or semicolons (;). 
        Keep the summary focused on the key medical advice and safety instructions.
    """
)

Voice Command Capability (STT Service)

To improve user interactiveness, it will be more handful if users can talk directly with the AI assistant rather than typing only, just like in ChatGPT or Gemini applications. To allow this capability, we can use a speech-to-text (STT) or automatic speech recognition (ASR) model to transcribe user’s speech into text, which will then be processed by MedGemma.

One of the most capable open-source ASR models is Whisper developed by OpenAI (Radford et al. 2022). It is an innovative STT focused on achieving robust speech recognition via large-scale weak supervision. Its primary distinguishing factor is the immense scale of its training data and resulting robustness:

• Massive data scale: Whisper was scaled up to 680,000 hours of labeled audio data, which includes weak supervision like multilingual and multitask supervision (see Figure 5 below).
• Focus on zero-shot transfer: The system aims to work reliably “out of the box” in a broad range of environments without requiring supervised fine-tuning for every new deployment distribution.
• Superior robustness: Whisper exhibits fundamentally different robustness properties; the best zero-shot Whisper models approach human accuracy and robustness when evaluated across various datasets.

Figure 5: Whisper leverages a seq2seq Transformer model trained on many different speech processing tasks, including multilingual speech recognition, speech translation, spoken language identification, and voice activity detection.

OpenAI offers several variants of the Whisper model to balance speed and accuracy, ranging from tiny (39M parameters) to large (1550M parameters), as well as the optimized turbo model (809M). Depending on the hardware constraints and accuracy requirements, we can select the appropriate size.

Here is a code example demonstrating how to implement audio transcription using Whisper in Python:

import whisper

model = whisper.load_model("turbo")
...
result = model.transcribe(audio_file_path) # transcribe text from an audio file

Voice Synthesis (TTS Service)

To elevate the user experience beyond simple text display, we can give our AI system a voice. It needs a functionality to synthesize speech from text. For this, I utilize Coqui TTS, a popular open-source deep learning toolkit designed to run advanced speech synthesis models directly on local hardware. While the company behind Coqui has sadly ceased operations, the library remains a powerful, community-maintained tool.

Coqui TTS offers a wide registry of pretrained deep learning models, making it easy to find the right fit for our needs. In this setup, I selected the VITS model (specifically tts_models/en/ljspeech/vits). VITS (Variational Inference with adversarial learning for end-to-end Text-to-Speech) is a parallel, end-to-end architecture known for generating highly natural sounding audio (Kim et al. 2021).

Unlike traditional TTS pipelines that require two separate stages — generating acoustic features (like mel-spectrograms) first, and then using a vocoder to create the waveform — VITS unifies these processes. It employs a Variational Autoencoder (VAE) to connect these modules through latent variables, allowing for efficient, high-fidelity waveform generation in a single framework.

Figure 6: Training and inference mechanism of VITS

Functionally, VITS employs Monotonic Alignment Search (MAS) during the training phase to automatically eliminate the alignment between the input text phonemes and the target speech, removing the need for external aligners — Figure 6 shows the difference between the training and inference phases of VITS. Because the model is non-autoregressive, it allows for parallel sampling, which significantly improves synthesis speed compared to sequential models.

The following code snippet illustrates the utilization of VITS model through Coqui TTS.

import torch
from TTS.api import TTS # Importing Coqui TTS
...

device = "cuda" if torch.cuda.is_available() else "cpu"
model_name = "tts_models/en/ljspeech/vits" # Using a VITS model
tts = TTS(model_name).to(device)

text = " ... "
tts.tts_to_file(text=text, file_path=file_path)

Spec-Driven Development with Spec Kit (Optional)

I’ve marked this section as “optional” because we don’t have to adopt this approach (Spec-Driven Development) when building the complete application — we could simply follow conventionally way of application development. I’m including it mainly out of curiosity on trying a new productivity tool.

To build the complete system locally, I use Spec Kit — an open-source toolkit that enables Spec-Driven Development (SDD), essentially extending CLI-driven AI agents into a more structured workflow. It’s a bit like the freedom of “vibe coding”, but with much more control and editability through clearly defined specifications. With Spec Kit, we begin with the specification first, and let the code follow — keeping product requirements, intentions, and architecture front and center. The unified source empowers both AI assistants and developers to implement features with consistent outcomes.

For me, adopting this approach has been a breakthrough. While many “vibe-coding” tools can spin off in unpredictable directions, Spec Kit introduces the right amount of structure and guardrails. It allows developers to collaborate more effectively especially with product managers / owners, align with real software-development phases, and retain control—even while harnessing AI capabilities.

Figure 7: Spec-Driven Development Life Cycle with Spec

Figure 7 illustrates the life cycle steps in SDD. In general, within the workflow we will:

1. Define the “what”: Using the /speckit.specify command, we describe the feature in plain English — what we want to build, and why — without worrying about the tech stack yet. Outputs: spec.md
2. Plan the “how”: With /speckit.plan , we pick the architecture, libraries, database, and overall structure. This step translates the spec into a technical implementation blueprint. Outputs: plan.md, quickstart.md, research.md, data-model.md, contracts/openapi.json.
3. Break into tasks & implement: Use /speckit.tasks to split the plan into actionable slices, then let /speckit.implement (with AI assistants or our team) generate the code. Specs remain alive and authoritative through process. Outputs: tasks.md and your generated codes.

The /speckit.constitution step is mostly static — you’ll probably only need to run it once at the very beginning. Think of it as setting the guiding principles: it defines the ground rules for what the AI should or shouldn’t do in subsequent steps. The real iterative action happens between /speckit.specify and /speckit.implement, where ideas evolve into actual features. Each time we run /speckit.specify, it spins up a new Git branch — just like when we’re introducing a quite major change to the application.

It is worth noting that the /speckit.clarify, which is an optional step, is useful to clarify the initially generated specification, involving a human-in-the-loop workflow for detailing some information or reducing ambiguity.

Under the hood, SpecKit relies on an AI code assistant tool as the core engine such as GitHub Copilot, Gemini CLI, Claude Code, Codex CLI, Cursor, and so on. I use Gemini CLI for my main AI code assistant in this entire project.

Conclusion

We have walked through a simple, complete pipeline of building a sovereign healthcare AI—from serving MedGemma locally with Ollama to enabling speech interaction via Whisper and Coqui TTS.

While this prototype highlights the potential of the MedGemma 4B model, the possibilities for expansion are vast. You could extend the agents to handle actual medical records, fine-tune the TTS voice for a more empathetic persona, or deploy the system on an edge device for portability. The tools for building private, powerful AI are now in your hands—it’s time to build something that protects as much as it assists.

The complete code implementation of this prototype is available on GitHub: https://github.com/ghif/personal-doctor.

Text-to-Speech

Have you ever had your phone read an article to you or heard audiobooks? That’s Text-to-Speech (TTS) in action, the tech that turns written text into spoken words. TTS has come a long way — from flat, robotic voices to ones that can whisper, pause, emphasize, or sound emotional. It’s everywhere now: voice assistants, screen readers, podcasts, navigation, etc.

TTS also has a long history, dating back to the early 20th century in the analog era. Here’s a brief overview of its evolution:

Modern TTS Pipeline

Modern TTS has evolved far beyond robotic-sounding voices of the past. Today, it can produce speech nearly indistinguishable from a human speaker. It typically involves 3 main stages:

1. Text Analysis: The input text is normalized (e.g., numbers —> words, punctuation —> pauses) and converted into phonemes or linguistic representations.
2. Acoustic Modeling: A neural model predicts mel-spectrograms — visual representations of sound frequency over time — which capture the tonal and rhythmic qualities of speech.
3. Vocoder: Another neural model converts the spectrogram into a raw audio waveform.

Figure 1: Modern TTS Pipeline

By harnessing the power of neural networks, recent advancements have simplified the TTS pipeline into a 2-stage approach. Some methods skip the explicit acoustic modeling step altogether — letting deep vocoders generate waveforms directly from linguistic features such as WaveNet (Oord et al., 2016), Parallel WaveNet (Oord et al., 2017), DeepVoice 1 (Arik et al., 2017), DeepVoice 2 (Arik et al., 2017), or HiFi-GAN (Kong et al. 2020). Others take a different route, using a single deep acoustic model to produce Mel-Spectrograms straight from text, as seen in systems like DeepVoice 3, TransformerTTS (Li et al., 2019), FastSpeech 1 (Ren et al., 2019), or FastSpeech 2 (Ren et al., 2022), which is a dominant approach recently. Another branch of work even attempts to solve the TTS task in an end-to-end manner such as WaveGrad 2 (Chen et al., 2021), FastDiff (Huang et al., 2022).

Voice Cloning

Speech synthesis becomes even more fascinating and useful when it can mimic a real person’s unique voice. Imagine this: you record yourself speaking for a few seconds, then type any sentence you want — and boom, it’s spoken back in your own voice. Voice cloning / conversion is turning that into reality. It’s one of the cooler (and also ethically complex) branches of speech AI.

Voice cloning can be understood as an extension of TTS systems, where the synthesized voice is personalized to match a specific individual. Achieving this requires careful modification of the acoustic modeling or, possibly, vocoder components to capture the unique vocal characteristics of the target speaker. However, the definition of voice cloning itself has varied across research efforts. (Azzuni and Saddik, 2025) proposed a more structured taxonomy that clarifies its various forms:

1. Voice Cloning: Replicating a specific person’s voice using a TTS system
2. Speaker Adaptation: Fine-tuning a pre-trained TTS model to replicate a specific user’s voice using limited data.
3. Few-shot Voice Cloning: A special case of speaker adaptation, where the amount of data or reference audio is much smaller — typically from just a few seconds up to a maximum of 5 minutes.
4. Zero-shot Voice Cloning: Adapting speaker’s voice without model finetuning at all, which is more challenging but highly desirable for real-world scalability and flexibility.

Figure 3: Taxonomy of Voice Cloning

Building a Simple TTS with Voice Cloning

As part of my personal content creation toolkit, I want to build a lightweight TTS that can run locally on my machine or on a cheap cloud environment like Google Colab. To begin, I utilize Chirp 3 to generate speech directly from scripted text.

Chirp 3 represents Google’s recent generation of its neural speech models, offering both Speech-to-Text (ASR / transcription) and Text-to-Speech (TTS) with high-definition voice synthesis. It is available through Google Cloud / Vertex AI, supporting a large number of pre-defined voices (~248 distinct voices) across many languages (~31 locales), each crafted to deliver realism and emotional expressiveness.

Although technical documentation about Chirp’s underlying architecture remains limited, it likely to follow the neutral TTS pipeline described earlier, with some modifications or advancements.

Chirp 3 also provides voice cloning via its Instant Custom Voice feature. However, since this capability is still in private preview, I opted to use another solution for the voice cloning process borrowed from an open-source model.

The overall pipeline can be illustrated as follows:

Figure 4: TTS with Voice Cloning

In this setup, the voice cloning module functions as a post-processing step, where it receives two audio inputs (source and reference speech) and generates a target speech output that conveys the linguistic content of the source while adopting the vocal characteristics of the reference.

from google.cloud import texttospeech_v1beta1 as texttospeech
from google.api_core.client_options import ClientOptions

...

def generate_speech(text, voice="Leda", language_code="id-ID"):
    """
    Generate speech using Google's Chirp 3 Text-to-Speech API

    Args:
        text (str): The text to be synthesized.
        voice (str): The name of the voice to use (e.g., "Leda", "Charon").
        language_code (str): The language code for the voice (e.g., "id-ID" for Indonesian, "en-US" for English).
    Returns:
        audio_content (bytes): The synthesized speech audio in MP3 format.
        temporary_file_name (str): The filename of the synthesized audio.
        
    """
    # Obtain default credentials and create an authorized session
    voice_name = f"{language_code}-Chirp3-HD-{voice}"
    voice = texttospeech.VoiceSelectionParams(
        name=voice_name,
        language_code=language_code,
    )

    # Initialize the TTS client with the specified API endpoint
    client = texttospeech.TextToSpeechClient(
        client_options=ClientOptions(api_endpoint=API_ENDPOINT)
    )

    # Perform the text-to-speech request on the text input with the selected
    # voice parameters and audio file type
    response = client.synthesize_speech(
        input=texttospeech.SynthesisInput(text=text),
        voice=voice,
        audio_config=texttospeech.AudioConfig(
            audio_encoding=texttospeech.AudioEncoding.MP3
        ),
    )

    # Create a temporary MP3 file
    audio_content = response.audio_content

    with tempfile.NamedTemporaryFile(delete=False, suffix=".mp3") as tmp_file:
        tmp_file.write(audio_content)
    
    schedule_cleanup(tmp_file.name)
    return audio_content, tmp_file.name

from pathlib import Path

import librosa
import torch
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file
from models.s3gen import S3Gen

S3_SR = 16_000
S3GEN_SR = 24_000

# Hugging Face repo ID for the pre-trained models
REPO_ID = "ResembleAI/chatterbox"

class VoiceCloner:
    ENC_COND_LEN = 6 * S3_SR
    DEC_COND_LEN = 10 * S3GEN_SR

    def __init__(
        self,
        s3gen: S3Gen,
        device: str,
        ref_dict: dict=None,
    ):
        self.sr = S3GEN_SR
        self.s3gen = s3gen
        self.device = device
        
        if ref_dict is None:
            self.ref_dict = None
        else:
            self.ref_dict = {
                k: v.to(device) if torch.is_tensor(v) else v
                for k, v in ref_dict.items()
            }

    @classmethod
    def from_local(cls, ckpt_dir, device) -> 'VoiceCloner':
        ckpt_dir = Path(ckpt_dir)
        
        # Always load to CPU first for non-CUDA devices to handle CUDA-saved models
        if device in ["cpu", "mps"]:
            map_location = torch.device('cpu')
        else:
            map_location = None
            
        ref_dict = None
        if (builtin_voice := ckpt_dir / "conds.pt").exists():
            states = torch.load(builtin_voice, map_location=map_location)
            ref_dict = states['gen']

        s3gen = S3Gen()
        s3gen.load_state_dict(
            load_file(ckpt_dir / "s3gen.safetensors"), strict=False
        )
        s3gen.to(device).eval()

        return cls(s3gen, device, ref_dict=ref_dict)

    @classmethod
    def from_pretrained(cls, device) -> 'VoiceCloner':
        # Check if MPS is available on macOS
        if device == "mps" and not torch.backends.mps.is_available():
            if not torch.backends.mps.is_built():
                print("MPS not available because the current PyTorch install was not built with MPS enabled.")
            else:
                print("MPS not available because the current MacOS version is not 12.3+ and/or you do not have an MPS-enabled device on this machine.")
            device = "cpu"
            
        for fpath in ["s3gen.safetensors", "conds.pt"]:
            local_path = hf_hub_download(repo_id=REPO_ID, filename=fpath)

        return cls.from_local(Path(local_path).parent, device)

    def set_target_voice(self, wav_fpath):
        ## Load reference wav
        s3gen_ref_wav, _sr = librosa.load(wav_fpath, sr=S3GEN_SR)
        s3gen_ref_wav = s3gen_ref_wav[:self.DEC_COND_LEN]
        self.ref_dict = self.s3gen.embed_ref(s3gen_ref_wav, S3GEN_SR, device=self.device)

    def generate(
        self,
        audio,
        target_voice_path=None,
    ):
        if target_voice_path:
            self.set_target_voice(target_voice_path)
        else:
            assert self.ref_dict is not None, "Please `prepare_conditionals` first or specify `target_voice_path`"

        with torch.inference_mode():
            audio_16, _ = librosa.load(audio, sr=S3_SR)
            audio_16 = torch.from_numpy(audio_16).float().to(self.device)[None, ]

            s3_tokens, _ = self.s3gen.tokenizer(audio_16)
            wav, _ = self.s3gen.inference(
                speech_tokens=s3_tokens,
                ref_dict=self.ref_dict,
            )
            wav = wav.squeeze(0).detach().cpu().numpy()
        return torch.from_numpy(wav).unsqueeze(0)

Dua Aspek Penyebab Halusinasi

Studi tersebut menguraikan penyebab halusinasi menjadi 2 tahap utama dalam paradigma pelatihan modern LLM: pra-pelatihan (pre-training) dan pasca-pelatihan (post-training).

Kesimpulan

Paper ini cukup membuka tabir halusinasi pada model bahasa AI, mulai dari asal-usulnya pada tahap pra-pelatihan dan mengapa tetap bertahan pada tahap pasca-pelatihan. Pada fase pra-pelatihan, halusinasi berkaitan erat dengan kesalahan klasifikasi apakah teks luaran LLM valid atau tidak valid. Pada fase pasca-pelatihan, halusinasi seperti “terpelihara” dengan prosedur sosio-teknikal dalam melakukan evaluasi yang memberikan penghargaan yang lebih terhadap model AI yang over-confident dibandingkan yang mengakui ketidaktahuan.

Sepertinya diperlukan kajian mendalam ke depan agar metode machine learning tidak hanya untuk membuat model AI menjadi pintar, namun juga bijak.

1. Yang terbaru pada Gemmaverse

Gemma merupakan sekumpulan model lightweight AI open-source yang dikembangkan dengan teknologi yang sama dengan Gemini namun dengan ukuran dan kapasitas yang lebih sederhana. Versi terbaru dari Gemma saat ini yaitu Gemma 3 yang terdiri dari beberapa jenis model: 1B, 4B, 12B, dan 27B, dengan peningkatan performa yang signifikan dibandingkan Gemma 2.

Terdapat pula beberapa varian dari Gemma yang di post-trained / fine-tuned pada domain tertentu yang menghasilkan special-purpose models:

Kontribusi dari komunitas terhadap Gemma juga semakin berkembang. Telah banyak aplikasi berbasis Gemma yang dikembangkan untuk berbagai use cases, baik di sektor publik maupun privat. Berbagai varian dari Gemma yang dikembangkan komunitas juga semakin banyak, seperti yang dapat ditemui di Hugging Face.

2. Small Language Models dengan Google AI Edge (keynote)

Saat ini cara yang paling mudah untuk memanfaatkan model AI adalah melalui API Cloud. Namun ada kalanya kita membutuhkan solusi AI on-device dengan ukuran model lebih kecil (small foundation models) untuk menjaga privasi data dan kebutuhan mode offline / tanpa koneksi internet.

Sesi ini menjelaskan tentang Google AI Edge, sekumpulan tech stacks yang dapat dimanfaatkan untuk merealisasikan deployment AI on-device. Framework utama yang digunakan adalah MediaPipe dan LiteRT (sebelumnya bernama TFLite). MediaPipe menyediakan low-code API untuk memudahkan pengaturan inferensi berbagai model AI off-the-shelves. LiteRT merupakan lingkungan runtime berkinerja tinggi untuk menjalankan inferensi model AI pada device tujuan secara efisien.

Google AI Edge memungkinkan pula untuk menginstalasi small language model (SLM) Gemma 3 secara on-device. Untuk mengimplementasikan Retrieval Augmented Generation (RAG) agar SLM dikondisikan dengan informasi atau konteks tertentu, disediakan juga RAG SDK sebagai bagian dari Google AI Edge. Saat ini kita juga dapat memanfaatkan Gemma 3n sebagai pilihan SLM yang lebih efektif dan teroptimalkan untuk mobile device.

Dengan terus berkembangnya pilihan dan kemampuan dari berbagai SLM beserta ekosistem dan tools pendukungnya, saat ini merupakan waktu yang tepat untuk membangun berbagai aplikasi berbasis AI on-device tanpa ketergantungan pada koneksi internet.

3. Membangun model Transformers dengan JAX (workshop)

Sesi workshop ini dibawakan oleh Wei Wei, Developer Advocate dari Google, yang menjelaskan secara terperinci implementasi arsitektur dan pelatihan model Transformer dari nol dengan menggunakan JAX dan Flax NNX. JAX merupakan sebuah framework open-source berkinerja tinggi untuk riset dan pengembangan AI, yang dikembangkan Google dengan dukungan dari NVIDIA dan komunitas. JAX didesain seperti NumPy, mendukung komputasi numerik diferensiabel dan kompilasi JIT (Just In Time) untuk akselerator seperti GPU dan TPU. Flax NNX merupakan library berisi berbagai API untuk memudahkan implementasi neural networks, mirip dengan modul-modul pada Keras dan PyTorch.

Selain Flax NNX, terdapat library lain pada ekosistem JAX untuk mendukung developer experience seperti Optax, Orbax, Grain, dan sebagainya. Untuk memudahkan instalasi, JAX AI Stack (https://github.com/jax-ml/jax-ai-stack) berisi all-in-one package yang kita butuhkan untuk melakukan pengembangan berbasis JAX.

Materi utama yang disediakan pada workshop ini berupa Google Colab Notebook, yang dapat diakses secara publik di sini. Instruktur mendemokan bagaimana melakukan pra-pelatihan terhadap model GPT-2 dengan dataset teks yang cukup sederhana, yang berjalan hanya dalam waktu beberapa menit jika menggunakan akselerator GPU atau TPU pada Colab. Dengan menggunakan Flax NNX, pembangunan layer penyusun utama seperti Transformer Block dan Multi-Head Attention dapat dilakukan lebih mudah.

Instrukturnya juga menyediakan contoh-contoh lain dalam menggunakan JAX untuk post-training dan instruction tuning GPT-2, yang dapat ditemui pada repositori Github: https://github.com/windmaple/LLM_from_scratch.JAX

4. Membangun Multi-Agent System dengan Agent Development Kit (workshop)

Agent Development Kit (ADK) merupakan sebuah kerangka kerja open-source dari Google untuk mengembangkan dan menyebarkan agentic AI, yang dirancang agar pengembangan dan pemanfaatan AI agents memiliki pengalaman seperti pengembangan perangkat lunak pada umumnya. Walaupun ADK dioptimalkan untuk menggunakan Gemini dan Google Cloud, desain kerangka kerjanya cukup model-agnostic dan deployment agnostic, yang berarti dapat menggunakan model GenAI dan infrastruktur lainnya.

Dibandingkan dengan kerangka kerja serupa seperti LangChain / LangGraph dan CrewAI, ADK relatif baru namun perkembangannya cukup pesat. Salah satu kekuatan dari ADK yaitu kemudahan implementasi arsitektur multi-agent systems (MAS). Dengan prinsip MAS, pengembang berfokus pada perancangan special-purpose, simple AI agent secara modular, dimana solusi dari suatu permasalahan dikerjakan oleh beberapa agen yang saling bekerjasama, dengan instruksi/prompt yang sederhana untuk tiap-tiap agen. Hal ini lebih menjamin keandalan sistem dan juga kemudahan pengelolaan dibandingkan menggunakan satu agen dengan prompt yang panjang dan kompleks.

Dengan menggunakan ADK, kita dapat merancang MAS dengan struktur Hierarchical Agent Tree, yang memudahkan untuk mengontrol dan memprediksi rute penalaran atau percakapan antar agen pada suatu pohon pengambilan keputusan. Struktur pohon hirarkis ini terinspirasi dari kolaborasi tim pada dunia nyata dan juga selaras dengan berbagai design pattern pada pengembangan perangkat lunak konvensional.

Pada workshop ini instruktur membagikan kelas Lab untuk membangun Multi-Agent System dengan ADK yang dapat diakses secara publik. Bagi yang ingin mencoba sendiri, bisa langsung mengunjungi cloudskillsboost.google.

5. Living Canvas (tech demo)

Di luar ruangan ballroom maupun kelas, tersedia area technical demo dimana beberapa stand siap dikunjungi untuk menyaksikan project showcase secara interaktif. Salah satu project yang menarik buat saya yaitu Living Canvas, sebuah aplikasi game puzzle berbasis web dimana pemain langsung menggambar sketsa pada kanvas, lalu mengkonversi sketsa tersebut menjadi objek yang relevan untuk menyelesaikan sebuah puzzle (mendekatkan kunci ke anak kunci).

Pembangunan aplikasi web ini memanfaatkan beberapa teknologi yaitu:

• Angular: Membuat antarmuka keseluruhan aplikasi
• Firebase App Hosting: Serverless web hosting untuk otomasi proses deployment.
• Vertex AI: Pintu masuk untuk mengakses Gemini dan Imagen melalui berbagai API yang tersedia.
• Gemini: Melakukan analisis terhadap gambar input sketsa dari pengguna dan memetakan ke tipe objek yang sudah terdefinisi.
• Imagen: Mengkonversi gambar input menjadi gambar dengan grafis berkualitas tinggi berdasarkan tipe objek yang dipetakan oleh Gemini.
• Veo: Membuat animasi pada gambar yang dihasilkan oleh Imagen.

Bagi yang ingin mencoba menjalankan langsung, dapat langsung mengakses ke repositori kode https://github.com/FirebaseExtended/solution-living-canvas.

Selain mendapatkan ilmu baru dari materi-materi yang disampaikan, dapat bertemu dan berinteraksi dengan Googlers, GDE, dan juga para peserta dari berbagai negara merupakan pengalaman yang tak kalah berharga. Terima kasih banyak kepada Googlers, terutama para Community Managers, dan juga penyelenggara yang telah memfasilitasi para GDE untuk bisa menghadiri langsung Google I/O Connect Shanghai 2025.

Novel View Synthesis with NeRF

Before delving into the detailed implementation, let’s formulate the novel view synthesis problem and how NeRF tries to solve it. Novel view synthesis is a task which consists of generating images of a specific scene from a specific point of view, when the only available information is pictures taken from different points of view.

Source: https://paperswithcode.com/task/novel-view-synthesis

Initially, we have a set of source images and their corresponding camera poses . The goal is to predict the novel view image from a new target pose .

Formally, novel view synthesis attempts to construct a function leveraging a dataset of image views such that

where are pixel coordinates and the space of represents a scene. Early methods factorize this into:

• 3D geometry estimation (e.g., depth or point clouds)
• View rendering (warping textures onto geometry)
• Hole-filling (inpainting unseen regions)

NeRF frames this problem into modeling a scene as a 5D continuous function , which a trainable neural network with parameter .

The 5D input space comprises a point location randomly sampled on the camera ray (which is basically a straight line), and the ray’s viewing angle . In practice, the 2D viewing angle can be converted into a 3D unit directional vector .

The 2D output space contains the volume density and the RGB color .

In other words, the neural network mapping looks like the following:

Through differentiable volume rendering (will be explained later), the model is optimized to match observed pixel colors—unifying geometry, view-dependent effects, and rendering in one trainable neural net.

By doing so, NeRF can yield:

• High photographic quality
• Realistic specular and lighting variations
• No explicit geometry required

Camera Ray Data as Inputs

One may ask where we can the data representing the inputs and . It all starts with the camera parameters (poses, intrinsics, and bounds) — those can be obtained by using COLMAP structure-from-motion (SfM) on real-world image samples or generated through 3D modeling software if using synthetic scenes.

from keras import ops

def get_rays(height, width, focal, pose):
    u, v = ops.meshgrid(
        ops.arange(width, dtype="float32"),
        ops.arange(height, dtype="float32"),
        indexing="xy",
    )
    transformed_u = (u - width * 0.5) / focal
    transformed_v = (v - height * 0.5) / focal
    directions = ops.stack(
        [transformed_u, -transformed_v, -ops.ones_like(u)], axis=-1
    )
    camera_matrix = pose[:3, :3]
    translations = pose[:3, -1]
    transformed_dirs = directions[..., None, :]
    camera_dirs = transformed_dirs * camera_matrix
    ray_directions = ops.sum(camera_dirs, axis=-1)
    ray_origins = ops.broadcast_to(translations, ops.shape(ray_directions))
    return (ray_origins, ray_directions)

                Camera
           ●
           |
           |   t1      t2      t3       tN
           |----●-------●-------●--------●----→ ray direction
               near                        far

import keras

def generate_t_vals(near, far, batch_size, num_samples, rand_sampling=True):
    t_vals = ops.linspace(near, far, num_samples, dtype="float32")
    if rand_sampling:
        noise = keras.random.uniform(shape=ops.shape(t_vals)) * (far - near) / num_samples
        t_vals = t_vals + noise
    
    t_vals = ops.broadcast_to(t_vals, (batch_size, num_samples))

    return t_vals

... 
t_vals = generate_t_vals(near, far, batch_size, num_samples, rand_sampling=True)
...

def sample_rays(ray_origins, ray_directions, t_vals):
    rays = ray_origins[..., None, :] + (
        ray_directions[..., None, :] * t_vals[..., :, None]
    )
    dir_shape = ops.shape(rays[..., :3])
    dirs = ops.broadcast_to(ray_directions[..., None, :], dir_shape)
    return rays, dirs

def encode_position(x, pos_encode_dims):
    positions = [x]
    for i in range(pos_encode_dims):
        for fn in [ops.sin, ops.cos]:
            positions.append(fn(2.0**i * x))
    return ops.concatenate(positions, axis=-1)

Network Architecture and Training

The NeRF network is a deep multilayer perceptron (MLP) featuring skip connections to enhance the gradient flow and preserve spatial detail deep into the network. By incorporating position embeddings, the network function becomes something like this:

def create_nerf_complete_model(num_layers, hidden_dim, skip_layer, lxyz, ldir, bn=False):
    ray_input = keras.Input(shape=(None, 2 * 3 * lxyz + 3))
    dir_input = keras.Input(shape=(None, 2 * 3 * ldir + 3))

    x = ray_input
    for i in range(num_layers):
        if bn:
            x = layers.Dense(hidden_dim)(x)
            x = layers.BatchNormalization()(x)
            x = layers.ReLU()(x)
        else:
            x = layers.Dense(hidden_dim, activation="relu")(x)

        # Check if we have to include residual connections
        if i % skip_layer == 0 and i > 0:
            x = layers.concatenate([x, ray_input], axis=-1)
        
    # Get the sigma value
    sigma = layers.Dense(1)(x)

    # Create a feature vector
    feature = layers.Dense(hidden_dim)(x)

    # Concatenate the feature vector with the direction input
    feature = layers.concatenate([feature, dir_input], axis=-1)
    if bn:
        x = layers.Dense(hidden_dim//2)(feature)
        x = layers.BatchNormalization()(x)
        x = layers.ReLU()(x)
    else:
        x = layers.Dense(hidden_dim//2, activation="relu")(feature)

    # Get the rgb value
    rgb = layers.Dense(3)(x)

    outputs = layers.concatenate([rgb, sigma], axis=-1)

    nerf_model = keras.Model(inputs=[ray_input, dir_input], outputs=outputs)
    return nerf_model

t_vals = generate_t_vals(near, far, batch_size, num_samples, rand_sampling=True) # generate depth parameters t
rays, dirs = sample_rays(ray_origins, ray_directions, t_vals) # generate ray samples r [batch_size, N * (6 * L + 3)] and directions d [batch_size, N * (6 * L + 3)]

model = create_nerf_complete_model(num_layers, hidden_dim, skip_layer, lxyz, ldir) # define NeRF network

(rgbs, sigmas) = model([rays, dirs]) # forward pass

def volume_render(preds, t_vals):
    # Get rgb and sigma from the predictions
    rgb = ops.sigmoid(preds[..., :-1])
    sigma_a = ops.relu(preds[..., -1])

    # Get the distance of adjacent intervals
    delta = t_vals[..., 1:] - t_vals[..., :-1]
    const = ops.broadcast_to([1e10], shape=(delta.shape[0], 1))
    delta = ops.concatenate([delta, const], axis=-1)

    alpha = 1.0 - ops.exp(-sigma_a * delta)
    exp_term = 1.0 - alpha
    epsilon = 1e-10

    # Compute transmittance: Cumulative prod with exclusive mode
    tm = ops.cumprod(exp_term + epsilon, axis=-1)
    tm = ops.roll(tm, shift=1, axis=-1)
    transmittance = ops.concatenate([ops.ones((tm.shape[0], 1)), tm[:, 1:]], axis=-1)

    # Compute weights
    weights = alpha * transmittance
    rgb_w = ops.sum(weights[..., None] * rgb, axis=-2)    
    depth_map = ops.sum(weights * t_vals, axis=-1)
    return (rgb_w, depth_map, weights)

# Define coarse and fine models
...
coarse_model = create_nerf_complete_model(
    num_layers=NUM_LAYERS,
    hidden_dim=HIDDEN_DIM,
    skip_layer=SKIP_LAYER,
    lxyz=L_XYZ,
    ldir=L_DIR,
    bn=BATCH_NORM
)
fine_model = create_nerf_complete_model(
    num_layers=NUM_LAYERS,
    hidden_dim=HIDDEN_DIM,
    skip_layer=SKIP_LAYER,
    lxyz=L_XYZ,
    ldir=L_DIR,
    bn=BATCH_NORM
)

...

# Coarse model forward pass 
t_vals = generate_t_vals(...)
rays, dirs = sample_rays(ray_origins, ray_directions, t_vals)
rays_enc = encode_position(rays, pos_encode_dims=l_xyz)
dirs_enc = encode_position(dirs, pos_encode_dims=l_dir)
predictions_coarse = coarse_model([rays_enc, dirs_enc], training=training)
rgb_coarse, depth_coarse, weights_coarse = volume_render(predictions_coarse, t_vals)

# Sample PDF
t_vals_coarse_mid = (0.5 * (t_vals[..., 1:] + t_vals[..., :-1]))
t_vals_fine = sample_pdf(t_vals_coarse_mid, weights_coarse, self.ns_fine)
t_vals_fine_all = ops.sort(ops.concatenate([t_vals, t_vals_fine], axis=-1), axis=-1)

# Fine model forward pass
rays_fine, dirs_fine = sample_rays(ray_origins, ray_directions, t_vals_fine_all)
rays_fine_enc = encode_position(rays_fine, pos_encode_dims=l_xyz)
dirs_fine_enc = encode_position(dirs_fine, pos_encode_dims=l_dir)
predictions_fine = fine_model([rays_fine_enc, dirs_fine_enc], training=training)

def sample_pdf(t_vals_mid, weights, ns_fine):
    # Get batch_size, H, W
    batch_size = ops.shape(weights)[0]
    if len(ops.shape(weights)) == 4: # (b, h, w, num_samples)
        image_height, image_width = ops.shape(weights)[1:3]
    
    # add a small value to the weights to prevent it from nan
    weights += 1e-5
    
    # normalize the weights to get the pdf
    pdf = weights / tf.reduce_sum(weights, axis=-1, keepdims=True)
    
    # from pdf to cdf transformation
    cdf = tf.cumsum(pdf, axis=-1)
    
    # start the cdf with 0sa
    cdf = tf.concat([tf.zeros_like(cdf[..., :1]), cdf], axis=-1)

    # get the sample points
    if len(ops.shape(weights)) == 4:
        u_shape = [batch_size, image_height, image_width, ns_fine]
    else:
        u_shape = [batch_size, ns_fine]

    u = tf.random.uniform(shape=u_shape)
    
    # get the indices of the points of u when u is inserted into cdf in a
    # sorted manner
    indices = tf.searchsorted(cdf, u, side="right")

    # define the boundaries
    below = tf.maximum(0, indices-1)
    above = tf.minimum(cdf.shape[-1]-1, indices)
    indices_g = tf.stack([below, above], axis=-1)
    
    # gather the cdf according to the indices
    cdf_g = tf.gather(cdf, indices_g, axis=-1, batch_dims=len(indices_g.shape)-2)

    # gather the tVals according to the indices
    indices_gt = tf.minimum(t_vals_mid.shape[-1] - 1, indices_g)
    t_vals_mid_g = tf.gather(t_vals_mid, indices_gt, axis=-1,
        batch_dims=len(indices_g.shape)-2)
    
    # create the samples by inverting the cdf
    denom = cdf_g[..., 1] - cdf_g[..., 0]
    denom = tf.where(denom < 1e-5, tf.ones_like(denom), denom)
    t = (u - cdf_g[..., 0]) / denom
    samples = (t_vals_mid_g[..., 0] + t * 
        (t_vals_mid_g[..., 1] - t_vals_mid_g[..., 0]))

    # return the samples
    return samples

def train_step(self, inputs):
    # Get the image and the rays
    (images, rays) = inputs
    (ray_origins, ray_directions, t_vals) = rays

    with tf.GradientTape() as tape:
        # Get the predictions from the model
        rgbs, _, _, _ = self.forward_pass(ray_origins, ray_directions, t_vals, self.l_xyz, self.l_dir, training=True)
        rgb_coarse, rgb_fine = rgbs
        loss_coarse = self.loss_fn(images, rgb_coarse)
        loss_fine = self.loss_fn(images, rgb_fine)
            
        # Combine the coarse and fine losses
    loss = loss_coarse + loss_fine

        # Apply gradient updates for the model
        tv_nerf = self.coarse_model.trainable_variables + self.fine_model.trainable_variables
        grads = tape.gradient(loss, tv_nerf)
        self.optimizer.apply_gradients(zip(grads, tv_nerf))

        # Get the PSNR of the reconstructed images and the source images
        psnr = ops.psnr(images, rgb_fine, max_val=1.0)

        # Compute the metrics
        self.loss_coarse_tracker.update_state(loss_coarse)
        self.loss_tracker.update_state(loss_fine)
        self.psnr_tracker.update_state(psnr)
        return {
            "loss_coarse": self.loss_coarse_tracker.result(),
            "loss": self.loss_tracker.result(),
            "psnr": self.psnr_tracker.result(),
        }

Experiments on a Synthetic Lego Scene

We conducted NeRF training on the synthetic Lego scene dataset, which consists of 106 images at a resolution of 100 x 100 pixels, each accompanied by its corresponding camera pose. The dataset was split into training and validation sets using an 80:20 ratio, resulting in 84 training samples and 22 validation samples. To further reduce the computational complexity, all images were downsamples to a resolution of (H = 50, W = 50) pixels.

Next, we executed the data processing pipeline described earlier to obtain the the positionally encoded coarse ray samples. By setting the number of coarse samples to and using positional embedding with (hence, the dimensionality of each ray becomes ), the shape of the full training dataset is:

Notes on batch training

Since NeRF operates at the pixel level, the effective number of training samples is not just the number of images #train = 84 , but rather the total number of rays in every pixel, which equals to #train x H x W = 84 x 50 x 50 = 210000. To enable efficient batch processing, the full training set is reshaped to

This reshaping step is essential for enabling mini-batch training — without it, feeding the full 5D tensor into the model would quickly exhaust GPU memory!

By setting the training epoch = 300 and batch_size = 1024 , we get the following results:

Rendering results during training, displayed per epoch

Rendering results during training, displayed per epoch

Comparison between the ground truth and reconstructed images

Conclusion

This article walked you through the complete implementation of a basic NeRF, demystifying every major component: problem setup, ray generation and sampling, positional embeddings, network architecture, volume rendering, hierarchical sampling, and training setup.

These days, NeRF has evolved remarkably and several exciting developments are pushing the field forward for better generalization, faster training, and real-time rendering, such as:

• (Cong et al. ICCV2023): Generalization and Volume Rendering advancements, such as transformer-based architectures for cross-scene NeRF, enable few-shot synthesis of new content and improved boilerplate-free deployments.
• (Wang et al. WACV2024): Hyb‑NeRF uses multiresolution hybrid encoding (e.g., hash grids) to drastically speed up inference while maintaining visual quality.
• (Ding et al. 2025): Neural Pruning methods apply structured parameter reduction during training, achieving up to 50% model size reduction and 35% faster training, with minimal accuracy loss.
• (Qazi et al. 2025): NeuGen introduces brain-inspired normalization to improve domain generalization, enabling NeRF-like models to perform robustly across diverse environments.

See also the curated list of papers: https://github.com/awesome-NeRF/awesome-NeRF.

Introduction

Augmented reality (AR), virtual reality (VR), and mixed/extended reality (XR) technologies have advanced rapidly in the past year, driven by major tech firms and innovative startups worldwide. The global XR market (encompassing AR and VR) was valued around $25.7 billion in 2024 , and hardware adoption is accelerating – for example, AR/VR headset shipments grew ~12.8% year-on-year in Q3 2024, with Meta’s Quest devices alone accounting for about 70% of units thanks to new releases . These “spatial computing” technologies are increasingly moving from hype to practical deployment across industries. This report details the latest state-of-the-art hardware and software, key company developments, and XR applications in gaming, healthcare, education, enterprise, and retail on a global scale.

Hardware Advancements in AR, VR, and XR

High-End Mixed Reality Headsets:  A milestone in 2024 was Apple’s entry into XR with the Vision Pro headset (available starting Feb 2024). Apple Vision Pro is a “spatial computer” that seamlessly blends digital content with the physical world, enabling powerful mixed-reality experiences in a user’s environment . It features advanced sensors (cameras and LiDAR), dual 4K micro-OLED displays, an Apple M2 chip plus a dedicated R1 chip for real-time sensor processing, and an innovative transparent lens design. The Vision Pro is controlled via natural inputs – eye tracking, hand gestures, and voice – with no controllers . Priced at$3,499, it targets developers and prosumer users initially. Apple’s visionOS platform allows Vision Pro to run thousands of iPad and iPhone apps plus new spatial apps via an all-new app store . Apple has been steadily improving visionOS (e.g. a September 2024 update and a March 2025 update adding AI-powered “Apple Intelligence” features) and expanding content, such as immersive 3D videos and concerts for Vision Pro . This marks a significant step forward in XR hardware, with TechCrunch calling it “the most advanced consumer electronics device ever created”, and it has spurred competition across the industry.

Meta’s VR/MR Headsets: Meta (Facebook) has solidified its lead in consumer VR with the Meta Quest 3, launched in October 2023. The Quest 3 is billed as “the world’s first mass-market mixed reality headset” . At $499, it delivers significant improvements over its predecessor: a 30% leap in visual resolution, 2× GPU performance (Snapdragon XR2 Gen 2 chip), 40% louder audio, and a slimmer, more comfortable design . Notably, Quest 3 introduces high-fidelity full-color passthrough cameras enabling mixed reality – users can see their surroundings in color and overlay virtual objects with 10× the pixel count of Quest 2’s passthrough . With a quick double-tap on the side, wearers can seamlessly toggle between fully immersive VR and mixed reality, blending virtual elements into their real space . This allows experiences like playing a virtual piano on your coffee table or opening a portal in your living room, as Meta showcases . The Quest 3, powered by inside-out tracking and self-contained, represents the mainstream trend toward standalone XR headsets that do both VR and passthrough AR. Meta also offers the higher-end Quest Pro (launched late 2022) oriented at enterprise use with features like eye and face tracking, but the Quest 3 has largely overtaken it for consumers. In addition, Meta introduced updated Ray-Ban Meta smart glasses in late 2023 – while these are camera-and-audio wearables (no display AR yet), Meta’s roadmap includes developing true AR glasses later in the decade.

Advanced VR for Enterprise: Outside the mass market, specialized manufacturers have pushed VR/MR hardware to extremes of fidelity. Notably, Finland’s Varjo released its next-generation Varjo XR-4 series headsets in late 2023, aimed at enterprise and government use. The XR-4 features the highest-resolution displays in any XR device (over 28 million pixels) combined with custom aspheric lenses, foveated eye-tracking, and LiDAR depth sensing . According to Varjo, the result is virtual and mixed reality visuals “practically indistinguishable from natural sight” – a huge claim enabled by human-eye resolution in the center of view and real-time depth scanning. The XR-4 (starting around$3,990) is not for gamers; it’s positioned to “go far beyond consumer-grade headsets” and serve industrial training, simulation, and design where top fidelity is required . In fact, over 25% of Fortune 100 companies use Varjo’s technology for use cases like training pilots and astronauts, speeding up automotive design, and powering medical breakthroughs . Varjo’s systems even achieved regulatory milestones – in 2024, a VR-based pilot training program using Varjo headsets became the first to earn FAA approval toward actual pilot certification . This underscores how professional-grade XR hardware is transforming workforce training. Other startups are also innovating in high-performance XR: for example, Bigscreen Beyond, a VR startup, began shipping an ultra-compact PC VR headset weighing just ~127g (much lighter than typical 500+g headsets) to improve comfort, and companies like Pimax (China) offer wide-FOV, high-resolution VR headsets for enthusiasts. While these remain niche, they showcase state-of-the-art advances in optics, displays, and ergonomics that will trickle down to consumer devices over time.

AR Glasses and Wearables: In parallel, there’s intense development of augmented reality glasses that are lightweight and all-day wearable. Many efforts in the past year focused on making AR glasses more practical for consumers and enterprises. For instance, California-based startup Xreal (formerly Nreal) introduced the Xreal Air 2 series of glasses in 2024 – lightweight smart glasses that connect to phones/laptops. The new Xreal Air 2 Ultra model adds onboard cameras and processors, enabling true spatial computing experiences comparable to pricier headsets from Apple or Meta, but in a 60 million for further development . This indicates a growing consumer appetite for affordable AR eyewear that leverages existing devices. Another notable player is Magic Leap, which shifted fully to enterprise AR with its Magic Leap 2 headset (released 2022 but rolled out broadly through 2023–24). Magic Leap 2 features one of the most advanced optic systems in see-through AR, with a wide field of view, high clarity, and robust occlusion of dark objects. It targets use cases like surgical visualization, remote expert support, and 3D design. Magic Leap made headlines in mid-2024 by partnering with Google to combine its optics and AR engineering expertise with Google’s software platforms . This strategic alliance aims to “advance the future of the XR ecosystem with unique and innovative product offerings” – possibly hinting that the next Magic Leap devices will integrate Google’s operating systems or AI capabilities (more on Google’s XR platform below). Beyond these, many other AR wearables emerged: Lenovo’s ThinkReality division, Vuzix, Rokid (China), and others released or announced new smart glasses for enterprise (typically offering heads-up info display or video calls for field workers). While true consumer AR glasses with full 3D graphics remain a challenge (due to optics, battery, and privacy concerns), the past year saw clear progress toward smaller, more powerful AR wearables that inch closer to everyday glasses.

Supportive Hardware and Haptics: XR hardware advances are not limited to headsets. There have been improvements in controllers and input devices – for example, Meta’s Quest 3 controllers feature improved haptic feedback and inside-out tracking (no external sensors needed), and some VR systems (HTC VIVE, Valve) support hand tracking so users can interact without controllers. In haptics, startups like HaptX and Manus are refining haptic gloves that provide force feedback for realistic touch in VR, which saw continued R&D and pilot uses (e.g. in VR training for feeling virtual objects). Additionally, full-body tracking and suits (such as Teslasuit or Sony’s Mocopi trackers) are being used for motion capture and more immersive experiences, though still niche. In spatial computing, there’s also focus on better spatial audio hardware – new headsets (Apple, Meta) incorporate advanced spatial speakers so virtual sounds come from proper directions, enhancing realism. All combined, the hardware ecosystem in 2024–2025 offers a spectrum from cutting-edge high-end XR systems to smaller, more accessible AR/VR devices, setting the stage for broader adoption.

Software and Platform Innovations

Hardware gains have been matched by major software advancements in AR/VR platforms and tools:

• Operating Systems and Toolkits: A significant development came from Google, which in December 2024 announced Android XR, a new open platform built specifically for XR devices . Developed in partnership with Samsung and Qualcomm, Android XR is essentially an XR-optimized version of Android that will power upcoming headsets and smart glasses. Google is providing a developer preview (with support for familiar tools like ARCore, Android Studio, Unity, OpenXR, etc.) so developers can start building apps for upcoming Android XR devices . The first device, codenamed “Project Moohan” by Samsung, is slated for release in 2025 . Android XR also integrates Google’s latest AI assistant (Project Gemini) to enable voice and vision intelligence – e.g. a user can ask the glasses about what they’re seeing, get contextual info, translations, or have the AI help with tasks . This shows a trend of convergence between XR and AI, making software more intuitive. Apple, similarly, has been evolving its visionOS for Vision Pro, adding features like persona avatars and incorporating on-device AI (referred to as “Apple Intelligence” in a 2025 update) for smarter interactions.
• Content and Developer Ecosystems: The XR content ecosystem expanded significantly. Both Unity and Unreal Engine – the two leading 3D engines – released updates to better support AR/VR development. For example, Unity’s 2023 and 2024 versions include built-in support for XR plugins, OpenXR standard, and even specific support for Apple Vision Pro’s unique features . Epic’s Unreal Engine 5 continues to be used for high-fidelity VR experiences and has templates for AR projects. Cross-platform frameworks like OpenXR have gained traction, allowing developers to write an app once and deploy to multiple VR/AR devices. This was underscored by companies like Varjo achieving OpenXR 1.0 compliance for their high-end headsets , and by Microsoft and Meta also aligning to OpenXR on their platforms.
• XR Cloud Services: The concept of the “AR cloud” (a persistent 3D map of the world for shared AR experiences) saw progress. Niantic’s Lightship VPS (Visual Positioning System) expanded to more cities, allowing phones and headsets to localize in specific real-world locations with cm-level accuracy for AR content. Niantic also updated its Lightship AR SDK (e.g. ARDK 3.0) to let developers blend Unity AR experiences with its real-world scanning and multiplayer features . Other players like Snap and 8th Wall (web AR platform) similarly grew their AR cloud offerings, enabling things like location-based AR advertising and city-scale AR games. Microsoft, while stepping back from its own VR hardware, continues to offer Azure Mixed Reality services (like Azure Object Anchors and Remote Rendering) to enterprise developers building cloud-connected AR applications (e.g. visualizing big 3D models via HoloLens or other devices by streaming from the cloud). We also saw more adoption of 3D content standards: USD (Universal Scene Description) is being championed by NVIDIA and others for metaverse interoperability, and glTF for 3D models on the web, helping standardize XR content pipelines.
• Social and Collaborative Platforms: In the past year, there have been upgrades to platforms enabling multi-user XR experiences. Meta’s Horizon Worlds social VR, while off to a slow start, got improvements and opened to web and mobile users to expand its reach (so people can join VR environments from a phone or PC). Microsoft’s Mesh for Teams – which blends mixed-reality avatars and meetings – was showcased and is gradually rolling out, allowing workers to meet as avatars in a shared virtual space via VR or computer. Startups like Spatial and Frame provided browser-based collaborative VR spaces, used for everything from virtual offices to conferences. The enterprise collaboration space is vibrant, with solutions for virtual training sessions, design reviews in VR, and remote assistance via AR (e.g. using a HoloLens or phone to get a remote expert’s guidance overlayed on machinery). The ongoing challenge has been interoperability, but initiatives through the Metaverse Standards Forum (launched 2022) continued through 2024 to align formats and protocols so virtual worlds can connect.
• Creation and AI Integration: Another notable trend is AI-powered content creation for XR. Generative AI is being leveraged to create 3D assets, virtual characters, and even entire virtual environments more efficiently. Companies like NVIDIA introduced tools to generate 3D models from 2D images or to assist in scene creation for simulation. Startups (e.g. Singularity 6 or Inworld AI) are making AI-driven NPCs that can interact naturally with users in VR. Snap’s Lens Studio now has AI features that let creators generate AR Lens effects or even create AR filters via text prompts. The Fast Company “most innovative companies” list for 2025 highlights that Snap’s platform now lets millions of users easily create their own AR Lenses with AI tools . This convergence of XR and AI lowers the barrier to content creation and will likely yield more dynamic, personalized XR experiences moving forward.