If you have ever tried to train a prostate lesion model on more than one public dataset, you know the problem. Every cohort was scanned differently and labeled by different rules, and the lesions themselves are small and low contrast on biparametric MRI (bpMRI). In arXiv:2611.08535, Thibault Heintz and colleagues at Weill Cornell Medicine and Cornell Tech pool five of those cohorts into one nnU-Net, add a deliberately destructive augmentation based on random wavelet filters, and release both the trainer code and the trained checkpoint, so you can download the model and run it on your own studies.

What the model does

You give it three co-registered sequences of a prostate bpMRI study: T2-weighted, the ADC map and the high b-value DWI, with ADC and DWI resampled onto the T2 grid. It returns a 3D mask with a lesion label. During training it also learns a prostate gland label, a trick called BackSplit (co-developed by one of the authors) that gives the network extra signal about the organ around a small lesion. At inference only the lesion label is used.

For a viewer or a reporting tool, the mask stands in for the contour a radiologist draws by hand today before an MRI-ultrasound fusion biopsy or a structured report. The paper cites only moderate agreement between readers for those contours, so a model-drafted first contour has a clear place in that workflow.

Five cohorts, five labeling rules

The training pool has 1,579 lesion-containing cases, which the authors describe as the largest collection used for this task. It combines a derived PROMIS set, PI-CAI (with the extra annotations by Pooch et al.), Prostate158, the UCLA Prostate-MRI-US-Biopsy collection (PMUB) and an internal Weill Cornell cohort. The cohorts differ in three ways that matter for training:

• Scanners. PROMIS is entirely 1.5T, PI-CAI mixes 1.5T and 3T, the rest are 3T. High b-values range from 1000 to 1400 s/mm².
• What counts as a lesion. Prostate158 annotates PI-RADS 4 and above, the internal cohort generally PI-RADS 3 and above, PMUB only the biopsy targets, and PI-CAI lesions of ISUP grade 2 or higher.
• Lesion size. Median lesion volume runs from 0.182 mL in the internal data to 1.121 mL in PROMIS, mostly a consequence of those labeling rules.

Destructive augmentation with random wavelets

With probability 0.5 per training sample, on top of the standard nnU-Net augmentations, each sequence is replaced by the response of a randomly built wavelet filter. The filter decomposes the volume at four scales, mixes the detail bands of the three axes along a random direction, sums the scales with random weights and rescales the result to the original intensity range. A new filter is drawn for each sequence.

The result looks nothing like an MRI. The featured image of this post shows Fig. 1 of the paper: the top row is a PI-CAI case as acquired, the bottom row is the same case after the sym4 filter. Textures are scrambled and edges are displaced. The authors argue that this helps because the labels are fuzzy and the cohorts differ, so the network gains little from fitting exact image-to-mask alignment. They compare four wavelet families (Haar, db2, db4, sym4) against random convolutions and nnU-Net’s stronger DA5 preset, and sym4 comes out best on average Dice. The gain is largest on PROMIS, the only cohort scanned entirely at 1.5T and a minority of the training data.

The second change is the loss. Plain Dice plus cross-entropy is computed over all lesion voxels at once, so one large lesion can dominate a case with several small ones. They add a connected-component loss (CC-Dice-CE) that splits the volume into Voronoi cells around each reference lesion and scores each lesion in its own cell, and they show it works together with BackSplit.

How well it works

On the official Prostate158 test split, against Reader 1 labels, the final model reaches a Dice of 0.614, compared with the 0.450 baseline reported by the dataset authors. Across all five test sets, the full recipe raises mean Dice from 0.391 for a plain nnU-Net ResEnc-M to 0.422.

Both the wavelet augmentation and the CC loss raise lesion-level recall and lower precision. The final model’s lesion F1 across all test sets is 0.571, below the 0.596 of the BackSplit-only model, though still above plain nnU-Net. The authors say they prefer recall because a missed lesion costs more than a false positive. In a viewer, that means more candidate blobs for the reader to dismiss.

Where it can fail

• No negative cases. Cases without an annotated lesion were excluded from training and testing, so there is no evidence on how often the model marks something in a lesion-free gland. Run it on your own negatives before you show its output to anyone.
• Lenient detection rule. A predicted lesion counts as found if it overlaps a reference lesion with Dice above 0.1. The detection numbers are therefore generous about boundary quality.
• Small public test set. The Prostate158 test split behind the headline number has 19 cases.
• Single fold, final checkpoint. Every configuration was trained once on one fold, so small differences between augmentations may sit inside run-to-run noise.
• Input preparation is on you. The model expects ADC and DWI already resampled to the T2 grid and stored as NIfTI in nnU-Net’s channel order. When you start from DICOM series, that conversion and resampling step is where most integration bugs will hide.
• Not tested out of domain. The authors list a true out-of-domain evaluation as future work.

Code, weights and license

The trainers live on GitHub at tppm/prostate_lesion_segmentation under the MIT license. They plug into nnU-Net v2 as extra trainer classes driven by small YAML configs, and the repo includes the evaluation script and the train, validation and test split lists for the public cohorts. The final checkpoint (ResEnc-M, 3d_fullres, fold 0, about 400 MB as safetensors) is on Hugging Face at heintz/prostate_lesion_seg, also marked MIT. A few rough edges as of 9 November 2026: the README still shows a placeholder for the Hugging Face repo name, the trainer imports a bicc_loss module that is not in the repo, and the README says the models were trained on an nnU-Net fork without linking the commit. Expect some plumbing before the first prediction.

The MIT tag covers the code and the weights as published. The training data come from public cohorts with their own terms plus an internal cohort that is not shared, so check those terms before building a commercial product on the checkpoint. The model is research software and is not cleared as a medical device.

Paper: arXiv:2611.08535 · Code: GitHub (MIT) · Weights: Hugging Face (MIT)

Originally published on pycad.co, the canonical version: https://pycad.co/blog/random-wavelet-nnunet-prostate-lesion-bpmri/

Random Wavelet nnU-Net: Prostate Lesion Masks Across bpMRI Sites was originally published in Python in Plain English on Medium, where people are continuing the conversation by highlighting and responding to this story.

Retrieval augmented generation, usually shortened to RAG, is often sold as the fix for made up answers. The model looks up documents first, then answers from them. In practice a RAG system can still give wrong answers, and the cause is rarely the language model itself.

Most failures start earlier in the pipeline. Understanding where they come from saves a lot of time, because tuning the prompt will not repair a retrieval problem.

The right passage never gets retrieved

If the answer is not in the retrieved text, the model has two options: admit it does not know, or fill the gap from its training data. Many models choose the second. So the first thing to check is whether the correct passage appears in the top results at all.

Common reasons it does not: the question uses different words than the document, the embedding model handles your domain terms poorly, or the search returns only a few results and the useful one ranks sixth. Combining keyword search with vector search often helps here, because exact terms such as product codes or error names are easy for keywords and hard for embeddings.

Chunks cut the answer in half

Documents are split into chunks before they are indexed. A fixed size split can separate a rule from its exception, or a table from its heading. The model then sees a fragment that looks complete and answers confidently from it.

Splitting along natural boundaries such as headings and paragraphs usually works better than counting characters. Keeping a small overlap between chunks, and storing the document title with each one, gives the model more context to work with.

Old or conflicting documents

A knowledge base that holds three versions of the same policy will return all three. The model has no way to know which one is current unless you tell it. Adding dates or version labels as metadata, and filtering out superseded files, removes a whole class of errors that look like model mistakes but are really data problems.

Too much context can hurt

It is tempting to retrieve twenty chunks and let the model sort it out. Long contexts dilute the relevant passage, and models tend to pay less attention to material in the middle. Retrieving fewer, better ranked chunks, perhaps with a reranking step, often improves accuracy and lowers cost at the same time.

The prompt does not demand grounding

Even with good retrieval, the instructions matter. A prompt that says “answer the question” invites the model to use whatever it knows. A prompt that says to answer only from the supplied passages, to quote the source, and to say so when the passages do not contain the answer gives a clearer rule to follow. Asking for citations also makes errors easier to spot, because a reader can check the quoted passage.

Nobody measures it

The most common problem is that the system is judged by a few demo questions. A small test set of real questions with known answers lets you measure two things separately: did retrieval find the right passage, and did the model use it correctly. If retrieval scores are low, work on chunking and search. If retrieval is fine and answers are still wrong, work on the prompt and the model.

A short debugging order

When an answer is wrong, check in this order. First, was the correct passage retrieved? Second, was it complete, or cut off by chunking? Third, was it outdated or contradicted by another document? Fourth, did the prompt tell the model to stay within the sources? Most bugs are found in the first two steps.

It also helps to log every request. Store the question, the retrieved chunks with their scores, the final prompt and the answer. When a user reports a wrong answer, you can see at once whether retrieval or generation failed, without guessing. Over a few weeks these logs also show which questions your documents cannot answer, which points to content that is missing rather than a bug.

Finally, decide what the system should do when it is unsure. A low retrieval score is a useful signal. Instead of answering anyway, the system can say that it could not find a reliable source and suggest where to look. Users tend to trust a tool more when it admits its limits than when it answers every question with the same confidence.

RAG reduces made up answers when each stage does its job. Treat it as a search system with a writer on the end, and test the search part first.

Why RAG Systems Still Give Wrong Answers, and Where to Look First was originally published in Python in Plain English on Medium, where people are continuing the conversation by highlighting and responding to this story.