Text-Conditioned IMLE

At APEX Lab (with Ke Li), I worked on extending IMLE (Implicit Maximum Likelihood Estimation) from unconditional image synthesis into the text-conditional setting, turning Adaptive IMLE into a testbed for one question: how should a text prompt steer a generator, and why does conditioning so often quietly kill sample diversity? Across more than a hundred runs on Oxford 102 Flowers, CelebA, and ImageNet, the answer kept pointing back at the same tension.

Why IMLE, and why in the small-data regime

GANs maximise a likelihood through an adversarial game, which makes them prone to mode collapse, especially when data is scarce. Diffusion models need a lot of data and sampling steps to behave. IMLE flips the objective around: instead of pushing generated samples toward the data, it pulls every real sample toward a nearby generated one, which directly optimises for mode coverage:

That difference matters most when training data is limited, which is exactly where the Adaptive IMLE line of work this project builds on shows its edge: GANs and diffusion degrade into artefacts and broken faces, while Adaptive IMLE stays sharp and diverse.

In the conditional version, the generator becomes $G(z,c)$ and we model $p(x\mid c)$, where the condition $c$ is a text embedding (a CLIP encoding of a caption). The objective keeps its shape but now matches within each condition:

Making the matching cheap

Every IMLE step solves, for each real sample $x_i$, a nearest-neighbour search over the generated pool:

This is naively $O(NM)$ in the number of real samples $N$ and latent draws $M$, and it dominates training cost. Before any conditioning experiments were possible, the matching had to be fast, so the first block of work was infrastructure:

• FAISS nearest-neighbour search replacing the original MDCI routine, giving GPU-accelerated approximate search that scales with batch size, latent count, and embedding dimension.
• A fused distance computation that expands the squared Euclidean distance into a single matrix multiply, the form GPUs and FAISS handle most efficiently:

• A revised latent resampling / assignment schedule (the force_resample cadence): how often to recompute matches versus reuse cached ones, trading raw compute against fresh assignments and training stability. Resampling every 2 epochs turned out to be a good default, slightly improving FID and precision at a small recall cost.

The matching is also non-differentiable, so IMLE deliberately separates the assignment step from the generator-optimisation step. The whole investigation rides on this separation working.

How should a prompt steer the generator?

With matching cheap, the core question became architectural: where and how do you inject the text embedding so the prompt actually controls the output? I built each variant as its own commit (the history reads like a changelog) and benchmarked them on Oxford 102 Flowers at a reduced latent count $k{=}9$ for speed. Metrics are FID (lower is better) and the improved precision / recall pair, where precision tracks fidelity and recall tracks diversity, all read at a matched training budget.

Conditioning architectures on Oxford 102 Flowers (k = 9), at a matched training budget. FiLM leads; conditional-StyleGAN noise injection collapses entirely.

Two results stood out. FiLM, where the text embedding modulates per-channel affine transforms, matched or beat every other injection scheme: the prompt gets to gate features rather than fight the latent for the same channels. At the opposite end, the conditional-StyleGAN style of injecting noisy text into the mapping network collapsed completely (FID near 275, recall zero), either fundamentally mismatched here or needing far more training than the budget allowed.

The unconditional reference is the quiet lesson in this table: its FID is much worse, yet its recall is healthy. Conditioning bought fidelity but spent diversity, and chasing that lost diversity back became the rest of the project.

Shaping the conditioning signal

Beyond architecture, I swept the signal itself: the auxiliary losses and the representation of the text embedding.

CLIP loss has to be gentle, and never alone. Adding a CLIP alignment loss between the prompt and the generated image helps only at a small weight, and only when paired with an L2 term; turn it up and quality craters.

Effect of the CLIP alignment-loss weight (Flowers, concatenation with per-block injection). Only a light CLIP term helps, and only alongside an L2 loss.

The text representation sets a floor on quality. Quantising captions into k-means codes makes the granularity tradeoff explicit. With only 8 codes the generator has too little to latch onto and never converges; push to 50 and then 100 codes and both fidelity and coverage recover, though pushing quantisation very fine eventually risks dropping the rarest flower types. Unit-normalising the CLIP embeddings, and randomly projecting them to low dimensions (64, 32, 5), both reliably hurt diversity too.

Quantising captions into k-means codes (Flowers, FiLM). Too few codes fail to converge; richer codebooks restore both fidelity and coverage.

Supporting machinery underneath all of this: a text sampler, CLIP and L2-CLIP losses, k-means clustering of captions, text unit normalisation, and dataset normalisation statistics, with W&B image and metric logging throughout, and early exploration on ImageNet (32×32) alongside the main Flowers and CelebA runs.

The diversity problem, and what actually moved it

The same pattern followed me from Flowers to CelebA: conditional models produced clean faces but a narrowing range of them. The clearest way to see it is to fix a prompt and draw several latents. A plain conditional model returns near-duplicates; the prompt has effectively pinned the output. The single most effective fix was not a loss or an injection trick, it was initialisation, starting the conditional generator from a strong unconditional IMLE prior and letting the text branch adapt on top of it.

Classifier-free guidance helped too, but proved brittle to its conditioning-dropout probability $p$.

Restoring diversity on CelebA. Initialising from an unconditional prior lifts recall the most; classifier-free guidance helps but is brittle to the dropout rate p.

Note the tension in the first rows: the base model has the best FID and precision but the worst within-prompt variety. Initialising from the unconditional prior trades a few FID points to lift recall from the mid-0.6s into the low-0.8s, the clearest single intervention in the whole study. CFG at $p{=}0.3$ lands in a similar place, but dropping $p$ to $0.1$ collapses the model entirely (precision 0.05): the conditioning-dropout rate is a cliff, not a slider. Smaller architectures, frozen backbones, and random horizontal flipping all underperformed.

What I took away

Every run, metric curve, and qualitative grid is public on Weights & Biases (organised into per-experiment reports), and the implementation, committed feature by feature, is on GitHub.